|
| 1 | +import { setTimeout as delay } from 'node:timers/promises'; |
| 2 | +import { sql } from 'drizzle-orm'; |
| 3 | +import type { AppDb } from './client.js'; |
| 4 | + |
| 5 | +const WRITE_LOCK_BUDGET_MS = 10_000; |
| 6 | +const WRITE_LOCK_ATTEMPT_MS = 50; |
| 7 | + |
| 8 | +/** Cancellation before the mutation callback starts; no write needs to be replayed. */ |
| 9 | +export class WriteLockWaitAbortedError extends Error { |
| 10 | + override readonly name = 'WriteLockWaitAbortedError'; |
| 11 | + |
| 12 | + constructor(readonly signal: AbortSignal) { |
| 13 | + super('SQLite write lock wait was cancelled', { cause: signal.reason }); |
| 14 | + } |
| 15 | +} |
| 16 | + |
| 17 | +/** |
| 18 | + * Retry only transaction admission: a callback that has started is never replayed. |
| 19 | + * The signal cancels contention waits, not an immediately available write. This |
| 20 | + * lets post-cancellation tool completion and cleanup messages remain durable. |
| 21 | + */ |
| 22 | +export async function runWithWriteLock<T>( |
| 23 | + db: AppDb, |
| 24 | + mutation: (tx: AppDb) => T, |
| 25 | + options: { readonly signal?: AbortSignal; readonly timeoutMs?: number } = {}, |
| 26 | +): Promise<T> { |
| 27 | + const budget = options.timeoutMs ?? WRITE_LOCK_BUDGET_MS; |
| 28 | + if (!Number.isFinite(budget) || budget <= 0) |
| 29 | + throw new RangeError('Invalid write lock budget'); |
| 30 | + const deadline = performance.now() + budget; |
| 31 | + let attempt = 0; |
| 32 | + let hasContended = false; |
| 33 | + let lastBusy: unknown = new Error('SQLite write lock wait exceeded its deadline'); |
| 34 | + for (;;) { |
| 35 | + if (hasContended) throwIfWaitAborted(options.signal); |
| 36 | + const remaining = deadline - performance.now(); |
| 37 | + if (remaining <= 0) throw lastBusy; |
| 38 | + const previous = db.get<{ timeout: number }>(sql`PRAGMA busy_timeout`).timeout; |
| 39 | + let entered = false; |
| 40 | + try { |
| 41 | + const nativeWaitMs = options.signal?.aborted |
| 42 | + ? 0 |
| 43 | + : Math.ceil(Math.min(WRITE_LOCK_ATTEMPT_MS, remaining)); |
| 44 | + db.run(sql.raw(`PRAGMA busy_timeout = ${nativeWaitMs}`)); |
| 45 | + return db.transaction( |
| 46 | + (tx) => { |
| 47 | + entered = true; |
| 48 | + // Only lock acquisition gets a short timeout. Restore the connection's |
| 49 | + // policy before callbacks (including nested transactions) can use it. |
| 50 | + db.run(sql.raw(`PRAGMA busy_timeout = ${previous}`)); |
| 51 | + return mutation(tx); |
| 52 | + }, |
| 53 | + { behavior: 'immediate' }, |
| 54 | + ); |
| 55 | + } catch (error) { |
| 56 | + if (entered || !isBusy(error)) throw error; |
| 57 | + lastBusy = error; |
| 58 | + } finally { |
| 59 | + // No await occurs while the shared connection has a temporary timeout. |
| 60 | + db.run(sql.raw(`PRAGMA busy_timeout = ${previous}`)); |
| 61 | + } |
| 62 | + hasContended = true; |
| 63 | + throwIfWaitAborted(options.signal); |
| 64 | + const wait = Math.min( |
| 65 | + deadline - performance.now(), |
| 66 | + 25 * 2 ** Math.min(attempt++, 3) + Math.random() * 25, |
| 67 | + ); |
| 68 | + if (wait <= 0) throw lastBusy; |
| 69 | + try { |
| 70 | + await delay(wait, undefined, { signal: options.signal }); |
| 71 | + } catch (error) { |
| 72 | + if (options.signal?.aborted && error instanceof Error && error.name === 'AbortError') { |
| 73 | + throw new WriteLockWaitAbortedError(options.signal); |
| 74 | + } |
| 75 | + throw error; |
| 76 | + } |
| 77 | + } |
| 78 | +} |
| 79 | + |
| 80 | +function throwIfWaitAborted(signal: AbortSignal | undefined): void { |
| 81 | + if (signal?.aborted) throw new WriteLockWaitAbortedError(signal); |
| 82 | +} |
| 83 | + |
| 84 | +function isBusy(error: unknown): boolean { |
| 85 | + return error instanceof Error && Reflect.get(error, 'code') === 'SQLITE_BUSY'; |
| 86 | +} |
0 commit comments