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.