diff --git a/packages/db/src/client.ts b/packages/db/src/client.ts index ee080641142..1f7c25f6c59 100644 --- a/packages/db/src/client.ts +++ b/packages/db/src/client.ts @@ -311,10 +311,24 @@ export function livenessDb(db: Db): Db { return created; } +const controllerLeaseDbFactories = new WeakMap Db>(); +const controllerLeaseDbs = new WeakMap(); + +export function controllerLeaseDb(db: Db): Db { + const existing = controllerLeaseDbs.get(db); + if (existing) return existing; + const factory = controllerLeaseDbFactories.get(db); + if (!factory) return db; + const created = factory(); + controllerLeaseDbs.set(db, created); + return created; +} + /** Ends the liveness pool if one was created. The ended pool stays cached, so * a late renewal fails like a query on the ended primary pool. */ export async function closeLivenessDb(db: Db, timeoutSeconds: number): Promise { - await livenessDbs.get(db)?.$client.end({ timeout: timeoutSeconds }); + await Promise.all([livenessDbs.get(db)?.$client.end({ timeout: timeoutSeconds }), + controllerLeaseDbs.get(db)?.$client.end({ timeout: timeoutSeconds })]); } // Control checkpoints must not occupy the request or lease-renewal pools @@ -389,6 +403,9 @@ function createDbClient(url: string, resolved: DatabaseClientOptions, owner?: ob livenessDbFactories.set(db, () => createDbClient(url, { ...resolved, maxConnections: LIVENESS_POOL_MAX_CONNECTIONS, applicationName: "paperclip-controller-lease", }, db)); + controllerLeaseDbFactories.set(db, () => createDbClient(url, { + ...resolved, maxConnections: 1, applicationName: "paperclip-legacy-controller-lease", + }, db)); checkpointDbFactories.set(db, () => createDbClient(url, { ...resolved, maxConnections: CHECKPOINT_POOL_MAX_CONNECTIONS, applicationName: "paperclip-runner-checkpoint", }, db)); diff --git a/packages/db/src/index.ts b/packages/db/src/index.ts index c513660e99e..feb4848b43e 100644 --- a/packages/db/src/index.ts +++ b/packages/db/src/index.ts @@ -4,6 +4,7 @@ export { createOwnedDatabaseListener, type OwnedDatabaseListener, livenessDb, + controllerLeaseDb, closeLivenessDb, checkpointDb, closeCheckpointDb, diff --git a/server/src/services/legacy-controller-lease.test.ts b/server/src/services/legacy-controller-lease.test.ts index 8d33b23cb42..0ff1028f38a 100644 --- a/server/src/services/legacy-controller-lease.test.ts +++ b/server/src/services/legacy-controller-lease.test.ts @@ -2,7 +2,7 @@ import { randomUUID } from "node:crypto"; 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 { agents, closeLivenessDb, companies, controllerLeaseDb, createDb, heartbeatRuns, livenessDb, subscribeDatabaseWork } from "@paperclipai/db"; import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../__tests__/helpers/embedded-postgres.js"; import { resolveExternalTestDatabaseUrl } from "../__tests__/helpers/external-test-database.js"; import { DELIVERY_QUEUES } from "./delivery-work-notifications.js"; @@ -122,7 +122,7 @@ const support = external ? { supported: true } : await getEmbeddedPostgresTestSu DELIVERY_QUEUES.activeExecutionRecovery, DELIVERY_QUEUES.settledNativeProcesses, DELIVERY_QUEUES.strandedAgentStatus, DELIVERY_QUEUES.activeRunSilence]; const stops = queues.map(queue => subscribeDatabaseWork(db, queue, event => { if (event === "settled") settled.push(queue); })); - const pool = livenessDb(db), prior = pool.$client.options.debug; + const pool = controllerLeaseDb(db), prior = pool.$client.options.debug; await pool.execute(sql`select 1`); pool.$client.options.debug = (_connection, query) => { queries.push(query); }; const watch = watchLegacyControllerLease(db, run, new AbortController()); @@ -224,19 +224,42 @@ const support = external ? { supported: true } : await getEmbeddedPostgresTestSu it("renews on one private pool per client that shutdown can end", async () => { const run = await seed(); const owner = createDb(external ?? database.connectionString, { maxConnections: 6 }); - const pool = livenessDb(owner); + const pool = controllerLeaseDb(owner); expect(pool).not.toBe(owner); - expect(livenessDb(owner)).toBe(pool); + expect(pool).not.toBe(livenessDb(owner)); + expect(controllerLeaseDb(owner)).toBe(pool); const watch = watchLegacyControllerLease(owner, run, new AbortController()); try { await watch.assertOwned(); const [{ name }] = await pool.execute<{ name: string }>(sql`select current_setting('application_name') as name`); - expect(name).toBe("paperclip-controller-lease"); + expect(name).toBe("paperclip-legacy-controller-lease"); } finally { watch.stop(); } await closeLivenessDb(owner, 1); + await expect(controllerLeaseDb(owner).execute(sql`select 1`)).rejects.toThrow(); await expect(livenessDb(owner).execute(sql`select 1`)).rejects.toThrow(); }); + it("renews without queueing behind native lease transactions that hold every liveness connection", async () => { + const run = await seed(); + const owner = createDb(external ?? database.connectionString, { maxConnections: 6 }); + const nativeRenewals = Array.from({ length: 4 }, () => + livenessDb(owner).transaction(async tx => { await tx.execute(sql`select pg_sleep(5)`); })); + const controller = new AbortController(); + const watch = watchLegacyControllerLease(owner, run, controller); + try { + await new Promise(resolve => setTimeout(resolve, 200)); + const startedAt = Date.now(); + await watch.assertOwned("dispatching"); + expect(Date.now() - startedAt).toBeLessThan(2_000); + expect(controller.signal.aborted).toBe(false); + } finally { + watch.stop(); + await Promise.all(nativeRenewals); + await closeLivenessDb(owner, 1); + await owner.$client.end({ timeout: 1 }); + } + }, 30_000); + it("retries a renewal that times out waiting on the run row instead of losing the lease", async () => { const run = await seed(); let lockTaken!: () => void, releaseLock!: () => void; @@ -265,7 +288,7 @@ const support = external ? { supported: true } : await getEmbeddedPostgresTestSu 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()`); + where application_name = 'paperclip-legacy-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); @@ -277,7 +300,7 @@ const support = external ? { supported: true } : await getEmbeddedPostgresTestSu 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; + const pool = controllerLeaseDb(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()); @@ -525,6 +548,20 @@ describe("legacy controller lease loss attribution", () => { } finally { watches.forEach(watch => watch.stop()); vi.useRealTimers(); } }); + it("keeps a controller owned when every renewal waits thirty-five seconds for a connection before its statement runs", async () => { + vi.useFakeTimers(); + const run = watchedRun(LEGACY_CONTROLLER_LEASE_MS); + const pool = slowPoolDb([run.id], { connections: 1, transactionMs: 35_000 }); + const controller = new AbortController(); + const watch = watchLegacyControllerLease(pool.db, run, controller); + try { + await vi.advanceTimersByTimeAsync(4 * LEGACY_CONTROLLER_LEASE_MS); + expect(controller.signal.reason?.leaseLossTiming).toBeUndefined(); + expect(controller.signal.aborted).toBe(false); + expect(pool.transactions()).toBeGreaterThanOrEqual(5); + } finally { 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 a6c39c31118..3565eef9a7d 100644 --- a/server/src/services/legacy-controller-lease.ts +++ b/server/src/services/legacy-controller-lease.ts @@ -1,6 +1,6 @@ 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 { controllerLeaseDb, heartbeatRuns, type Db } from "@paperclipai/db"; import type { LegacyControllerLeaseLossCause, LegacyControllerLeaseLossTiming } from "./run-cancellation.js"; import { randomUUID } from "node:crypto"; import { @@ -110,19 +110,20 @@ function isStatementTimeout(error: unknown): boolean { return false; } +type RenewalResult = { outcome: RenewalOutcome; issuedAt: number }; type RenewalRequest = { target: RenewalTarget; statementTimeoutMs: number; - resolve(outcome: RenewalOutcome): void; + resolve(result: RenewalResult): void; reject(error: unknown): void; }; type RenewalBatch = { pending: RenewalRequest[]; flushing: boolean }; const renewalBatches = new WeakMap(); -function renewThroughBatch(db: Db, target: RenewalTarget, statementTimeoutMs: number): Promise { +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) => { + const outcome = new Promise((resolve, reject) => { batch.pending.push({ target, statementTimeoutMs, resolve, reject }); }); if (!batch.flushing) void flushRenewals(db, batch); @@ -142,13 +143,15 @@ async function flushRenewals(db: Db, batch: RenewalBatch) { } try { const statementTimeoutMs = Math.min(...requests.map(request => request.statementTimeoutMs)); - const outcomes = await livenessDb(db).transaction(async (tx) => { + let issuedAt = 0; + const outcomes = await controllerLeaseDb(db).transaction(async (tx) => { + issuedAt = Date.now(); 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"); + request.resolve({ outcome: sameCompany ? outcomes.get(request.target.id) ?? "rejected" : "rejected", issuedAt }); } } catch (error) { for (const request of requests) request.reject(error); @@ -351,13 +354,13 @@ export function watchLegacyControllerLease( // deadline is measured from that attempt rather than the first one. // Failed attempts are retried until the deadline aborts the controller: // the UPDATE is idempotent and still requires this controller's boot ID. - let attemptStartedAt = startedAt; + let renewedFrom = startedAt; const renew = async () => { for (;;) { - attemptStartedAt = Date.now(); try { - const outcome = await renewThroughBatch(db, { id: run.id, companyId: run.companyId, stage }, + const { outcome, issuedAt } = await renewThroughBatch(db, { id: run.id, companyId: run.companyId, stage }, renewalStatementTimeoutMs); + renewedFrom = issuedAt; 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; @@ -389,7 +392,7 @@ export function watchLegacyControllerLease( lastRenewedAt = Date.now(); renewalCount += 1; renewalFailed = false; - armDeadline(Math.max(0, LEGACY_CONTROLLER_LEASE_MS - (Date.now() - attemptStartedAt))); + armDeadline(Math.max(0, LEGACY_CONTROLLER_LEASE_MS - (Date.now() - renewedFrom))); } }; const timer = setInterval(() => {