diff --git a/server/src/__tests__/heartbeat-active-run-output-watchdog.test.ts b/server/src/__tests__/heartbeat-active-run-output-watchdog.test.ts index eb04d057a4..197dc260ed 100644 --- a/server/src/__tests__/heartbeat-active-run-output-watchdog.test.ts +++ b/server/src/__tests__/heartbeat-active-run-output-watchdog.test.ts @@ -1,5 +1,5 @@ import { randomUUID } from "node:crypto"; -import { and, eq, sql } from "drizzle-orm"; +import { and, eq } from "drizzle-orm"; import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "vitest"; import { activityLog, @@ -19,6 +19,7 @@ import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase, } from "./helpers/embedded-postgres.js"; +import { truncateTablesWithDeadlockRetry } from "./helpers/truncate-with-deadlock-retry.js"; import { appendHeartbeatRunEvent } from "../services/heartbeat-run-events.js"; import { ACTIVE_RUN_OUTPUT_CONTINUE_REARM_MS, @@ -43,29 +44,6 @@ if (!embeddedPostgresSupport.supported) { ); } -function errorHasPostgresCode(error: unknown, code: string): boolean { - let current: unknown = error; - for (let depth = 0; depth < 4; depth += 1) { - if (!current || typeof current !== "object") return false; - const record = current as { code?: unknown; cause?: unknown }; - if (record.code === code) return true; - current = record.cause; - } - return false; -} - -async function truncateCompaniesWithDeadlockRetry(db: ReturnType) { - for (let attempt = 0; attempt < 5; attempt += 1) { - try { - await db.execute(sql.raw(`TRUNCATE TABLE "companies" CASCADE`)); - return; - } catch (error) { - if (!errorHasPostgresCode(error, "40P01") || attempt === 4) throw error; - await new Promise((resolve) => setTimeout(resolve, 50 * (attempt + 1))); - } - } -} - describeEmbeddedPostgres("active-run output watchdog", () => { let tempDb: Awaited> | null = null; let db: ReturnType; @@ -77,7 +55,10 @@ describeEmbeddedPostgres("active-run output watchdog", () => { afterEach(async () => { mockedAppendHeartbeatRunEvent.mockClear(); - await truncateCompaniesWithDeadlockRetry(db); + await truncateTablesWithDeadlockRetry(db, `TRUNCATE TABLE "companies" CASCADE`, { + attempts: 5, + delayMs: (attempt) => 50 * (attempt + 1), + }); }); afterAll(async () => { diff --git a/server/src/__tests__/heartbeat-stale-queue-invalidation.test.ts b/server/src/__tests__/heartbeat-stale-queue-invalidation.test.ts index 882a81e5d2..99efb01d32 100644 --- a/server/src/__tests__/heartbeat-stale-queue-invalidation.test.ts +++ b/server/src/__tests__/heartbeat-stale-queue-invalidation.test.ts @@ -22,6 +22,8 @@ import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase, } from "./helpers/embedded-postgres.js"; +import { drainHeartbeatRunsToQuiescence } from "./helpers/drain-heartbeat-runs.js"; +import { truncateTablesWithDeadlockRetry } from "./helpers/truncate-with-deadlock-retry.js"; import { MAX_TURN_CONTINUATION_RETRY_REASON, MAX_TURN_CONTINUATION_WAKE_REASON, @@ -89,42 +91,30 @@ async function waitForCondition(fn: () => Promise, timeoutMs = 3_000) { } async function cleanupHeartbeatInvalidationFixture(db: ReturnType) { - for (let attempt = 0; attempt < 10; attempt += 1) { - try { - await db.execute(sql.raw(` - TRUNCATE TABLE - "company_skills", - "issue_comments", - "issue_documents", - "document_revisions", - "documents", - "issue_relations", - "issue_tree_holds", - "issues", - "heartbeat_run_events", - "cost_events", - "activity_log", - "heartbeat_runs", - "agent_wakeup_requests", - "agent_runtime_state", - "agents", - "companies" - RESTART IDENTITY CASCADE - `)); - return; - } catch (error) { - const isLateCommentRace = - error instanceof Error && - error.message.includes("issue_comments_issue_id_issues_id_fk"); - if (!isLateCommentRace || attempt === 9) { - throw error; - } - - // Heartbeat completion can write issue-thread comments shortly after the - // run leaves queued/running. Retry the dependent deletes once those land. - await new Promise((resolve) => setTimeout(resolve, 100)); - } - } + await truncateTablesWithDeadlockRetry( + db, + ` + TRUNCATE TABLE + "company_skills", + "issue_comments", + "issue_documents", + "document_revisions", + "documents", + "issue_relations", + "issue_tree_holds", + "issues", + "heartbeat_run_events", + "cost_events", + "activity_log", + "heartbeat_runs", + "agent_wakeup_requests", + "agent_runtime_state", + "agents", + "companies" + RESTART IDENTITY CASCADE + `, + { attempts: 10, delayMs: () => 100 }, + ); } type SeedOptions = { @@ -182,21 +172,12 @@ describeEmbeddedPostgres("heartbeat stale queued-run invalidation", () => { model: "test-model", })); runningProcesses.clear(); - let idlePolls = 0; - for (let attempt = 0; attempt < 100; attempt += 1) { - const runs = await db - .select({ status: heartbeatRuns.status }) - .from(heartbeatRuns); - const hasActiveRun = runs.some((run) => run.status === "queued" || run.status === "running"); - if (!hasActiveRun) { - idlePolls += 1; - if (idlePolls >= 3) break; - } else { - idlePolls = 0; - } - await new Promise((resolve) => setTimeout(resolve, 50)); - } - await new Promise((resolve) => setTimeout(resolve, 50)); + // A table poll for "queued" or "running" rows cannot see a wakeup promise + // that has not yet written its run row. The shared drain helper awaits + // both in-flight wakeup promises and in-flight run executions first, then + // re-checks the run table, so a late writer cannot still be in flight when + // the TRUNCATE below runs. + await drainHeartbeatRunsToQuiescence(db, heartbeat); await cleanupHeartbeatInvalidationFixture(db); }); diff --git a/server/src/__tests__/helpers/truncate-with-deadlock-retry.test.ts b/server/src/__tests__/helpers/truncate-with-deadlock-retry.test.ts new file mode 100644 index 0000000000..98d9d657ee --- /dev/null +++ b/server/src/__tests__/helpers/truncate-with-deadlock-retry.test.ts @@ -0,0 +1,96 @@ +import { describe, expect, it, vi } from "vitest"; +import { errorHasPostgresCode, truncateTablesWithDeadlockRetry } from "./truncate-with-deadlock-retry.js"; + +function deadlockError(): Error & { code: string } { + return Object.assign(new Error("deadlock detected"), { code: "40P01" }); +} + +function lateCommentForeignKeyError(): Error { + return new Error( + 'insert or update on table "issue_comments" violates foreign key constraint "issue_comments_issue_id_issues_id_fk"', + ); +} + +function unrelatedError(): Error & { code: string } { + return Object.assign(new Error("relation does not exist"), { code: "42P01" }); +} + +function fakeDb(execute: ReturnType) { + return { execute } as unknown as Parameters[0]; +} + +describe("truncateTablesWithDeadlockRetry", () => { + it("retries a deadlock (40P01) and succeeds once the lock clears", async () => { + const execute = vi.fn() + .mockRejectedValueOnce(deadlockError()) + .mockRejectedValueOnce(deadlockError()) + .mockResolvedValueOnce(undefined); + + await truncateTablesWithDeadlockRetry(fakeDb(execute), "TRUNCATE TABLE \"companies\" CASCADE", { + attempts: 5, + delayMs: () => 0, + }); + + expect(execute).toHaveBeenCalledTimes(3); + }); + + it("retries the late-comment foreign key violation and succeeds once the dependent row lands", async () => { + const execute = vi.fn() + .mockRejectedValueOnce(lateCommentForeignKeyError()) + .mockResolvedValueOnce(undefined); + + await truncateTablesWithDeadlockRetry(fakeDb(execute), "TRUNCATE TABLE \"issue_comments\" CASCADE", { + attempts: 5, + delayMs: () => 0, + }); + + expect(execute).toHaveBeenCalledTimes(2); + }); + + it("rethrows once attempts are exhausted on a persistent deadlock", async () => { + const error = deadlockError(); + const execute = vi.fn().mockRejectedValue(error); + + await expect( + truncateTablesWithDeadlockRetry(fakeDb(execute), "TRUNCATE TABLE \"companies\" CASCADE", { + attempts: 3, + delayMs: () => 0, + }), + ).rejects.toBe(error); + + expect(execute).toHaveBeenCalledTimes(3); + }); + + it("rethrows immediately on an error that is neither a deadlock nor the late-comment race", async () => { + const error = unrelatedError(); + const execute = vi.fn().mockRejectedValue(error); + + await expect( + truncateTablesWithDeadlockRetry(fakeDb(execute), "TRUNCATE TABLE \"companies\" CASCADE", { + attempts: 5, + delayMs: () => 0, + }), + ).rejects.toBe(error); + + expect(execute).toHaveBeenCalledTimes(1); + }); +}); + +describe("errorHasPostgresCode", () => { + it("finds the code on the error itself", () => { + expect(errorHasPostgresCode(deadlockError(), "40P01")).toBe(true); + }); + + it("finds the code on a wrapped cause", () => { + const wrapped = new Error("wrapper", { cause: deadlockError() }); + expect(errorHasPostgresCode(wrapped, "40P01")).toBe(true); + }); + + it("returns false when the code never matches", () => { + expect(errorHasPostgresCode(unrelatedError(), "40P01")).toBe(false); + }); + + it("returns false for a non-object error", () => { + expect(errorHasPostgresCode("plain string failure", "40P01")).toBe(false); + }); +}); diff --git a/server/src/__tests__/helpers/truncate-with-deadlock-retry.ts b/server/src/__tests__/helpers/truncate-with-deadlock-retry.ts new file mode 100644 index 0000000000..98f64b7894 --- /dev/null +++ b/server/src/__tests__/helpers/truncate-with-deadlock-retry.ts @@ -0,0 +1,49 @@ +import { sql } from "drizzle-orm"; +import type { createDb } from "@paperclipai/db"; + +type Db = ReturnType; + +const LATE_COMMENT_FOREIGN_KEY = "issue_comments_issue_id_issues_id_fk"; +const DEADLOCK_DETECTED_CODE = "40P01"; + +export function errorHasPostgresCode(error: unknown, code: string): boolean { + let current: unknown = error; + for (let depth = 0; depth < 4; depth += 1) { + if (!current || typeof current !== "object") return false; + const record = current as { code?: unknown; cause?: unknown }; + if (record.code === code) return true; + current = record.cause; + } + return false; +} + +function isLateWriterRace(error: unknown): boolean { + if (errorHasPostgresCode(error, DEADLOCK_DETECTED_CODE)) return true; + return error instanceof Error && error.message.includes(LATE_COMMENT_FOREIGN_KEY); +} + +// A fixture cleanup can race a background write that the test's drain step +// cannot see yet. Two forms of that race show up in this suite: a PostgreSQL +// deadlock (code 40P01) when a late writer holds a row lock that the +// TRUNCATE's lock ordering also needs, and a foreign-key violation when a +// dependent row lands between two deletes. Retry a fixed number of times with +// a backoff delay before the caller gives up and rethrows. +export async function truncateTablesWithDeadlockRetry( + db: Db, + truncateStatement: string, + opts: { attempts?: number; delayMs?: (attempt: number) => number } = {}, +): Promise { + const attempts = opts.attempts ?? 10; + const delayMs = opts.delayMs ?? ((attempt: number) => 50 * (attempt + 1)); + for (let attempt = 0; attempt < attempts; attempt += 1) { + try { + await db.execute(sql.raw(truncateStatement)); + return; + } catch (error) { + if (!isLateWriterRace(error) || attempt === attempts - 1) { + throw error; + } + await new Promise((resolve) => setTimeout(resolve, delayMs(attempt))); + } + } +}