danvics/openreader
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source 
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
This commit is contained in:
Richard R 2026-06-06 20:30:41 -06:00
parent 0df57b62b3
commit 523774f549
9 changed files with 158 additions and 208 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
src/app/api/documents/blob/upload/finalize/route.ts
View file

@ -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) {

23
src/db/run-in-transaction.ts Normal file
View file

@ -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);
}

26
src/lib/server/admin/settings.ts
View file

@ -264,25 +264,13 @@ export async function setRuntimeConfigKey<K extends RuntimeConfigKey>(
}
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). */

94
src/lib/server/documents/delete-owned.ts
View file

@ -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> {
await db.insert(documents).values(row).onConflictDoNothing();
// eslint-disable-next-line @typescript-eslint/no-explicit-any
async function restoreDocumentOwnership(conn: any, row: DocumentRow): Promise<void> {
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<boolean> {
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;
}

59
src/lib/server/documents/document-lock.ts Normal file
View file

@ -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));
}

40
src/lib/server/documents/mutation-lock.ts
View file

@ -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();
});
});
}

20
src/lib/server/rate-limit/rate-limiter.ts
View file

@ -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,

84
src/lib/server/user/data-cleanup.ts
View file

@ -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);
}
});
}

16
tests/unit/user-data-cleanup.vitest.spec.ts
View file

@ -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();
});
});