Skip to content

fix(storage): retry contended writes across the whole Turn path #319

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zch0227 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(storage): retry contended writes across the whole Turn path #319

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
120 changes: 59 additions & 61 deletions packages/local-runtime-v2/src/infra/db/write-transaction.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,12 @@
import { setTimeout as delay } from 'node:timers/promises';
import { sql } from 'drizzle-orm';
import {
retrySqliteWrite,
SqliteWriteAttemptFailedError,
SqliteWriteRetryAbortedError,
} from '@mavis/shared/sqlite-write-retry';

import type { AppDb } from './client.js';

const WRITE_LOCK_BUDGET_MS = 10_000;
const WRITE_LOCK_ATTEMPT_MS = 50;

/** Cancellation before the mutation callback starts; no write needs to be replayed. */
Expand All @@ -15,72 +19,66 @@ export class WriteLockWaitAbortedError extends Error {
}

/**
* Retry only transaction admission: a callback that has started is never replayed.
* The signal cancels contention waits, not an immediately available write. This
* lets post-cancellation tool completion and cleanup messages remain durable.
* The transaction handle drizzle hands to a `transaction()` callback. Callers
* thread this through their `*InTransaction` helpers, so the retrying call has
* to expose the same type the raw drizzle call used to.
*/
type DrizzleWriteTransaction = Parameters<Parameters<AppDb['transaction']>[0]>[0];

