From 612cc7e3e6b3ed33ab9b5346bb4b416dd66577ed Mon Sep 17 00:00:00 2001 From: Dotta <34892728+cryppadotta@users.noreply.github.com> Date: Wed, 11 Nov 2026 15:57:20 -0600 Subject: [PATCH] fix(heartbeat): renew every legacy controller lease in one shared statement (#20924) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Thinking Path > - Paperclip is the open source app people use to manage AI agents for work > - A legacy run holds a durable controller lease while it prepares. The server renews the lease every 10 seconds, and it stops the run if no renewal succeeds within the 60 second lease. > - Each run renewed in its own transaction on a private pool of two connections. Each renewal used six round trips (begin, statement timeout, transaction id, notify, update, commit) and woke live-run recovery for that run. > - When many runs prepare at the same time on a loaded database, the renewals queue behind each other. The queue grows with the number of runs, and every run loses its lease at about the same time. > - A healthy server must not fail its own runs because it has many of them. > - This pull request renews all leases of one server process through one shared queue. Requests that arrive while a renewal is in flight share the next statement. > - The benefit is that the renewal cost no longer grows with the number of live runs, so a burst cannot starve the lease. ## Linked Issues or Issue Description **What happened?** On a staging server that did not restart, 22 runs failed with `controller_lease_lost` in about 13 minutes. Every run was a legacy run in the `preparing` stage with `providerWorkStarted: false`. Every loss had cause `renewal_deadline`. The renewal had been pending 28 to 48 seconds (`renewalPendingMs`). The deadline callback ran on time (`deadlineLateMs` 3 to 861 ms), so the event loop was not blocked. The successful renewals before the loss had each taken 20 to 30 seconds, because the deadline fired only 25 to 48 seconds after the last renewal (`sinceRenewalMs`). About 30 runs were live in the process. The first wave started when a user retried about 14 blocked tasks at once. The automatic retries of the first wave then failed the same way, which made the second wave. 12 of 16 affected tasks became blocked. **Steps to reproduce** 1. Start about 30 legacy runs in one server process. 2. Make each renewal transaction slow (for example, a loaded database). 3. Wait one lease period. **Expected behavior** The leases stay owned. The renewal cost does not grow with the number of runs. **Actual behavior** All runs fail with `controller_lease_lost` (`renewal_deadline`). ## What Changed - `server/src/services/legacy-controller-lease.ts`: All controllers of one process renew through one queue per database. One flush is in flight at a time. Requests that arrive during a flush share the next flush, so a renewal waits behind at most one other statement. A flush is one transaction on the private lease pool: `statement_timeout` and one statement that renews all requested rows. - The statement locks the requested rows with `for update skip locked`. A row that another transaction holds is reported busy, and only that run retries after 250 ms. It does not stall the other runs in the batch. A row that is not this boot's live legacy lease is reported rejected, as before. - Renewal no longer emits the `active-execution-recovery` notification. A renewal only moves the deadline later. Recovery already has a deadline at the earlier expiry and rereads the moved deadline then. This removes one recovery pass, `pg_current_xact_id` and `pg_notify` per renewal. - The lease contract does not change: 60 second lease, 10 second renewal, boot-id fence, expired leases cannot renew, the local deadline is measured from the attempt start, and the loss causes and timings are the same. - `doc/run-log-events.md`: Documents the shared renewal queue. Upgrade story: no persisted shape changes. The row columns, lease length and boot-id fence are the same. A lease held by the previous release has a different boot id, so the new batch rejects it and leaves it untouched. A test covers a previous-release lease in the same batch. ## Verification - Red on `origin/master`, green with this change: - `keeps thirty preparing controllers owned when every renewal takes four seconds on a two-connection pool`: before, all 30 controllers lost the lease with `renewal_deadline` (the staging shape). After, none lose it. - `renews fifteen concurrent controllers in one statement` (real Postgres): before, 15 update statements. After, 1. - `renews through one private statement without waking live-run recovery`: before, the renewal woke `active-execution-recovery`. After, it does not. - `live-recovery-work.test.ts` `renews without waking recovery and re-reads the moved controller deadline at the old one`: before, 2 recovery passes per renewal. After, none, and recovery still moves to the new deadline at the old one. - Safety cases: a row lock on one run does not delay another run; a previous-release lease in the same batch stays untouched; a terminated lease connection is retried; existing deadline, retry, rejection and revocation tests pass. - `npx vitest run src/services/legacy-controller-lease.test.ts src/__tests__/live-recovery-work.test.ts`: 88 passed. - `pnpm verify`: see the CI checks on this head. ## Risks - Medium-low. All renewals of a process now share one statement. A failed flush fails every request in it; each watcher retries independently, as it did before. - A busy row is retried every 250 ms until the lease deadline. Before, the renewal waited on the row lock for up to the statement timeout. The loss behavior is the same. - Recovery wakes at the old lease expiry instead of after each renewal. It finds the live lease and moves to the new expiry. - Not in scope: the downstream workspace recovery that blocked tasks after a lost lease (`legacy_execution_requires_reconciliation`, `source_unavailable`). > For core feature work, check [`ROADMAP.md`](ROADMAP.md) first and discuss it in `#dev` before opening the PR. Feature PRs that overlap with planned core work may need to be redirected — check the roadmap first. See `CONTRIBUTING.md`. ## Model Used - Claude Opus 5.5 (claude-opus-5-5), Anthropic, with tool use and code execution in Claude Code. ## Checklist - [x] I have included a thinking path that traces from project context to this change - [x] I have specified the model used (with version and capability details) - [x] I have checked ROADMAP.md and confirmed this PR does not duplicate planned core work - [x] I have searched GitHub for duplicate or related PRs and linked them above - [x] I have either (a) linked existing issues with `Fixes: #` / `Closes #` / `Refs #` OR (b) described the issue in-PR following the relevant issue template - [x] I have not referenced internal/instance-local Paperclip issues or links (only public GitHub `#NNN` / `github.com/paperclipai/paperclip` URLs) - [x] My branch name describes the change (e.g. `docs/...`, `fix/...`) and contains no internal Paperclip ticket id or instance-derived details - [x] I have run tests locally and they pass - [x] I have added or updated tests where applicable - [x] I have updated relevant documentation to reflect my changes - [x] I have considered and documented any risks above - [ ] All Paperclip CI gates are green - [ ] Greptile is 5/5 with no open P2s, recommendations, or follow-ups - [x] I will address all Greptile and reviewer comments before requesting merge 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 5.5 --- doc/run-log-events.md | 8 + .../src/__tests__/live-recovery-work.test.ts | 4 +- .../services/legacy-controller-lease.test.ts | 142 ++++++++++++++---- .../src/services/legacy-controller-lease.ts | 118 +++++++++++---- 4 files changed, 214 insertions(+), 58 deletions(-) diff --git a/doc/run-log-events.md b/doc/run-log-events.md index cbf3b62079b..5269e61926d 100644 --- a/doc/run-log-events.md +++ b/doc/run-log-events.md @@ -786,6 +786,14 @@ after it lost the durable controller lease; the reason names a closed cause (`renewal_rejected`, `renewal_deadline`, or `renewal_error`). The watchdog retries a failed renewal until the lease deadline, so `renewal_error` means the attempts still failed with an error when the deadline passed. +All legacy controllers in one server process renew through a single queue on +the private lease connection pool. Renewals that arrive while one is in flight +share the next statement, so a burst of runs costs one renewal statement per +round instead of one transaction per run, and a renewal waits behind at most +one other. A run whose row another transaction holds is skipped and retried +shortly, without delaying the other runs. A renewal only moves the lease +deadline later, so it does not wake live-run recovery; recovery rereads the +moved deadline when the earlier one comes due. A lease lost during run setup (for example a rejected renewal while preparing the workspace) settles the run `failed` with error code `controller_lease_lost`, the same `controller_lease` cancellation evidence, and diff --git a/server/src/__tests__/live-recovery-work.test.ts b/server/src/__tests__/live-recovery-work.test.ts index 4bdacaf6c7f..1e78c73c160 100644 --- a/server/src/__tests__/live-recovery-work.test.ts +++ b/server/src/__tests__/live-recovery-work.test.ts @@ -560,7 +560,7 @@ describe("residual recovery production matrix", () => { eq(heartbeatRunEvents.eventType, "native.local_process_stopped")))).toEqual([]); await quiet(); }); - it("renews only the active recovery owner and re-reads the moved controller deadline", async () => { + it("renews without waking recovery and re-reads the moved controller deadline at the old one", async () => { const f = await fixture(), old = new Date(Date.now() - 6 * 60_000), expiry = new Date(Date.now() + 20_000); const [run] = await db.insert(heartbeatRuns).values({ companyId: f.companyId, agentId: f.agent.id, runtimeMode: "legacy", status: "running", controllerBootId: legacyControllerBootId, @@ -573,7 +573,7 @@ describe("residual recovery production matrix", () => { expect(moved.getTime()).toBeGreaterThan(expiry.getTime()); await advance(1_000); expect(timeline.slice(before).filter(event => Array.isArray(event) && event[0] === "pass") - .map(event => (event as unknown[])[1])).toEqual(["active-execution-recovery", "active-execution-recovery"]); + .map(event => (event as unknown[])[1])).toEqual([]); await advance(expiry.getTime() - Date.now()); expect((await getRun(run.id)).status).toBe("running"); expect(deadlines.get("active-execution-recovery")!()).toBe(moved.getTime()); diff --git a/server/src/services/legacy-controller-lease.test.ts b/server/src/services/legacy-controller-lease.test.ts index 283b345ff50..8d33b23cb42 100644 --- a/server/src/services/legacy-controller-lease.test.ts +++ b/server/src/services/legacy-controller-lease.test.ts @@ -1,5 +1,6 @@ import { randomUUID } from "node:crypto"; -import { and, eq, ne, sql } from "drizzle-orm"; +import { and, eq, ne, sql, type SQL } from "drizzle-orm"; +import { PgDialect } from "drizzle-orm/pg-core"; import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; import { agents, closeLivenessDb, companies, createDb, heartbeatRuns, livenessDb, subscribeDatabaseWork } from "@paperclipai/db"; import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../__tests__/helpers/embedded-postgres.js"; @@ -28,13 +29,40 @@ import { createOwnershipReportState, runOwnershipReconciliationPass } from "./re import { OWNABLE_KINDS } from "./recovery/ownership-registry.js"; type WatchArgs = Parameters; +const dialect = new PgDialect(); // A renewal test double: each attempt runs in a transaction on the liveness pool. function fakeRenewalDb(returning: () => Promise) { - const fake = { update: () => ({ set: () => ({ where: () => ({ returning }) }) }), - execute: async () => [], transaction: (action: (tx: unknown) => Promise) => action(fake) }; + const execute = async (query: SQL) => { + const { sql: text, params } = dialect.sqlToQuery(query); + if (/set_config/.test(text)) return []; + const rows = await returning(); + const renewed = Array.isArray(rows) && rows.length > 0; + return params.filter(param => typeof param === "string").map(id => ({ id, renewed, examined: true })); + }; + const fake = { execute, transaction: (action: (tx: unknown) => Promise) => action(fake) }; return fake as unknown as WatchArgs[0]; } +function slowPoolDb(runIds: string[], { connections, transactionMs }: { connections: number; transactionMs: number }) { + let active = 0, transactions = 0; + const waiting: Array<() => void> = []; + const tx = { + execute: async () => runIds.map(id => ({ id, renewed: true, examined: true })), + update: () => ({ set: () => ({ where: () => ({ returning: async () => [{ id: runIds[0] }] }) }) }), + }; + const fake = { + transaction: async (action: (tx: unknown) => Promise) => { + while (active >= connections) await new Promise(resolve => waiting.push(resolve)); + active += 1; transactions += 1; + try { + await new Promise(resolve => setTimeout(resolve, transactionMs)); + return await action(tx); + } finally { active -= 1; waiting.shift()?.(); } + }, + }; + return { db: fake as unknown as WatchArgs[0], transactions: () => transactions }; +} + const external = resolveExternalTestDatabaseUrl("PAPERCLIP_LEGACY_CONTROLLER_TEST_DATABASE_URL"); const support = external ? { supported: true } : await getEmbeddedPostgresTestSupport(); (support.supported ? describe : describe.skip)("durable legacy controller ownership", () => { @@ -88,7 +116,7 @@ const support = external ? { supported: true } : await getEmbeddedPostgresTestSu expect(saved.executionStage).toBe("dispatching"); expect(await revokeExpiredLegacyController(db, run)).toBe(false); }); - it("renews through one private transaction with only the active execution topic", async () => { + it("renews through one private statement without waking live-run recovery", async () => { const run = await seed(), settled: string[] = [], queries: string[] = []; const queues = [DELIVERY_QUEUES.nativeProcessSessions, DELIVERY_QUEUES.nativeQuestionCancellations, DELIVERY_QUEUES.activeExecutionRecovery, DELIVERY_QUEUES.settledNativeProcesses, @@ -100,12 +128,12 @@ const support = external ? { supported: true } : await getEmbeddedPostgresTestSu const watch = watchLegacyControllerLease(db, run, new AbortController()); try { await watch.assertOwned(); - expect(settled).toEqual([DELIVERY_QUEUES.activeExecutionRecovery]); + expect(settled).toEqual([]); expect(queries.filter(query => /savepoint/i.test(query))).toEqual([]); - expect(queries.filter(query => /pg_current_xact_id/.test(query))).toHaveLength(1); + expect(queries.filter(query => /pg_current_xact_id|pg_notify/.test(query))).toEqual([]); expect(queries.filter(query => /set_config.*statement_timeout/.test(query))).toHaveLength(1); - expect(queries.filter(query => /^update /i.test(query))).toHaveLength(1); - expect(queries).toHaveLength(5); + expect(queries.filter(query => /update "heartbeat_runs"/i.test(query))).toHaveLength(1); + expect(queries).toHaveLength(4); } finally { watch.stop(); pool.$client.options.debug = prior; stops.forEach(stop => stop()); } }); it("an expired controller cannot renew or dispatch even before a reaper claims it", async () => { @@ -232,36 +260,82 @@ const support = external ? { supported: true } : await getEmbeddedPostgresTestSu it("retries a renewal whose connection was terminated instead of losing the lease", async () => { const run = await seed(); + const controller = new AbortController(); + const watch = watchLegacyControllerLease(db, run, controller); + try { + await watch.assertOwned(); + const terminated = await db.execute<{ pid: number }>(sql`select pg_terminate_backend(pid) as pid from pg_stat_activity + where application_name = 'paperclip-controller-lease' and pid <> pg_backend_pid()`); + expect(terminated.length).toBeGreaterThan(0); + await expect(watch.assertOwned("dispatching")).resolves.toBeUndefined(); + expect(controller.signal.aborted).toBe(false); + const [saved] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, run.id)); + expect(saved).toMatchObject({ controllerBootId: legacyControllerBootId, executionStage: "dispatching" }); + } finally { watch.stop(); } + }, 15_000); + + it("renews fifteen concurrent controllers in one statement", async () => { + const runs = await Promise.all(Array.from({ length: 15 }, () => seed())); + const queries: string[] = []; + const pool = livenessDb(db), prior = pool.$client.options.debug; + await pool.execute(sql`select 1`); + pool.$client.options.debug = (_connection, query) => { queries.push(query); }; + const controllers = runs.map(() => new AbortController()); + const watches = runs.map((run, index) => watchLegacyControllerLease(db, run, controllers[index]!)); + try { + await Promise.all(watches.map(watch => watch.assertOwned())); + expect(queries.filter(query => /update "heartbeat_runs"/i.test(query))).toHaveLength(1); + expect(controllers.filter(controller => controller.signal.aborted)).toEqual([]); + for (const run of runs) { + const [saved] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, run.id)); + expect(saved.controllerLeaseExpiresAt!.getTime()).toBeGreaterThan(run.controllerLeaseExpiresAt!.getTime()); + } + } finally { watches.forEach(watch => watch.stop()); pool.$client.options.debug = prior; } + }); + + it("does not stall other controllers behind a run row another transaction holds", async () => { + const [held, free] = [await seed(), await seed()]; let lockTaken!: () => void, releaseLock!: () => void; const taken = new Promise((resolve) => { lockTaken = resolve; }); const released = new Promise((resolve) => { releaseLock = resolve; }); const holder = db.transaction(async (tx) => { - await tx.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(eq(heartbeatRuns.id, run.id)).for("update"); + await tx.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(eq(heartbeatRuns.id, held.id)).for("update"); lockTaken(); await released; }); await taken; - const controller = new AbortController(); - const watch = watchLegacyControllerLease(db, run, controller); + const heldWatch = watchLegacyControllerLease(db, held, new AbortController()); + const freeWatch = watchLegacyControllerLease(db, free, new AbortController()); try { - const renewal = watch.assertOwned(); - // Drop the renewal's own connection while it waits on the run row. - let blocked: { pid: number } | undefined; - for (let attempt = 0; !blocked && attempt < 100; attempt++) { - [blocked] = await db.execute<{ pid: number }>(sql`select pid from pg_stat_activity - where application_name = 'paperclip-controller-lease' and wait_event_type = 'Lock'`); - if (!blocked) await new Promise((resolve) => setTimeout(resolve, 20)); - } - expect(blocked).toBeDefined(); - await db.execute(sql`select pg_terminate_backend(${blocked!.pid})`); - setTimeout(releaseLock, LEGACY_CONTROLLER_RENEWAL_RETRY_DELAY_MS + 500); - await expect(renewal).resolves.toBeUndefined(); - expect(controller.signal.aborted).toBe(false); - const [saved] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, run.id)); - expect(saved.controllerBootId).toBe(legacyControllerBootId); - } finally { releaseLock(); watch.stop(); await holder; } + const startedAt = Date.now(); + const heldRenewal = heldWatch.assertOwned(); + await freeWatch.assertOwned(); + expect(Date.now() - startedAt).toBeLessThan(1_000); + setTimeout(releaseLock, 500); + await expect(heldRenewal).resolves.toBeUndefined(); + } finally { releaseLock(); heldWatch.stop(); freeWatch.stop(); await holder; } }, 15_000); + it("leaves a lease held by the previous release untouched while renewing its own in the same batch", async () => { + const [own, previous] = [await seed(), await seed()]; + await db.update(heartbeatRuns).set({ controllerBootId: randomUUID(), controllerProcess: null }) + .where(eq(heartbeatRuns.id, previous.id)); + const [before] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, previous.id)); + const ownController = new AbortController(), previousController = new AbortController(); + const ownWatch = watchLegacyControllerLease(db, own, ownController); + const previousWatch = watchLegacyControllerLease(db, previous, previousController); + try { + const [ownResult, previousResult] = await Promise.allSettled([ownWatch.assertOwned(), previousWatch.assertOwned("dispatching")]); + expect(ownResult.status).toBe("fulfilled"); + expect(previousResult.status).toBe("rejected"); + expect(ownController.signal.aborted).toBe(false); + expect(previousController.signal.reason.leaseLossCause).toBe("renewal_rejected"); + const [after] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, previous.id)); + expect(after).toMatchObject({ controllerBootId: before.controllerBootId, executionStage: before.executionStage, + controllerLeaseExpiresAt: before.controllerLeaseExpiresAt }); + } finally { ownWatch.stop(); previousWatch.stop(); } + }); + it("leaves native controller ownership to the native coordinator", () => { expect(legacyControllerClaim("native")).toEqual({}); }); @@ -437,6 +511,20 @@ describe("legacy controller lease loss attribution", () => { } finally { watch.stop(); vi.useRealTimers(); } }); + it("keeps thirty preparing controllers owned when every renewal takes four seconds on a two-connection pool", async () => { + vi.useFakeTimers(); + const runs = Array.from({ length: 30 }, () => watchedRun(LEGACY_CONTROLLER_LEASE_MS)); + const pool = slowPoolDb(runs.map(run => run.id), { connections: 2, transactionMs: 4_000 }); + const controllers = runs.map(() => new AbortController()); + const watches = runs.map((run, index) => watchLegacyControllerLease(pool.db, run, controllers[index]!)); + try { + await vi.advanceTimersByTimeAsync(3 * LEGACY_CONTROLLER_LEASE_MS); + expect(controllers.filter(controller => controller.signal.aborted).map(controller => controller.signal.reason?.leaseLossCause)) + .toEqual([]); + expect(pool.transactions()).toBeLessThanOrEqual(2 * (3 * LEGACY_CONTROLLER_LEASE_MS / LEGACY_CONTROLLER_RENEW_MS)); + } finally { watches.forEach(watch => watch.stop()); vi.useRealTimers(); } + }); + const watchedRun = (leaseMs: number) => ({ id: randomUUID(), companyId: randomUUID(), runtimeMode: "legacy", controllerBootId: legacyControllerBootId, controllerLeaseExpiresAt: new Date(Date.now() + leaseMs) }) as WatchArgs[1]; diff --git a/server/src/services/legacy-controller-lease.ts b/server/src/services/legacy-controller-lease.ts index 8eab71e3bd9..a6c39c31118 100644 --- a/server/src/services/legacy-controller-lease.ts +++ b/server/src/services/legacy-controller-lease.ts @@ -1,4 +1,4 @@ -import { notifyActiveExecutionRecovery, notifyLiveRunRecovery } from "./live-run-recovery-notifications.js"; +import { notifyLiveRunRecovery } from "./live-run-recovery-notifications.js"; import { and, eq, gt, lte, sql } from "drizzle-orm"; import { heartbeatRuns, livenessDb, type Db } from "@paperclipai/db"; import type { LegacyControllerLeaseLossCause, LegacyControllerLeaseLossTiming } from "./run-cancellation.js"; @@ -48,22 +48,39 @@ export function legacyControllerClaim(runtimeMode: string) { }; } -async function renewLegacyControllerLeaseInTransaction( +type RenewalTarget = Pick & { stage?: "dispatching" }; +type RenewalOutcome = "renewed" | "rejected" | "busy"; + +async function renewLegacyControllerLeasesInTransaction( tx: Db, - run: Pick, - stage?: "dispatching", -): Promise { - await notifyActiveExecutionRecovery(tx, { companyId: run.companyId, runId: run.id }); - const [renewed] = await tx.update(heartbeatRuns).set({ - controllerLeaseExpiresAt: sql`clock_timestamp() + interval '60 seconds'`, - ...(stage ? { executionStage: stage } : {}), - }).where(and( - eq(heartbeatRuns.id, run.id), eq(heartbeatRuns.companyId, run.companyId), - eq(heartbeatRuns.runtimeMode, "legacy"), eq(heartbeatRuns.status, "running"), - eq(heartbeatRuns.controllerBootId, legacyControllerBootId), - gt(heartbeatRuns.controllerLeaseExpiresAt, sql`clock_timestamp()`), - )).returning({ id: heartbeatRuns.id }); - return Boolean(renewed); + targets: RenewalTarget[], +): Promise> { + const requested = sql.join(targets.map(target => sql`(${target.id}::uuid, ${target.companyId}::uuid)`), sql`, `); + const dispatching = targets.filter(target => target.stage === "dispatching").map(target => sql`${target.id}::uuid`); + const stage = dispatching.length + ? sql`case when r.id in (${sql.join(dispatching, sql`, `)}) then 'dispatching' else r.execution_stage end` + : sql`r.execution_stage`; + const rows = await tx.execute<{ id: string; renewed: boolean; examined: boolean }>(sql` + with requested(id, company_id) as (values ${requested}), + examined as ( + select r.id from ${heartbeatRuns} r join requested q on r.id = q.id and r.company_id = q.company_id + for update of r skip locked + ), + renewed as ( + update ${heartbeatRuns} r set controller_lease_expires_at = clock_timestamp() + interval '60 seconds', + execution_stage = ${stage} + from examined e + where r.id = e.id and r.runtime_mode = 'legacy' and r.status = 'running' + and r.controller_boot_id = ${legacyControllerBootId} + and r.controller_lease_expires_at > clock_timestamp() + returning r.id + ) + select q.id::text as id, exists (select 1 from renewed n where n.id = q.id) as renewed, + exists (select 1 from examined e where e.id = q.id) as examined + from requested q`); + const outcomes = new Map(); + for (const row of rows) outcomes.set(row.id, row.renewed ? "renewed" : row.examined ? "rejected" : "busy"); + return outcomes; } export async function renewLegacyControllerLease( @@ -71,7 +88,9 @@ export async function renewLegacyControllerLease( run: Pick, stage?: "dispatching", ): Promise { - return db.transaction(tx => renewLegacyControllerLeaseInTransaction(tx as unknown as Db, run, stage)); + const outcomes = await db.transaction(tx => renewLegacyControllerLeasesInTransaction(tx as unknown as Db, + [{ id: run.id, companyId: run.companyId, stage }])); + return outcomes.get(run.id) === "renewed"; } /** One attempt may wait at most this long; a timed-out attempt is retried @@ -82,6 +101,8 @@ export const LEGACY_CONTROLLER_RENEWAL_STATEMENT_TIMEOUT_MS = 10_000; * connection is retried without spinning against an unavailable database. */ export const LEGACY_CONTROLLER_RENEWAL_RETRY_DELAY_MS = 1_000; +export const LEGACY_CONTROLLER_RENEWAL_BUSY_RETRY_DELAY_MS = 250; + function isStatementTimeout(error: unknown): boolean { for (let cause = error, depth = 0; cause && depth < 4; cause = (cause as { cause?: unknown }).cause, depth++) { if ((cause as { code?: unknown }).code === "57014") return true; @@ -89,18 +110,53 @@ function isStatementTimeout(error: unknown): boolean { return false; } -/** Renew on the private liveness pool so the write cannot queue behind an - * unrelated slow statement on a busy shared connection. */ -async function renewOnLivenessPool( - db: Db, - run: Pick, - stage: "dispatching" | undefined, - statementTimeoutMs: number, -): Promise { - return livenessDb(db).transaction(async (tx) => { - await tx.execute(sql`select set_config('statement_timeout', ${`${statementTimeoutMs}ms`}, true)`); - return renewLegacyControllerLeaseInTransaction(tx as unknown as Db, run, stage); +type RenewalRequest = { + target: RenewalTarget; + statementTimeoutMs: number; + resolve(outcome: RenewalOutcome): void; + reject(error: unknown): void; +}; +type RenewalBatch = { pending: RenewalRequest[]; flushing: boolean }; +const renewalBatches = new WeakMap(); + +function renewThroughBatch(db: Db, target: RenewalTarget, statementTimeoutMs: number): Promise { + let batch = renewalBatches.get(db); + if (!batch) renewalBatches.set(db, batch = { pending: [], flushing: false }); + const outcome = new Promise((resolve, reject) => { + batch.pending.push({ target, statementTimeoutMs, resolve, reject }); }); + if (!batch.flushing) void flushRenewals(db, batch); + return outcome; +} + +async function flushRenewals(db: Db, batch: RenewalBatch) { + batch.flushing = true; + try { + await Promise.resolve(); + while (batch.pending.length) { + const requests = batch.pending.splice(0); + const targets = new Map(); + for (const { target } of requests) { + const existing = targets.get(target.id); + targets.set(target.id, existing ? { ...existing, stage: existing.stage ?? target.stage } : target); + } + try { + const statementTimeoutMs = Math.min(...requests.map(request => request.statementTimeoutMs)); + const outcomes = await livenessDb(db).transaction(async (tx) => { + await tx.execute(sql`select set_config('statement_timeout', ${`${statementTimeoutMs}ms`}, true)`); + return renewLegacyControllerLeasesInTransaction(tx as unknown as Db, [...targets.values()]); + }); + for (const request of requests) { + const sameCompany = targets.get(request.target.id)!.companyId === request.target.companyId; + request.resolve(sameCompany ? outcomes.get(request.target.id) ?? "rejected" : "rejected"); + } + } catch (error) { + for (const request of requests) request.reject(error); + } + } + } finally { + batch.flushing = false; + } } export async function liveLegacyControllerExpiresAt(db: Db, run: Run): Promise { @@ -300,7 +356,11 @@ export function watchLegacyControllerLease( for (;;) { attemptStartedAt = Date.now(); try { - return await renewOnLivenessPool(db, run, stage, renewalStatementTimeoutMs); + const outcome = await renewThroughBatch(db, { id: run.id, companyId: run.companyId, stage }, + renewalStatementTimeoutMs); + if (outcome !== "busy") return outcome === "renewed"; + await new Promise((resolve) => setTimeout(resolve, LEGACY_CONTROLLER_RENEWAL_BUSY_RETRY_DELAY_MS)); + if (stopped || controller.signal.aborted) return false; } catch (error) { if (stopped || controller.signal.aborted) throw error; // A statement timeout is a stall; it stays a `renewal_deadline` loss.