diff --git a/src/app/api/documents/blob/upload/finalize/route.ts b/src/app/api/documents/blob/upload/finalize/route.ts index 76c6e24..428f166 100644 --- a/src/app/api/documents/blob/upload/finalize/route.ts +++ b/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) { diff --git a/src/db/run-in-transaction.ts b/src/db/run-in-transaction.ts new file mode 100644 index 0000000..1ec6ecd --- /dev/null +++ b/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( + // eslint-disable-next-line @typescript-eslint/no-explicit-any + fn: (conn: any) => Promise, +): Promise { + 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); +} diff --git a/src/lib/server/admin/settings.ts b/src/lib/server/admin/settings.ts index 9f65e4a..99e1b87 100644 --- a/src/lib/server/admin/settings.ts +++ b/src/lib/server/admin/settings.ts @@ -264,25 +264,13 @@ export async function setRuntimeConfigKey( } const serialized = serializeForStorage(validated); const now = Date.now(); - // Upsert with onConflict. - const isPg = !!process.env.POSTGRES_URL; - if (isPg) { - 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 }, - }); - } 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 }, - }); - } + 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). */ diff --git a/src/lib/server/documents/delete-owned.ts b/src/lib/server/documents/delete-owned.ts index c7deba9..302a15c 100644 --- a/src/lib/server/documents/delete-owned.ts +++ b/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 { - await db.insert(documents).values(row).onConflictDoNothing(); +// eslint-disable-next-line @typescript-eslint/no-explicit-any +async function restoreDocumentOwnership(conn: any, row: DocumentRow): Promise { + await conn.insert(documents).values(row).onConflictDoNothing(); } -async function removeOwnershipAndCheckLastOwner(input: { - userId: string; - documentId: string; -}): Promise<{ removed: DocumentRow | null; isLastOwner: boolean }> { +async function removeOwnershipAndCheckLastOwner( // eslint-disable-next-line @typescript-eslint/no-explicit-any - const database = db as any; - const runAsync = async (tx: typeof database) => { - await tx - .select({ id: documents.id }) - .from(documents) - .where(eq(documents.id, input.documentId)) - .for('update'); + conn: any, + input: { userId: string; documentId: string }, +): Promise<{ removed: DocumentRow | null; isLastOwner: boolean }> { + const [removed] = await conn + .delete(documents) + .where(and( + eq(documents.id, input.documentId), + eq(documents.userId, input.userId), + )) + .returning(); + if (!removed) return { removed: null, isLastOwner: false }; - const [removed] = await tx - .delete(documents) - .where(and( - eq(documents.id, input.documentId), - eq(documents.userId, input.userId), - )) - .returning(); - if (!removed) return { removed: null, isLastOwner: false }; - - 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' }); + const otherOwners = await conn + .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 }; } export async function deleteOwnedDocument(input: { @@ -81,8 +45,8 @@ export async function deleteOwnedDocument(input: { documentId: string; namespace: string | null; }): Promise { - return withDocumentMutationLock(input.documentId, async () => { - const { removed, isLastOwner } = await removeOwnershipAndCheckLastOwner(input); + return withDocumentLock(input.documentId, async (conn) => { + 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; } diff --git a/src/lib/server/documents/document-lock.ts b/src/lib/server/documents/document-lock.ts new file mode 100644 index 0000000..64099df --- /dev/null +++ b/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>(); + +async function withLocalDocumentLock(documentId: string, fn: () => Promise): Promise { + const previous = localTails.get(documentId) ?? Promise.resolve(); + let release!: () => void; + const current = new Promise((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( + documentId: string, + // eslint-disable-next-line @typescript-eslint/no-explicit-any + fn: (conn: any) => Promise, +): Promise { + 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)); +} diff --git a/src/lib/server/documents/mutation-lock.ts b/src/lib/server/documents/mutation-lock.ts deleted file mode 100644 index 15856e2..0000000 --- a/src/lib/server/documents/mutation-lock.ts +++ /dev/null @@ -1,40 +0,0 @@ -import { sql } from 'drizzle-orm'; -import { db } from '@/db'; - -const localTails = new Map>(); - -async function withLocalDocumentLock(documentId: string, fn: () => Promise): Promise { - const previous = localTails.get(documentId) ?? Promise.resolve(); - let release!: () => void; - const current = new Promise((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( - documentId: string, - fn: () => Promise, -): Promise { - 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(); - }); - }); -} diff --git a/src/lib/server/rate-limit/rate-limiter.ts b/src/lib/server/rate-limit/rate-limiter.ts index 1b2d3a0..a526d0d 100644 --- a/src/lib/server/rate-limit/rate-limiter.ts +++ b/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(fn: (conn: any) => Promise): Promise { - 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 { 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, diff --git a/src/lib/server/user/data-cleanup.ts b/src/lib/server/user/data-cleanup.ts index 52d7154..80c8f8c 100644 --- a/src/lib/server/user/data-cleanup.ts +++ b/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) => { - const run = async (tx: typeof database) => { - // Lock every ownership row for this document so concurrent account - // deletions cannot both observe the other owner's uncommitted row. - await tx - .select({ id: documents.id }) - .from(documents) - .where(eq(documents.id, doc.id)) - .for('update'); + // The mutation lock already serializes all mutations for this document, so a + // plain read-modify-write is safe without an inner transaction or FOR UPDATE. + const removeOwnershipAndCheckLastOwner = async ( + // eslint-disable-next-line @typescript-eslint/no-explicit-any + conn: any, + doc: DocumentRow, + ) => { + const [removedDoc] = await conn + .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 - .delete(documents) - .where(and(eq(documents.id, doc.id), eq(documents.userId, userId))) - .returning(); - if (!removedDoc) return { removedDoc: null, isLastOwner: false }; - - const otherOwners = await tx - .select({ id: documents.id }) - .from(documents) - .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' }); + const otherOwners = await conn + .select({ id: documents.id }) + .from(documents) + .where(and( + eq(documents.id, doc.id), + ne(documents.userId, userId), + )) + .limit(1); + return { removedDoc, isLastOwner: otherOwners.length === 0 }; }; for (const doc of userDocs) { - await withDocumentMutationLock(doc.id, async () => { - const { removedDoc, isLastOwner } = await removeOwnershipAndCheckLastOwner(doc); + await withDocumentLock(doc.id, async (conn) => { + 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); } }); } diff --git a/tests/unit/user-data-cleanup.vitest.spec.ts b/tests/unit/user-data-cleanup.vitest.spec.ts index 05cee91..95a8784 100644 --- a/tests/unit/user-data-cleanup.vitest.spec.ts +++ b/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(); }); });