refactor(data): unify document mutation locking and transaction handling

Replace mutation-lock with document-lock to centralize document mutation
serialization and database transaction management. Introduce runInDbTransaction
utility to abstract SQLite/Postgres transaction differences. Update document
deletion, user data cleanup, and rate limiter logic to delegate transaction
handling and locking to shared helpers. Remove dialect-specific branching and
inline transaction logic for improved maintainability and testability.

BREAKING CHANGE: withDocumentMutationLock is removed in favor of withDocumentLock and runInDbTransaction

src/app/api/documents/blob/upload/finalize/route.ts

@@ -24,7 +24,7 @@ import { errorToLog, serverLogger } from '@/lib/server/logger';
 import { getOpenReaderTestNamespace } from '@/lib/server/testing/test-namespace';
 import { isS3Configured } from '@/lib/server/storage/s3';
 import type { BaseDocument, DocumentType } from '@/types/documents';
-import { withDocumentMutationLock } from '@/lib/server/documents/mutation-lock';
+import { withDocumentLock } from '@/lib/server/documents/document-lock';
 
 export const dynamic = 'force-dynamic';
 
@@ -150,7 +150,7 @@ async function finalizeOne(input: {
     : input.upload.name;
   const documentId = createHash('sha256').update(finalizedBody).digest('hex');
 
-  const stored = await withDocumentMutationLock(documentId, async () => {
+  const stored = await withDocumentLock(documentId, async () => {
     try {
       await headDocumentBlob(documentId, input.namespace);
     } catch (error) {

src/db/run-in-transaction.ts

@@ -0,0 +1,23 @@
+import { db } from '@/db';
+
+/**
+ * Run `fn` with a database connection, wrapped in a transaction on Postgres.
+ *
+ * This is the single definition of the SQLite/Postgres bridge, so callers never
+ * branch on `process.env.POSTGRES_URL` themselves:
+ *   - Postgres: opens a real transaction so a multi-statement read-modify-write
+ *     is atomic, and passes the transaction handle as `conn`.
+ *   - SQLite: better-sqlite3 transactions require synchronous callbacks (they
+ *     cannot be awaited) and the database is a single writer anyway, so `fn`
+ *     runs directly on the shared connection.
+ */
+export async function runInDbTransaction<T>(
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  fn: (conn: any) => Promise<T>,
+): Promise<T> {
+  if (process.env.POSTGRES_URL) {
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    return (db as any).transaction(async (tx: any) => fn(tx));
+  }
+  return fn(db);
+}

src/lib/server/admin/settings.ts

@@ -264,25 +264,13 @@ export async function setRuntimeConfigKey<K extends RuntimeConfigKey>(
   }
   const serialized = serializeForStorage(validated);
   const now = Date.now();
-  // Upsert with onConflict.
+  await db
-  const isPg = !!process.env.POSTGRES_URL;
+    .insert(adminSettings)
-  if (isPg) {
+    .values({ key, valueJson: serialized as never, source: 'admin', updatedAt: now })
-    await db
+    .onConflictDoUpdate({
-      .insert(adminSettings)
+      target: adminSettings.key,
-      .values({ key, valueJson: serialized as never, source: 'admin', updatedAt: now })
+      set: { valueJson: serialized as never, source: 'admin', updatedAt: now },
-      .onConflictDoUpdate({
+    });
-        target: adminSettings.key,
-        set: { valueJson: serialized as never, source: 'admin', updatedAt: now },
-      });
-  } else {
-    await db
-      .insert(adminSettings)
-      .values({ key, valueJson: serialized as never, source: 'admin', updatedAt: now })
-      .onConflictDoUpdate({
-        target: adminSettings.key,
-        set: { valueJson: serialized as never, source: 'admin', updatedAt: now },
-      });
-  }
 }
 
 /** Delete a runtime config row (resets to default/implicit behavior). */

src/lib/server/documents/delete-owned.ts

@@ -1,5 +1,4 @@
 import { and, eq, ne } from 'drizzle-orm';
-import { db } from '@/db';
 import { documents } from '@/db/schema';
 import { deleteDocumentBlob } from '@/lib/server/documents/blobstore';
 import {
@@ -7,73 +6,38 @@ import {
   deleteDocumentPreviewRows,
 } from '@/lib/server/documents/previews';
 import { deleteDocumentTtsSegmentCache } from '@/lib/server/tts/segments-cache';
-import { withDocumentMutationLock } from '@/lib/server/documents/mutation-lock';
+import { withDocumentLock } from '@/lib/server/documents/document-lock';
 
 type DocumentRow = typeof documents.$inferSelect;
 
-async function restoreDocumentOwnership(row: DocumentRow): Promise<void> {
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
-  await db.insert(documents).values(row).onConflictDoNothing();
+async function restoreDocumentOwnership(conn: any, row: DocumentRow): Promise<void> {
+  await conn.insert(documents).values(row).onConflictDoNothing();
 }
 
-async function removeOwnershipAndCheckLastOwner(input: {
+async function removeOwnershipAndCheckLastOwner(
-  userId: string;
-  documentId: string;
-}): Promise<{ removed: DocumentRow | null; isLastOwner: boolean }> {
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  const database = db as any;
+  conn: any,
-  const runAsync = async (tx: typeof database) => {
+  input: { userId: string; documentId: string },
-    await tx
+): Promise<{ removed: DocumentRow | null; isLastOwner: boolean }> {
-      .select({ id: documents.id })
+  const [removed] = await conn
-      .from(documents)
+    .delete(documents)
-      .where(eq(documents.id, input.documentId))
+    .where(and(
-      .for('update');
+      eq(documents.id, input.documentId),
+      eq(documents.userId, input.userId),
+    ))
+    .returning();
+  if (!removed) return { removed: null, isLastOwner: false };
 
-    const [removed] = await tx
+  const otherOwners = await conn
-      .delete(documents)
+    .select({ id: documents.id })
-      .where(and(
+    .from(documents)
-        eq(documents.id, input.documentId),
+    .where(and(
-        eq(documents.userId, input.userId),
+      eq(documents.id, input.documentId),
-      ))
+      ne(documents.userId, input.userId),
-      .returning();
+    ))
-    if (!removed) return { removed: null, isLastOwner: false };
+    .limit(1);
-
+  return { removed, isLastOwner: otherOwners.length === 0 };
-    const otherOwners = await tx
-      .select({ id: documents.id })
-      .from(documents)
-      .where(and(
-        eq(documents.id, input.documentId),
-        ne(documents.userId, input.userId),
-      ))
-      .limit(1);
-    return { removed, isLastOwner: otherOwners.length === 0 };
-  };
-
-  if (process.env.POSTGRES_URL) {
-    return database.transaction(runAsync);
-  }
-
-  return database.transaction((tx: typeof database) => {
-    const removed = tx
-      .delete(documents)
-      .where(and(
-        eq(documents.id, input.documentId),
-        eq(documents.userId, input.userId),
-      ))
-      .returning()
-      .all()[0] ?? null;
-    if (!removed) return { removed: null, isLastOwner: false };
-
-    const otherOwners = tx
-      .select({ id: documents.id })
-      .from(documents)
-      .where(and(
-        eq(documents.id, input.documentId),
-        ne(documents.userId, input.userId),
-      ))
-      .limit(1)
-      .all();
-    return { removed, isLastOwner: otherOwners.length === 0 };
-  }, { behavior: 'immediate' });
 }
 
 export async function deleteOwnedDocument(input: {
@@ -81,8 +45,8 @@ export async function deleteOwnedDocument(input: {
   documentId: string;
   namespace: string | null;
 }): Promise<boolean> {
-  return withDocumentMutationLock(input.documentId, async () => {
+  return withDocumentLock(input.documentId, async (conn) => {
-    const { removed, isLastOwner } = await removeOwnershipAndCheckLastOwner(input);
+    const { removed, isLastOwner } = await removeOwnershipAndCheckLastOwner(conn, input);
     if (!removed) return false;
 
     try {
@@ -92,7 +56,7 @@ export async function deleteOwnedDocument(input: {
         await cleanupDocumentPreviewArtifacts(input.documentId, input.namespace);
         await deleteDocumentPreviewRows(input.documentId, input.namespace);
 
-        const [newOwner] = await db
+        const [newOwner] = await conn
           .select({ id: documents.id })
           .from(documents)
           .where(eq(documents.id, input.documentId))
@@ -102,7 +66,7 @@ export async function deleteOwnedDocument(input: {
         }
       }
     } catch (error) {
-      await restoreDocumentOwnership(removed);
+      await restoreDocumentOwnership(conn, removed);
       throw error;
     }
 

src/lib/server/documents/document-lock.ts

@@ -0,0 +1,59 @@
+import { sql } from 'drizzle-orm';
+import { runInDbTransaction } from '@/db/run-in-transaction';
+
+// In-process serialization for the single-writer SQLite deployment. SQLite runs
+// in one process, so an in-memory promise chain per document is sufficient (and
+// is the only lock available — there is no advisory lock). On Postgres this is
+// unused: the transaction-scoped advisory lock below provides exclusion that
+// works across a stateless/serverless fleet.
+const localTails = new Map<string, Promise<void>>();
+
+async function withLocalDocumentLock<T>(documentId: string, fn: () => Promise<T>): Promise<T> {
+  const previous = localTails.get(documentId) ?? Promise.resolve();
+  let release!: () => void;
+  const current = new Promise<void>((resolve) => {
+    release = resolve;
+  });
+  localTails.set(documentId, current);
+
+  await previous;
+  try {
+    return await fn();
+  } finally {
+    release();
+    if (localTails.get(documentId) === current) {
+      localTails.delete(documentId);
+    }
+  }
+}
+
+/**
+ * Serialize all mutations to a single document and run `fn` with exclusive
+ * access, passing it the connection to do its work on.
+ *
+ * The SQLite/Postgres transaction bridge is delegated to `runInDbTransaction`;
+ * this helper only adds the exclusion on top:
+ *   - Postgres: a transaction-scoped advisory lock keyed on the document id,
+ *     acquired as the first statement of the shared transaction so it covers
+ *     the whole read-modify-write. It releases on commit, so it is safe under
+ *     transaction-mode poolers and gives fleet-wide exclusion in a stateless
+ *     deployment.
+ *   - SQLite: an in-process lock; the single WAL writer needs nothing more.
+ *
+ * Because this guarantees exclusivity for the document, `fn` does not need its
+ * own transaction or `SELECT ... FOR UPDATE` to read-modify-write safely.
+ */
+export async function withDocumentLock<T>(
+  documentId: string,
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  fn: (conn: any) => Promise<T>,
+): Promise<T> {
+  if (process.env.POSTGRES_URL) {
+    return runInDbTransaction(async (conn) => {
+      await conn.execute(sql`SELECT pg_advisory_xact_lock(hashtextextended(${documentId}, 0))`);
+      return fn(conn);
+    });
+  }
+
+  return withLocalDocumentLock(documentId, () => runInDbTransaction(fn));
+}

src/lib/server/documents/mutation-lock.ts

@@ -1,40 +0,0 @@
-import { sql } from 'drizzle-orm';
-import { db } from '@/db';
-
-const localTails = new Map<string, Promise<void>>();
-
-async function withLocalDocumentLock<T>(documentId: string, fn: () => Promise<T>): Promise<T> {
-  const previous = localTails.get(documentId) ?? Promise.resolve();
-  let release!: () => void;
-  const current = new Promise<void>((resolve) => {
-    release = resolve;
-  });
-  localTails.set(documentId, current);
-
-  await previous;
-  try {
-    return await fn();
-  } finally {
-    release();
-    if (localTails.get(documentId) === current) {
-      localTails.delete(documentId);
-    }
-  }
-}
-
-export async function withDocumentMutationLock<T>(
-  documentId: string,
-  fn: () => Promise<T>,
-): Promise<T> {
-  return withLocalDocumentLock(documentId, async () => {
-    if (!process.env.POSTGRES_URL) {
-      return fn();
-    }
-
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    return (db as any).transaction(async (tx: any) => {
-      await tx.execute(sql`SELECT pg_advisory_xact_lock(hashtextextended(${documentId}, 0))`);
-      return fn();
-    });
-  });
-}

src/lib/server/rate-limit/rate-limiter.ts

@@ -1,4 +1,5 @@
 import { db } from '@/db';
+import { runInDbTransaction } from '@/db/run-in-transaction';
 import { userTtsChars } from '@/db/schema';
 import { eq, and, lt, sql } from 'drizzle-orm';
 import { nextUtcMidnightTimestampMs, nowTimestampMs } from '@/lib/shared/timestamps';
@@ -127,25 +128,10 @@ function getRowsAffected(result: unknown): number {
 export class RateLimiter {
   constructor() { }
 
-  private isPostgres(): boolean {
-    return Boolean(process.env.POSTGRES_URL);
-  }
-
   private getUpdatedAtValue(): number {
     return nowTimestampMs();
   }
 
-  // Use a transaction only when running with Postgres.
-  // better-sqlite3 transactions require sync callbacks and cannot be awaited.
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  private async runMutation<T>(fn: (conn: any) => Promise<T>): Promise<T> {
-    if (this.isPostgres()) {
-      // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      return safeDb().transaction(async (tx: any) => fn(tx));
-    }
-    return fn(safeDb());
-  }
-
   /**
    * Check if a user can use TTS and increment their char count if allowed
    */
@@ -189,7 +175,7 @@ export class RateLimiter {
 
     try {
       const updatedAt = this.getUpdatedAtValue() as unknown as UserTtsCharsUpdatedAtValue;
-      return await this.runMutation(async (conn) => {
+      return await runInDbTransaction(async (conn) => {
         // Ensure records exist for each bucket
         for (const bucket of buckets) {
           await conn.insert(userTtsChars)
@@ -322,7 +308,7 @@ export class RateLimiter {
   async transferAnonymousUsage(anonymousUserId: string, authenticatedUserId: string): Promise<void> {
     const updatedAt = this.getUpdatedAtValue() as unknown as UserTtsCharsUpdatedAtValue;
 
-    await this.runMutation(async (conn) => {
+    await runInDbTransaction(async (conn) => {
       const anonymousRows = await conn.select({
         date: userTtsChars.date,
         charCount: userTtsChars.charCount,

src/lib/server/user/data-cleanup.ts

@@ -21,7 +21,7 @@ import { audiobookPrefix, deleteAudiobookPrefix } from '@/lib/server/audiobooks/
 import { deleteTtsSegmentPrefix } from '@/lib/server/tts/segments-blobstore';
 import { hashForLog, serverLogger } from '@/lib/server/logger';
 import { logDegraded } from '@/lib/server/errors/logging';
-import { withDocumentMutationLock } from '@/lib/server/documents/mutation-lock';
+import { withDocumentLock } from '@/lib/server/documents/document-lock';
 
 type DocumentRow = typeof documents.$inferSelect;
 type AudiobookRow = { id: string };
@@ -55,10 +55,11 @@ export async function deleteUserStorageData(
 
   let docsDeleted = 0;
   const removedDocs: DocumentRow[] = [];
-  const restoreRemovedDocs = async () => {
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  const restoreRemovedDocs = async (conn: any = database) => {
     if (removedDocs.length === 0) return;
     const docsToRestore = [...removedDocs];
-    await database.insert(documents).values(docsToRestore).onConflictDoNothing();
+    await conn.insert(documents).values(docsToRestore).onConflictDoNothing();
     removedDocs.length = 0;
   };
   // Restore before throwing the diagnostic AggregateError, but never let a
@@ -76,62 +77,33 @@ export async function deleteUserStorageData(
       });
     }
   };
-  const removeOwnershipAndCheckLastOwner = async (doc: DocumentRow) => {
+  // The mutation lock already serializes all mutations for this document, so a
-    const run = async (tx: typeof database) => {
+  // plain read-modify-write is safe without an inner transaction or FOR UPDATE.
-      // Lock every ownership row for this document so concurrent account
+  const removeOwnershipAndCheckLastOwner = async (
-      // deletions cannot both observe the other owner's uncommitted row.
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      await tx
+    conn: any,
-        .select({ id: documents.id })
+    doc: DocumentRow,
-        .from(documents)
+  ) => {
-        .where(eq(documents.id, doc.id))
+    const [removedDoc] = await conn
-        .for('update');
+      .delete(documents)
+      .where(and(eq(documents.id, doc.id), eq(documents.userId, userId)))
+      .returning();
+    if (!removedDoc) return { removedDoc: null, isLastOwner: false };
 
-      const [removedDoc] = await tx
+    const otherOwners = await conn
-        .delete(documents)
+      .select({ id: documents.id })
-        .where(and(eq(documents.id, doc.id), eq(documents.userId, userId)))
+      .from(documents)
-        .returning();
+      .where(and(
-      if (!removedDoc) return { removedDoc: null, isLastOwner: false };
+        eq(documents.id, doc.id),
-
+        ne(documents.userId, userId),
-      const otherOwners = await tx
+      ))
-        .select({ id: documents.id })
+      .limit(1);
-        .from(documents)
+    return { removedDoc, isLastOwner: otherOwners.length === 0 };
-        .where(and(
-          eq(documents.id, doc.id),
-          ne(documents.userId, userId),
-        ))
-        .limit(1);
-      return { removedDoc, isLastOwner: otherOwners.length === 0 };
-    };
-
-    if (process.env.POSTGRES_URL) {
-      return database.transaction(run);
-    }
-
-    return database.transaction((tx: typeof database) => {
-      const removedRows = tx
-        .delete(documents)
-        .where(and(eq(documents.id, doc.id), eq(documents.userId, userId)))
-        .returning()
-        .all();
-      const removedDoc = removedRows[0] ?? null;
-      if (!removedDoc) return { removedDoc: null, isLastOwner: false };
-
-      const otherOwners = tx
-        .select({ id: documents.id })
-        .from(documents)
-        .where(and(
-          eq(documents.id, doc.id),
-          ne(documents.userId, userId),
-        ))
-        .limit(1)
-        .all();
-      return { removedDoc, isLastOwner: otherOwners.length === 0 };
-    }, { behavior: 'immediate' });
   };
 
   for (const doc of userDocs) {
-    await withDocumentMutationLock(doc.id, async () => {
+    await withDocumentLock(doc.id, async (conn) => {
-      const { removedDoc, isLastOwner } = await removeOwnershipAndCheckLastOwner(doc);
+      const { removedDoc, isLastOwner } = await removeOwnershipAndCheckLastOwner(conn, doc);
       if (!removedDoc) return;
       removedDocs.push(removedDoc);
 
@@ -193,7 +165,7 @@ export async function deleteUserStorageData(
       }
 
       if (namespace !== null) {
-        await restoreRemovedDocs();
+        await restoreRemovedDocs(conn);
       }
     });
   }

tests/unit/user-data-cleanup.vitest.spec.ts

@@ -5,7 +5,7 @@ const mocks = vi.hoisted(() => ({
   documentDeleteResults: [] as unknown[][],
   deleteWhere: vi.fn(async () => undefined),
   insertValues: vi.fn(() => ({ onConflictDoNothing: vi.fn(async () => undefined) })),
-  forUpdate: vi.fn(),
+  execute: vi.fn(async () => undefined),
   deleteDocumentBlob: vi.fn(async () => undefined),
   deleteDocumentPrefix: vi.fn(async () => 0),
   deleteDocumentPreviewArtifacts: vi.fn(async () => 0),
@@ -22,10 +22,6 @@ function resultBuilder(result: unknown[]) {
   };
   return {
     all: () => result,
-    for: vi.fn((mode: string) => {
-      mocks.forUpdate(mode);
-      return Promise.resolve(result);
-    }),
     limit: () => limitedResult,
     then: (resolve: (value: unknown[]) => unknown, reject: (error: unknown) => unknown) =>
       Promise.resolve(result).then(resolve, reject),
@@ -52,7 +48,7 @@ vi.mock('@/db', () => {
     insert: vi.fn(() => ({
       values: mocks.insertValues,
     })),
-    execute: vi.fn(async () => undefined),
+    execute: mocks.execute,
     transaction: vi.fn((callback: (tx: unknown) => unknown) => callback(database)),
   };
   return { db: database };
@@ -112,7 +108,7 @@ describe('user data cleanup', () => {
     }
     mocks.deleteWhere.mockResolvedValue(undefined);
     mocks.insertValues.mockReturnValue({ onConflictDoNothing: vi.fn(async () => undefined) });
-    mocks.forUpdate.mockReset();
+    mocks.execute.mockResolvedValue(undefined);
     mocks.deleteDocumentBlob.mockResolvedValue(undefined);
     mocks.deleteDocumentPrefix.mockResolvedValue(0);
     mocks.deleteDocumentPreviewArtifacts.mockResolvedValue(0);
@@ -177,7 +173,7 @@ describe('user data cleanup', () => {
     expect(mocks.insertValues).toHaveBeenCalledWith([{ id: 'doc-1' }]);
   });
 
-  test('locks all document owners before the Postgres ownership decision', async () => {
+  test('acquires a Postgres advisory lock before the ownership decision', async () => {
     vi.stubEnv('POSTGRES_URL', 'postgres://test');
     mocks.selectResults = [
       [{ id: 'doc-1' }],
@@ -189,6 +185,8 @@ describe('user data cleanup', () => {
 
     await deleteUserStorageData('user-1', null);
 
-    expect(mocks.forUpdate).toHaveBeenCalledWith('update');
+    // The mutation lock serializes the read-modify-write via a transaction-scoped
+    // advisory lock instead of SELECT ... FOR UPDATE.
+    expect(mocks.execute).toHaveBeenCalled();
   });
 });