Cut sequential D1/KV round trips from the execute hot path

packages/worker/src/account/deletion-state.node.test.ts

@@ -86,3 +86,152 @@ test('non-expiring lease blocks deletion until audited repair fences callback',
 		reason: 'Inspected worker crash and confirmed process termination.',
 	})
 })
+
+function createLeaseTestDb() {
+	const sqlite = new DatabaseSync(':memory:')
+	sqlite.exec(`
+		CREATE TABLE users (
+			id INTEGER PRIMARY KEY,
+			stable_user_id TEXT UNIQUE,
+			deleting_at TEXT,
+			active_write_count INTEGER NOT NULL DEFAULT 0,
+			updated_at TEXT
+		);
+		CREATE TABLE account_write_leases (
+			token TEXT PRIMARY KEY,
+			user_id TEXT NOT NULL,
+			holder TEXT NOT NULL,
+			acquired_at TEXT NOT NULL,
+			released_at TEXT
+		);
+		INSERT INTO users (id, stable_user_id) VALUES (1, 'user-a');
+		INSERT INTO users (id, stable_user_id) VALUES (2, 'user-b');
+	`)
+	return { sqlite, db: createD1FromSqlite(sqlite) }
+}
+
+function countLeaseRows(sqlite: DatabaseSync) {
+	const row = sqlite
+		.prepare(`SELECT COUNT(*) AS count FROM account_write_leases`)
+		.get() as { count: number }
+	return Number(row.count)
+}
+
+test('nested same-user lease reuses the outer lease instead of re-acquiring', async () => {
+	const { sqlite, db } = createLeaseTestDb()
+	const result = await withAccountWriteLease({
+		db,
+		stableUserId: 'user-a',
+		holder: 'test:outer',
+		async write() {
+			expect(countLeaseRows(sqlite)).toBe(1)
+			return await withAccountWriteLease({
+				db,
+				stableUserId: 'user-a',
+				holder: 'test:nested',
+				async write() {
+					// Still only the outer lease: the nested call must not pay
+					// another acquire/release round trip.
+					expect(countLeaseRows(sqlite)).toBe(1)
+					const active = await listActiveAccountWriteLeases(db, 'user-a')
+					expect(active).toHaveLength(1)
+					expect(active[0]).toEqual(
+						expect.objectContaining({ holder: 'test:outer' }),
+					)
+					return 'nested-result'
+				},
+			})
+		},
+	})
+	expect(result).toBe('nested-result')
+	// The outer release is the only release.
+	expect(await listActiveAccountWriteLeases(db, 'user-a')).toHaveLength(0)
+	expect(
+		sqlite.prepare(`SELECT active_write_count FROM users WHERE id = 1`).get(),
+	).toEqual({ active_write_count: 0 })
+})
+
+test('nested lease for a different user still acquires its own lease', async () => {
+	const { sqlite, db } = createLeaseTestDb()
+	await withAccountWriteLease({
+		db,
+		stableUserId: 'user-a',
+		holder: 'test:outer',
+		async write() {
+			await withAccountWriteLease({
+				db,
+				stableUserId: 'user-b',
+				holder: 'test:other-user',
+				async write() {
+					expect(countLeaseRows(sqlite)).toBe(2)
+					expect(await listActiveAccountWriteLeases(db, 'user-b')).toHaveLength(
+						1,
+					)
+				},
+			})
+		},
+	})
+	expect(await listActiveAccountWriteLeases(db, 'user-a')).toHaveLength(0)
+	expect(await listActiveAccountWriteLeases(db, 'user-b')).toHaveLength(0)
+})
+
+test('detached work spawned inside write re-acquires after the outer lease releases', async () => {
+	const { sqlite, db } = createLeaseTestDb()
+	let detached: Promise<void> = Promise.resolve()
+	let releaseOuter: () => void = () => undefined
+	const outerReleased = new Promise<void>((resolve) => {
+		releaseOuter = resolve
+	})
+	await withAccountWriteLease({
+		db,
+		stableUserId: 'user-a',
+		holder: 'test:outer',
+		async write() {
+			// Simulates a waitUntil callback: created inside the lease scope
+			// (inheriting the AsyncLocalStorage context) but running only
+			// after the outer lease has been released.
+			detached = (async () => {
+				await outerReleased
+				await withAccountWriteLease({
+					db,
+					stableUserId: 'user-a',
+					holder: 'test:detached',
+					async write() {
+						const active = await listActiveAccountWriteLeases(db, 'user-a')
+						expect(active).toHaveLength(1)
+						expect(active[0]).toEqual(
+							expect.objectContaining({ holder: 'test:detached' }),
+						)
+					},
+				})
+			})()
+		},
+	})
+	releaseOuter()
+	await detached
+	expect(countLeaseRows(sqlite)).toBe(2)
+	expect(await listActiveAccountWriteLeases(db, 'user-a')).toHaveLength(0)
+	expect(
+		sqlite.prepare(`SELECT active_write_count FROM users WHERE id = 1`).get(),
+	).toEqual({ active_write_count: 0 })
+})
+
+test('sequential sibling leases each acquire after the previous released', async () => {
+	const { sqlite, db } = createLeaseTestDb()
+	await withAccountWriteLease({
+		db,
+		stableUserId: 'user-a',
+		async write() {
+			expect(countLeaseRows(sqlite)).toBe(1)
+		},
+	})
+	await withAccountWriteLease({
+		db,
+		stableUserId: 'user-a',
+		async write() {
+			expect(countLeaseRows(sqlite)).toBe(2)
+			expect(await listActiveAccountWriteLeases(db, 'user-a')).toHaveLength(1)
+		},
+	})
+	expect(await listActiveAccountWriteLeases(db, 'user-a')).toHaveLength(0)
+})