/**
* Run a `BEGIN IMMEDIATE` mutation under the shared bounded-stall retry policy
* (see `@mavis/shared/sqlite-write-retry`). Only lock acquisition is retried:
* once the callback has started, any error — busy or not — rolls the
* transaction back and propagates without a replay.
*
* The optional signal cancels contention waits only. An immediately available
* write still commits, which keeps post-cancellation tool completion and
* cleanup messages durable.
*/
export async function runWithWriteLock<T>(
db: AppDb,
mutation: (tx: AppDb) => T,
mutation: (tx: DrizzleWriteTransaction) => T,
options: { readonly signal?: AbortSignal; readonly timeoutMs?: number } = {},
): Promise<T> {
const budget = options.timeoutMs ?? WRITE_LOCK_BUDGET_MS;
if (!Number.isFinite(budget) || budget <= 0)
throw new RangeError('Invalid write lock budget');
const deadline = performance.now() + budget;
let attempt = 0;
let hasContended = false;
let lastBusy: unknown = new Error('SQLite write lock wait exceeded its deadline');
for (;;) {
if (hasContended) throwIfWaitAborted(options.signal);
const remaining = deadline - performance.now();
if (remaining <= 0) throw lastBusy;
const previous = db.get<{ timeout: number }>(sql`PRAGMA busy_timeout`).timeout;
let entered = false;
try {
const nativeWaitMs = options.signal?.aborted
? 0
: Math.ceil(Math.min(WRITE_LOCK_ATTEMPT_MS, remaining));
db.run(sql.raw(`PRAGMA busy_timeout = ${nativeWaitMs}`));
return db.transaction(
(tx) => {
entered = true;
// Only lock acquisition gets a short timeout. Restore the connection's
// policy before callbacks (including nested transactions) can use it.
try {
return await retrySqliteWrite(
({ remainingMs }) => {
const previous = db.get<{ timeout: number }>(sql`PRAGMA busy_timeout`).timeout;
let entered = false;
try {
// Only lock acquisition gets a short native wait; the callback below
// restores the connection policy before any statement runs.
db.run(
sql.raw(
`PRAGMA busy_timeout = ${Math.min(WRITE_LOCK_ATTEMPT_MS, Math.ceil(remainingMs))}`,
),
);
return db.transaction(
(tx) => {
entered = true;
db.run(sql.raw(`PRAGMA busy_timeout = ${previous}`));
return mutation(tx);
},
{ behavior: 'immediate' },
);
} catch (error) {
// A callback that already started owns its own rollback; only a
// failure to acquire the write lock is safe to replay.
if (entered) throw new SqliteWriteAttemptFailedError(error);
throw error;
} finally {
// No await occurs while the shared connection has a temporary timeout.
db.run(sql.raw(`PRAGMA busy_timeout = ${previous}`));
return mutation(tx);
},
{ behavior: 'immediate' },
);
} catch (error) {
if (entered || !isBusy(error)) throw error;
lastBusy = error;
} finally {
// No await occurs while the shared connection has a temporary timeout.
db.run(sql.raw(`PRAGMA busy_timeout = ${previous}`));
}
hasContended = true;
throwIfWaitAborted(options.signal);
const wait = Math.min(
deadline - performance.now(),
25 * 2 ** Math.min(attempt++, 3) + Math.random() * 25,
}
},
{ timeoutMs: options.timeoutMs, signal: options.signal },
);
if (wait <= 0) throw lastBusy;
try {
await delay(wait, undefined, { signal: options.signal });
} catch (error) {
if (options.signal?.aborted && error instanceof Error && error.name === 'AbortError') {
throw new WriteLockWaitAbortedError(options.signal);
}
throw error;
} catch (error) {
// Only an abandoned wait becomes a Turn-facing cancellation; an error the
// callback raised after it started keeps its own identity.
if (error instanceof SqliteWriteRetryAbortedError) {
throw new WriteLockWaitAbortedError(options.signal as AbortSignal);
}
throw error;
}
}

function throwIfWaitAborted(signal: AbortSignal | undefined): void {
if (signal?.aborted) throw new WriteLockWaitAbortedError(signal);
}

function isBusy(error: unknown): boolean {
return error instanceof Error && Reflect.get(error, 'code') === 'SQLITE_BUSY';
}
66 changes: 34 additions & 32 deletions packages/local-runtime-v2/src/service/session-system/messages/repo/drizzle.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ class DrizzleMessageRepository implements MessageRepository {
}

async get(sessionId: string, msgId: string): Promise<DisplayMessageRecord | undefined> {
this.ensureReady(sessionId);
await this.ensureReady(sessionId);
const db = this.options.db;
let query = messageReads.get(db);
if (!query) {
Expand All @@ -86,7 +86,7 @@ class DrizzleMessageRepository implements MessageRepository {
}

async list(sessionId: string, options: ListMessagesOptions = {}): Promise<ListMessagesResult> {
this.ensureReady(sessionId);
await this.ensureReady(sessionId);
const anchorId = options.before ? this.messageRowId(sessionId, options.before) : undefined;
const limit = normalizeLimit(options.limit, 0);
const query = this.options.db
Expand All @@ -110,7 +110,7 @@ class DrizzleMessageRepository implements MessageRepository {
}

async listAfter(sessionId: string, afterMsgId?: string): Promise<MessageReplayResult> {
this.ensureReady(sessionId);
await this.ensureReady(sessionId);
const anchorId = afterMsgId ? this.messageRowId(sessionId, afterMsgId) : undefined;
if (afterMsgId && anchorId === undefined) return { status: 'missing-anchor', messages: [] };
const rows = this.options.db
Expand All @@ -128,7 +128,7 @@ class DrizzleMessageRepository implements MessageRepository {
}

async listTurn(sessionId: string, turnId: string): Promise<DisplayMessageRecord[]> {
this.ensureReady(sessionId);
await this.ensureReady(sessionId);
const db = this.options.db;
let query = turnReads.get(db);
if (!query) {
Expand All @@ -147,7 +147,7 @@ class DrizzleMessageRepository implements MessageRepository {
},
): Promise<DisplayMessageRecord[]> {
if (!Number.isFinite(options.limit) || options.limit <= 0) return [];
this.ensureReady(sessionId);
await this.ensureReady(sessionId);
const rows = this.options.db
.select()
.from(messageRows)
Expand Down Expand Up @@ -175,7 +175,8 @@ class DrizzleMessageRepository implements MessageRepository {
generatedIdDiscriminator: `user:${input.turnId}`,
nowMs: this.nowMs,
});
return this.options.db.transaction(
return await runWithWriteLock(
this.options.db,
(tx) => {
const rejection = this.options.userMessageAdmission?.rejectionInTransaction(tx, {
sessionId: input.sessionId,
Expand Down Expand Up @@ -235,7 +236,6 @@ class DrizzleMessageRepository implements MessageRepository {
message: decodeDisplayMessage(inserted),
};
},
{ behavior: 'immediate' },
);
}

Expand Down Expand Up @@ -282,7 +282,7 @@ class DrizzleMessageRepository implements MessageRepository {

async replace(input: MessageReplaceInput): Promise<void> {
const normalized = this.normalizeReplacementBatch(input.messages, 0);
this.writeReplacementBatch(input.sessionId, normalized, true);
await this.writeReplacementBatch(input.sessionId, normalized, true);
}

async replaceStream(input: MessageReplaceStreamInput): Promise<void> {
Expand All @@ -291,15 +291,16 @@ class DrizzleMessageRepository implements MessageRepository {
for await (const batch of input.batches) {
if (batch.length === 0) continue;
const normalized = this.normalizeReplacementBatch(batch, messageIndex);
this.writeReplacementBatch(input.sessionId, normalized, replaceExisting);
await this.writeReplacementBatch(input.sessionId, normalized, replaceExisting);
messageIndex += normalized.length;
replaceExisting = false;
}
if (replaceExisting) this.writeReplacementBatch(input.sessionId, [], true);
if (replaceExisting) await this.writeReplacementBatch(input.sessionId, [], true);
}

async rewindInclusive(input: MessageRewindInclusiveInput): Promise<MessageRewindInclusiveResult> {
return this.options.db.transaction(
return await runWithWriteLock(
this.options.db,
(tx) => {
this.ensureReadyInTransaction(tx, input.sessionId);
if (!input.fromMessageId.startsWith('msg-user-v1-')) {
Expand Down Expand Up @@ -366,12 +367,11 @@ class DrizzleMessageRepository implements MessageRepository {
this.markSessionAssetIndexCurrent(tx, input.sessionId);
return { deletedMessageIds };
},
{ behavior: 'immediate' },
);
}

async rewind(input: MessageRewindInput): Promise<void> {
this.mutationTransaction(input.sessionId, (tx) => {
await this.mutationTransaction(input.sessionId, (tx) => {
const ids = new Set(input.messageIds ?? []);
const anchorId = input.afterMessageId
? this.messageRowId(input.sessionId, input.afterMessageId, tx)
Expand Down Expand Up @@ -401,7 +401,7 @@ class DrizzleMessageRepository implements MessageRepository {
}

async resolveTurnSource(sessionId: string, turnId: string) {
this.ensureReady(sessionId);
await this.ensureReady(sessionId);
const row = this.options.db
.select()
.from(messageRows)
Expand All @@ -419,7 +419,7 @@ class DrizzleMessageRepository implements MessageRepository {
}

async latestDisplayRowId(sessionId: string): Promise<number> {
this.ensureReady(sessionId);
await this.ensureReady(sessionId);
return (
this.options.db
.select({ id: messageRows.id })
Expand All @@ -444,12 +444,12 @@ class DrizzleMessageRepository implements MessageRepository {
targetSessionId: string;
throughMessageId: string;
}): Promise<void> {
this.options.db.transaction(
await runWithWriteLock(
this.options.db,
(tx) => {
this.copyPrefixInTransaction(tx, input);
this.appendForkOriginInTransaction(tx, input);
},
{ behavior: 'immediate' },
);
}

Expand Down Expand Up @@ -514,7 +514,8 @@ class DrizzleMessageRepository implements MessageRepository {
targetSessionId: string;
sourceSessionId: string;
}): Promise<void> {
this.options.db.transaction((tx) => {
await runWithWriteLock(
this.options.db,(tx) => {
this.ensureReadyInTransaction(tx, input.targetSessionId);
this.appendForkOriginInTransaction(tx, input);
});
Expand Down Expand Up @@ -544,7 +545,8 @@ class DrizzleMessageRepository implements MessageRepository {
}

async deleteSessionData(sessionId: string): Promise<void> {
this.options.db.transaction((tx) => {
await runWithWriteLock(
this.options.db,(tx) => {
tx.delete(sessionAssets).where(eq(sessionAssets.sessionId, sessionId)).run();
tx.delete(sessionAssetIndexState)
.where(eq(sessionAssetIndexState.sessionId, sessionId))
Expand All @@ -567,23 +569,23 @@ class DrizzleMessageRepository implements MessageRepository {
});
}

private ensureReady(sessionId: string): void {
private async ensureReady(sessionId: string): Promise<void> {
const marker = this.options.db
.select({ sessionId: messageRowMigrations.sessionId })
.from(messageRowMigrations)
.where(eq(messageRowMigrations.sessionId, sessionId))
.get();
if (marker) return;
this.mutationTransaction(sessionId, () => undefined);
await this.mutationTransaction(sessionId, () => undefined);
}
private mutationTransaction<T>(sessionId: string, mutation: (db: AppDb) => T): T {
return this.options.db.transaction(
(tx) => {
this.ensureReadyInTransaction(tx, sessionId);
return mutation(tx);
},
{ behavior: 'immediate' },
);
private mutationTransaction<T>(
sessionId: string,
mutation: (db: AppDb) => T,
): Promise<T> {
return runWithWriteLock(this.options.db, (tx) => {
this.ensureReadyInTransaction(tx, sessionId);
return mutation(tx);
});
}
private ensureReadyInTransaction(db: AppDb, sessionId: string): void {
ensureMessageRowsReadyInTransaction(db, sessionId, this.nowMs(), (message, index) => {
Expand Down Expand Up @@ -653,12 +655,12 @@ class DrizzleMessageRepository implements MessageRepository {
}),
);
}
private writeReplacementBatch(
private async writeReplacementBatch(
sessionId: string,
messages: readonly NormalizedDisplayMessage[],
replaceExisting: boolean,
): void {
this.mutationTransaction(sessionId, (tx) => {
): Promise<void> {
await this.mutationTransaction(sessionId, (tx) => {
if (replaceExisting) {
tx.delete(messageRows).where(eq(messageRows.sessionId, sessionId)).run();
tx.delete(sessionAssets).where(eq(sessionAssets.sessionId, sessionId)).run();
Expand Down
Loading
Loading