diff --git a/docs/deploy/database.md b/docs/deploy/database.md index 5f25489d8eb..4f72b70c097 100644 --- a/docs/deploy/database.md +++ b/docs/deploy/database.md @@ -77,7 +77,7 @@ The server opens one postgres.js pool for its own queries (and a second one when | `DATABASE_APPLICATION_NAME` | `paperclip` | Value of `application_name` in `pg_stat_activity`, so you can find Paperclip's backends: `SELECT * FROM pg_stat_activity WHERE application_name = 'paperclip';` | | `DATABASE_PREPARED_STATEMENTS` | `true` (driver) | Set `false` behind a transaction-mode pooler (see above). | -The server ends its pools during shutdown (SIGINT/SIGTERM) and when startup fails after the pool was opened, so a restarting server does not leave idle backends behind. Size `max_connections` on the PostgreSQL side for at least `DATABASE_POOL_MAX` per server process plus your other clients. +The server ends its pools during shutdown (SIGINT/SIGTERM) and when startup fails after the pool was opened, so a restarting server does not leave idle backends behind. Size `max_connections` on the PostgreSQL side for at least `DATABASE_POOL_MAX` per server process plus your other clients. Each server process also opens up to 2 connections (`application_name=paperclip-controller-lease`) for agent-run lease renewals, so they never queue behind slow application queries; count those too. ## Switching Between Modes diff --git a/packages/db/src/client.ts b/packages/db/src/client.ts index 2673eae1088..5149c1c9b27 100644 --- a/packages/db/src/client.ts +++ b/packages/db/src/client.ts @@ -267,6 +267,32 @@ export async function withDedicatedDbConnection(db: Db, action: (dedicated: D finally { await dedicated.$client.end({ timeout: 1 }); } } +// postgres.js pipelines a query onto a busy connection when none is idle, so +// a short liveness write can wait behind an unrelated slow statement with no +// client-side bound. Liveness writes (controller-lease renewal) get a small +// private pool, created on first use and shared for the owner's lifetime. +const livenessDbFactories = new WeakMap Db>(); +const livenessDbs = new WeakMap(); +const LIVENESS_POOL_MAX_CONNECTIONS = 2; + +/** The private liveness pool for a `createDb` client; other clients (test + * doubles, transactions) are returned unchanged. */ +export function livenessDb(db: Db): Db { + const existing = livenessDbs.get(db); + if (existing) return existing; + const factory = livenessDbFactories.get(db); + if (!factory) return db; + const created = factory(); + livenessDbs.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 }); +} + export function createDb(url: string, options?: DatabaseClientOptions) { const resolved = resolveDatabaseClientOptions(options ?? databaseClientOptionsFromEnv()); return createDbClient(url, resolved); @@ -284,6 +310,9 @@ function createDbClient(url: string, resolved: DatabaseClientOptions, owner?: ob dedicatedDbFactories.set(db, () => createDbClient(url, { ...resolved, maxConnections: 1, applicationName: "paperclip-workspace-finalization-lock", }, db)); + livenessDbFactories.set(db, () => createDbClient(url, { + ...resolved, maxConnections: LIVENESS_POOL_MAX_CONNECTIONS, applicationName: "paperclip-controller-lease", + }, db)); return db; } diff --git a/packages/db/src/index.ts b/packages/db/src/index.ts index 73c6028d067..d65d39527ac 100644 --- a/packages/db/src/index.ts +++ b/packages/db/src/index.ts @@ -1,6 +1,8 @@ export { createDb, withDedicatedDbConnection, + livenessDb, + closeLivenessDb, closeRegisteredClients, getPostgresDataDirectory, ensurePostgresDatabase, diff --git a/server/src/__tests__/server-startup-feedback-export.test.ts b/server/src/__tests__/server-startup-feedback-export.test.ts index 99658852b21..7794789e382 100644 --- a/server/src/__tests__/server-startup-feedback-export.test.ts +++ b/server/src/__tests__/server-startup-feedback-export.test.ts @@ -258,6 +258,7 @@ vi.mock("@paperclipai/db", () => ({ reconcilePendingMigrationHistory: vi.fn(async () => ({ repairedMigrations: [] })), formatDatabaseBackupResult: vi.fn(() => "ok"), runDatabaseBackup: vi.fn(), + closeLivenessDb: vi.fn(async () => {}), authUsers: {}, companies: {}, companyMemberships: {}, diff --git a/server/src/index.ts b/server/src/index.ts index 287ae6c34bd..1259003df67 100644 --- a/server/src/index.ts +++ b/server/src/index.ts @@ -66,6 +66,7 @@ import { companies, companyMemberships, instanceUserRoles, + closeLivenessDb, } from "@paperclipai/db"; import detectPort from "detect-port"; import { createApp } from "./app.js"; @@ -670,7 +671,7 @@ async function startServerWithDatabaseTeardown( const closeDatabaseClients = async () => { setSentryRunLinkResolver(null); const clients = pluginMigrationDb === db ? [db] : [db, pluginMigrationDb]; - await Promise.all(clients.map((client) => endDatabaseClient(client, 5))); + await Promise.all([...clients.map((client) => endDatabaseClient(client, 5)), closeLivenessDb(db, 5)]); }; startupDatabase.close = closeDatabaseClients; setSentryRunLinkResolver(createSentryRunLinkResolver(db)); diff --git a/server/src/services/legacy-controller-lease.test.ts b/server/src/services/legacy-controller-lease.test.ts index 354d88923f0..da54835709f 100644 --- a/server/src/services/legacy-controller-lease.test.ts +++ b/server/src/services/legacy-controller-lease.test.ts @@ -1,13 +1,21 @@ import { randomUUID } from "node:crypto"; import { eq, sql } from "drizzle-orm"; import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; -import { agents, companies, createDb, heartbeatRuns } from "@paperclipai/db"; +import { agents, closeLivenessDb, companies, createDb, heartbeatRuns, livenessDb } from "@paperclipai/db"; import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../__tests__/helpers/embedded-postgres.js"; import { heartbeatService } from "./heartbeat.js"; import { hasLiveLegacyController, LEGACY_CONTROLLER_LEASE_MS, LEGACY_CONTROLLER_RENEW_MS, legacyControllerBootId, legacyControllerClaim, LegacyControllerLeaseLostError, renewLegacyControllerLease, revokeExpiredLegacyController, revokeExpiredLegacyControllerWithState, watchLegacyControllerLease } from "./legacy-controller-lease.js"; +type WatchArgs = Parameters; +// 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) }; + return fake as unknown as WatchArgs[0]; +} + const support = await getEmbeddedPostgresTestSupport(); (support.supported ? describe : describe.skip)("durable legacy controller ownership", () => { let database: Awaited>; @@ -115,7 +123,7 @@ const support = await getEmbeddedPostgresTestSupport(); }); it("rejects dispatch at the lease deadline even if the database query never settles", async () => { const run = await seed(); - const hungDb = { update: () => ({ set: () => ({ where: () => ({ returning: () => new Promise(() => {}) }) }) }) } as unknown as typeof db; + const hungDb = fakeRenewalDb(() => new Promise(() => {})); vi.useFakeTimers(); const controller = new AbortController(); const watch = watchLegacyControllerLease(hungDb, { ...run, controllerLeaseExpiresAt: new Date(Date.now() + 100) }, controller); @@ -128,6 +136,60 @@ const support = await getEmbeddedPostgresTestSupport(); } finally { watch.stop(); vi.useRealTimers(); } }); + it("renews without queueing behind a slow statement on a saturated shared pool", async () => { + const run = await seed(); + // postgres.js pipelines onto a busy connection when none is idle. + const shared = createDb(database.connectionString, { maxConnections: 1 }); + const slow = shared.execute(sql`select pg_sleep(5)`).then(() => undefined); // drizzle sends on then() + const controller = new AbortController(); + const watch = watchLegacyControllerLease(shared, run, controller); + try { + const startedAt = Date.now(); + await watch.assertOwned("dispatching"); + expect(Date.now() - startedAt).toBeLessThan(2_000); + expect(controller.signal.aborted).toBe(false); + const [saved] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, run.id)); + expect(saved.executionStage).toBe("dispatching"); + } finally { watch.stop(); await slow; } + }, 15_000); + + it("renews on one private pool per client that shutdown can end", async () => { + const run = await seed(); + const owner = createDb(database.connectionString); + const pool = livenessDb(owner); + expect(pool).not.toBe(owner); + expect(livenessDb(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"); + } finally { watch.stop(); } + await closeLivenessDb(owner, 1); + await expect(livenessDb(owner).execute(sql`select 1`)).rejects.toThrow(); + }); + + 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; + 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"); + lockTaken(); + await released; + }); + await taken; + const controller = new AbortController(); + const watch = watchLegacyControllerLease(db, run, controller, { renewalStatementTimeoutMs: 200 }); + try { + const renewal = watch.assertOwned(); + setTimeout(releaseLock, 1_000); + await expect(renewal).resolves.toBeUndefined(); + expect(controller.signal.aborted).toBe(false); + } finally { releaseLock(); watch.stop(); await holder; } + }, 15_000); + it("leaves native controller ownership to the native coordinator", () => { expect(legacyControllerClaim("native")).toEqual({}); }); @@ -135,9 +197,7 @@ const support = await getEmbeddedPostgresTestSupport(); describe("legacy controller lease loss attribution", () => { it("attributes a failing renewal query to renewal_error", async () => { - type WatchArgs = Parameters; - const failingDb = { update: () => ({ set: () => ({ where: () => ({ - returning: () => Promise.reject(new Error("connection reset")) }) }) }) } as unknown as WatchArgs[0]; + const failingDb = fakeRenewalDb(() => Promise.reject(new Error("connection reset"))); vi.useFakeTimers(); const controller = new AbortController(); const run = { id: randomUUID(), companyId: randomUUID(), runtimeMode: "legacy", controllerBootId: legacyControllerBootId, @@ -151,16 +211,15 @@ describe("legacy controller lease loss attribution", () => { } finally { watch.stop(); vi.useRealTimers(); } }); - type WatchArgs = Parameters; const watchedRun = (leaseMs: number) => ({ id: randomUUID(), companyId: randomUUID(), runtimeMode: "legacy", controllerBootId: legacyControllerBootId, controllerLeaseExpiresAt: new Date(Date.now() + leaseMs) }) as WatchArgs[1]; it("records a renewal still in flight when a stalled query reaches the deadline", async () => { let calls = 0; - const stallingDb = { update: () => ({ set: () => ({ where: () => ({ returning: () => { + const stallingDb = fakeRenewalDb(() => { calls += 1; return calls <= 2 ? Promise.resolve([{ id: "run" }]) : new Promise(() => {}); - } }) }) }) } as unknown as WatchArgs[0]; + }); vi.useFakeTimers(); const controller = new AbortController(); const watch = watchLegacyControllerLease(stallingDb, watchedRun(LEGACY_CONTROLLER_LEASE_MS), controller); @@ -174,8 +233,33 @@ describe("legacy controller lease loss attribution", () => { } finally { watch.stop(); vi.useRealTimers(); } }); + it("measures the next deadline from the retry that renewed the lease", async () => { + const timeoutAfterMs = 45_000; + let calls = 0; + const retryingDb = fakeRenewalDb(() => { + calls += 1; + if (calls === 1) { + return new Promise((_, reject) => setTimeout(() => reject(Object.assign(new Error("statement timeout"), { code: "57014" })), timeoutAfterMs)); + } + return calls === 2 ? Promise.resolve([{ id: "run" }]) : new Promise(() => {}); + }); + vi.useFakeTimers(); + const controller = new AbortController(); + const watch = watchLegacyControllerLease(retryingDb, watchedRun(LEGACY_CONTROLLER_LEASE_MS), controller); + try { + await vi.advanceTimersByTimeAsync(LEGACY_CONTROLLER_RENEW_MS + timeoutAfterMs); + expect(calls).toBe(2); + // The retry wrote a fresh 60s lease; the first attempt's start must not shorten it. + await vi.advanceTimersByTimeAsync(LEGACY_CONTROLLER_LEASE_MS - 1); + expect(controller.signal.aborted).toBe(false); + await vi.advanceTimersByTimeAsync(2); + expect(controller.signal.reason.leaseLossCause).toBe("renewal_deadline"); + expect(controller.signal.reason.leaseLossTiming).toMatchObject({ sinceRenewalMs: LEGACY_CONTROLLER_LEASE_MS, renewalCount: 1 }); + } finally { watch.stop(); vi.useRealTimers(); } + }); + it("records how late the deadline ran when the event loop was blocked", async () => { - const idleDb = { update: () => ({ set: () => ({ where: () => ({ returning: () => new Promise(() => {}) }) }) }) } as unknown as WatchArgs[0]; + const idleDb = fakeRenewalDb(() => new Promise(() => {})); const controller = new AbortController(); const watch = watchLegacyControllerLease(idleDb, watchedRun(10), controller); try { diff --git a/server/src/services/legacy-controller-lease.ts b/server/src/services/legacy-controller-lease.ts index fe6b7e68f28..55a8439debc 100644 --- a/server/src/services/legacy-controller-lease.ts +++ b/server/src/services/legacy-controller-lease.ts @@ -1,6 +1,6 @@ import { randomUUID } from "node:crypto"; import { and, eq, gt, lte, sql } from "drizzle-orm"; -import { heartbeatRuns, type Db } from "@paperclipai/db"; +import { heartbeatRuns, livenessDb, type Db } from "@paperclipai/db"; import type { LegacyControllerLeaseLossCause, LegacyControllerLeaseLossTiming } from "./run-cancellation.js"; // A boot UUID has meaning across containers; a numeric PID does not. @@ -52,6 +52,31 @@ export async function renewLegacyControllerLease( return Boolean(renewed); } +/** One attempt may wait at most this long; a timed-out attempt is retried + * until the lease deadline instead of holding the renewal for the whole lease. */ +export const LEGACY_CONTROLLER_RENEWAL_STATEMENT_TIMEOUT_MS = 10_000; + +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; + } + 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 renewLegacyControllerLease(tx as unknown as Db, run, stage); + }); +} + export async function hasLiveLegacyController(db: Db, run: Run): Promise { if (run.runtimeMode === "native" || !run.controllerBootId) return false; const [owner] = await db.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(and( @@ -98,7 +123,12 @@ export async function revokeExpiredLegacyController(db: Db, run: Run): Promise reject(controller.signal.reason); controller.signal.addEventListener("abort", onAbort, { once: true }); }); + // The successful attempt's UPDATE writes a fresh lease, so the local + // deadline is measured from that attempt rather than the first one. + let attemptStartedAt = startedAt; + const renew = async () => { + for (;;) { + attemptStartedAt = Date.now(); + try { + return await renewOnLivenessPool(db, run, stage, renewalStatementTimeoutMs); + } catch (error) { + if (!isStatementTimeout(error) || stopped || controller.signal.aborted) throw error; + } + } + }; let renewed: boolean; try { - renewed = await Promise.race([renewLegacyControllerLease(db, run, stage), aborted]); + renewed = await Promise.race([renew(), aborted]); } finally { renewalsInFlight.delete(renewal); controller.signal.removeEventListener("abort", onAbort); @@ -157,7 +200,7 @@ export function watchLegacyControllerLease(db: Db, run: Run, controller: AbortCo if (!stopped) { lastRenewedAt = Date.now(); renewalCount += 1; - armDeadline(Math.max(0, LEGACY_CONTROLLER_LEASE_MS - (Date.now() - startedAt))); + armDeadline(Math.max(0, LEGACY_CONTROLLER_LEASE_MS - (Date.now() - attemptStartedAt))); } }; const timer = setInterval(() => {