packages/worker/src/account/deletion-state.ts

@@ -1,3 +1,4 @@
+import { AsyncLocalStorage } from 'node:async_hooks'
 import { utcSqliteTimestamp } from '@kody-internal/shared/date-keys.ts'
 
 export class AccountDeletionInProgressError extends Error {
@@ -80,11 +81,55 @@ export async function assertAccountWritable(env: Env, stableUserId: string) {
 	await assertAccountWritableDb(env.APP_DB, stableUserId)
 }
 
+/**
+ * Live lease frames keyed by stable user id, propagated down the current
+ * async call chain. Nested {@link withAccountWriteLease} calls for the same
+ * user reuse the outer lease while its frame is still active instead of
+ * paying another acquire/release round trip to D1 (~5 statements each): the
+ * outer lease spans the nested write, so deletion stays blocked for exactly
+ * as long as it does today. MCP requests, app requests, job runs, and
+ * package invocations all take a lease at their boundary, so before this
+ * reuse a single execute call that invoked one package export paid for two
+ * full leases.
+ *
+ * Frames deactivate when the outer lease releases. Detached work (for
+ * example `waitUntil` callbacks spawned inside `write`) inherits this
+ * AsyncLocalStorage context, and without the active flag it would keep
+ * skipping acquisition after the lease row was already released.
+ */
+type AccountWriteLeaseFrame = { active: boolean }
+const heldAccountWriteLeaseStorage = new AsyncLocalStorage<
+	ReadonlyMap<string, AccountWriteLeaseFrame>
+>()
+
 export async function withAccountWriteLease<T>(input: {
 	db: D1Database
 	stableUserId: string
 	holder?: string
 	write: () => Promise<T>
+}) {
+	const heldLeases = heldAccountWriteLeaseStorage.getStore()
+	if (heldLeases?.get(input.stableUserId)?.active) {
+		return await input.write()
+	}
+	const frame: AccountWriteLeaseFrame = { active: true }
+	const nextHeldLeases = new Map(heldLeases)
+	nextHeldLeases.set(input.stableUserId, frame)
+	try {
+		return await heldAccountWriteLeaseStorage.run(nextHeldLeases, async () =>
+			acquireAccountWriteLeaseAndWrite({ ...input, frame }),
+		)
+	} finally {
+		frame.active = false
+	}
+}
+
+async function acquireAccountWriteLeaseAndWrite<T>(input: {
+	db: D1Database
+	stableUserId: string
+	holder?: string
+	frame: AccountWriteLeaseFrame
+	write: () => Promise<T>
 }) {
 	if (typeof input.db.batch !== 'function') {
 		const acquired = await input.db
@@ -101,6 +146,7 @@ export async function withAccountWriteLease<T>(input: {
 		try {
 			return await input.write()
 		} finally {
+			input.frame.active = false
 			await input.db
 				.prepare(
 					`UPDATE users
@@ -159,6 +205,7 @@ export async function withAccountWriteLease<T>(input: {
 		if (held?.held !== 1) throw new AccountWriteLeaseLostError()
 		return result
 	} finally {
+		input.frame.active = false
 		const releasedAt = utcSqliteTimestamp()
 		const released = await input.db
 			.prepare(

packages/worker/src/entitlements/entitlements.node.test.ts

@@ -31,6 +31,8 @@ import {
 	countRunningPackageServices,
 	estimateEntitlementStorageEntryByteDelta,
 	estimateEntitlementStorageEntryBytes,
+	findCachedUserAccountByStableUserId,
+	getCachedUserPlan,
 	getUserPlan,
 	incrementDailyEntitlementCounter,
 	packageServiceStateStaleMs,
@@ -136,6 +138,21 @@ function createEntitlementsTestDb(
 										: null
 								) as T | null
 							}
+							if (query.includes('SELECT email, plan, email_verified_at')) {
+								const stableUserId = params[0]
+								const user = users.find(
+									(row) => row.stable_user_id === stableUserId,
+								)
+								return (
+									user
+										? {
+												email: user.email,
+												plan: user.plan,
+												email_verified_at: null,
+											}
+										: null
+								) as T | null
+							}
 							if (query.includes('FROM entitlement_daily_counters')) {
 								const row = counters.find(
 									(counter) =>
@@ -374,6 +391,95 @@ test('getUserPlan resolves plans via email+stable id and short-circuits invalid
 	expect(consoleWarn).toHaveBeenCalledWith(unknownStoredPlanWarningTag)
 })
 
+test('getCachedUserPlan caches per db binding and never caches failures', async () => {
+	const userId = await createStableUserIdFromEmail(plannedEmail)
+	const users = [{ email: plannedEmail, plan: 'pro', stable_user_id: userId }]
+	const { db, queries } = createEntitlementsTestDb({ users })
+
+	expect(await getCachedUserPlan(db, { userId, email: plannedEmail })).toBe(
+		'pro',
+	)
+	expect(await getCachedUserPlan(db, { userId, email: plannedEmail })).toBe(
+		'pro',
+	)
+	const planQueries = () =>
+		queries.filter((query) =>
+			query.sql.includes('email = ? AND stable_user_id = ?'),
+		)
+	expect(planQueries()).toHaveLength(1)
+
+	// A plan change is visible to the uncached lookup immediately and to the
+	// cached lookup only after the TTL: quota checks tolerate that staleness.
+	users[0]!.plan = 'free'
+	expect(await getUserPlan(db, { userId, email: plannedEmail })).toBe('free')
+	expect(await getCachedUserPlan(db, { userId, email: plannedEmail })).toBe(
+		'pro',
+	)
+
+	// Another db binding (fresh test database) never shares cache entries.
+	const second = createEntitlementsTestDb({
+		users: [{ email: plannedEmail, plan: 'free', stable_user_id: userId }],
+	})
+	expect(
+		await getCachedUserPlan(second.db, { userId, email: plannedEmail }),
+	).toBe('free')
+
+	// Anonymous / invalid ids short-circuit without touching the cache or D1.
+	expect(await getCachedUserPlan(db, { userId, email: null })).toBe('max')
+	expect(
+		await getCachedUserPlan(db, { userId: 'user-1', email: plannedEmail }),
+	).toBe('max')
+
+	// Failures are not pinned for the TTL: the next call retries D1.
+	let firstCall = true
+	const flaky = {
+		prepare() {
+			return {
+				bind() {
+					return {
+						async first() {
+							if (firstCall) {
+								firstCall = false
+								throw new Error('D1 blip')
+							}
+							return { plan: 'pro', stripe_plan: null }
+						},
+					}
+				},
+			}
+		},
+	} as unknown as D1Database
+	await expect(
+		getCachedUserPlan(flaky, { userId, email: plannedEmail }),
+	).rejects.toThrow('D1 blip')
+	expect(await getCachedUserPlan(flaky, { userId, email: plannedEmail })).toBe(
+		'pro',
+	)
+})
+
+test('findCachedUserAccountByStableUserId caches the account reverse-resolution per db', async () => {
+	const userId = await createStableUserIdFromEmail(plannedEmail)
+	const { db, queries } = createEntitlementsTestDb({
+		users: [{ email: plannedEmail, plan: 'pro', stable_user_id: userId }],
+	})
+	const accountQueries = () =>
+		queries.filter((query) =>
+			query.sql.includes('SELECT email, plan, email_verified_at'),
+		)
+	expect(await findCachedUserAccountByStableUserId(db, userId)).toEqual({
+		email: plannedEmail,
+		plan: 'pro',
+		emailVerified: false,
+	})
+	expect(await findCachedUserAccountByStableUserId(db, userId)).toEqual({
+		email: plannedEmail,
+		plan: 'pro',
+		emailVerified: false,
+	})
+	expect(accountQueries()).toHaveLength(1)
+	expect(await findCachedUserAccountByStableUserId(db, '  ')).toBeNull()
+})
+
 test('assertWithinEntitlement passes under the limit, throws at it, and enforces finite max ordinary limits', async () => {
 	const userId = await createStableUserIdFromEmail(plannedEmail)
 	const freeLimit = planLimits.free.maxScheduledJobs