diff --git a/scripts/benchmark-run-event-checkpoint-lock.ts b/scripts/benchmark-run-event-checkpoint-lock.ts new file mode 100644 index 00000000000..42e7d5cadd2 --- /dev/null +++ b/scripts/benchmark-run-event-checkpoint-lock.ts @@ -0,0 +1,123 @@ +/** + * Local proof for the run-event append lock and controller checkpoint saves. + * + * A native run appends every PRP event through `appendHeartbeatRunEvent`, + * which locks the run row. The controller then saves its checkpoint: a new + * `event:N` row plus the `cursor` row. The new row's foreign key check takes + * FOR KEY SHARE on the same run row. This script replays that pattern with the + * real append and the real checkpoint writer, on separate connections, and + * times each checkpoint save. It compares a baseline append module (FOR UPDATE) + * with the current one (FOR NO KEY UPDATE), in alternating order. + * + * Overlap is the normal controller order. The frame chain awaits each event's + * append (onCommittedEvent) and local commit, but not its checkpoint save: the + * ACK goes through #sendJson, which queues the save and does not block the + * next frame. Only semantic inputs and command results await the save. So the + * save for event N runs while event N+1 is appended. SERIAL=1 awaits every + * save before the next append, as a control without overlap. + * + * HOLD_MS keeps each append transaction open that much longer. It stands in + * for hosted round trips, which make the real append transaction last longer + * than on loopback. + * + * git show :server/src/services/heartbeat-run-events.ts \ + * > server/src/services/heartbeat-run-events.baseline.ts + * pnpm --filter @paperclipai/server exec tsx ../scripts/benchmark-run-event-checkpoint-lock.ts + * + * Optional: EVENTS (default 400 per pass), PASSES (4), HOLD_MS (5), EVENT_KIB (2), SERIAL (0). + * Delete the baseline copy afterwards. + */ +import assert from "node:assert/strict"; +import { randomBytes, randomUUID } from "node:crypto"; +import { existsSync } from "node:fs"; +import { createRequire } from "node:module"; +import { fileURLToPath } from "node:url"; +import { agents, companies, createDb, heartbeatRunEvents, heartbeatRuns, issues, nativeRunFinalizations, + nativeRunnerCheckpoints, startEmbeddedPostgresTestDatabase } from "../packages/db/src/index.js"; +import * as candidate from "../server/src/services/heartbeat-run-events.js"; +import { nativeRunnerCheckpointWriter } from "../server/src/services/native-runtime/native-runner-checkpoint.js"; + +const { eq, sql: drizzleSql } = createRequire(new URL("../server/package.json", import.meta.url))("drizzle-orm") as typeof import("drizzle-orm"); + +const baselinePath = fileURLToPath(new URL("../server/src/services/heartbeat-run-events.baseline.ts", import.meta.url)); +assert(existsSync(baselinePath), "Write the baseline module first (see the header)"); +const baseline = await import(baselinePath) as typeof candidate; +const size = (name: string, fallback: number) => Number(process.env[name] ?? fallback); +const EVENTS = size("EVENTS", 400), PASSES = size("PASSES", 4), HOLD_MS = size("HOLD_MS", 5), EVENT_KIB = size("EVENT_KIB", 2); +const SERIAL = process.env.SERIAL === "1"; +process.env.PAPERCLIP_SECRETS_MASTER_KEY ??= randomBytes(32).toString("base64"); + +const keepAlive = setInterval(() => {}, 1_000); +const temporary = await startEmbeddedPostgresTestDatabase("run-event-checkpoint-lock-"); +const appendDb = createDb(temporary.connectionString); +const checkpointDb = createDb(temporary.connectionString); +try { + const [company] = await appendDb.insert(companies).values({ name: "Lock", issuePrefix: "LCK" }).returning(); + const [agent] = await appendDb.insert(agents).values({ companyId: company!.id, name: "Native" }).returning(); + const [issue] = await appendDb.insert(issues).values({ companyId: company!.id, title: "Streaming task" }).returning(); + const quantile = (values: number[], q: number) => [...values].sort((a, b) => a - b)[Math.min(values.length - 1, Math.floor(values.length * q))]!; + const round = (value: number) => Number(value.toFixed(2)); + + async function pass(label: "baseline" | "candidate") { + const append = (label === "baseline" ? baseline : candidate).appendHeartbeatRunEvent; + const runId = randomUUID(), bootId = randomUUID(), leaseOwner = randomUUID(); + await appendDb.insert(heartbeatRuns).values({ id: runId, companyId: company!.id, agentId: agent!.id, status: "running", + runtimeMode: "native", nativeIssueId: issue!.id, issueId: issue!.id, + runnerProfileJson: { padding: randomBytes(3 * 1024).toString("base64") } }); + await appendDb.insert(nativeRunFinalizations).values({ runId, companyId: company!.id, issueId: issue!.id, phase: "observed", + controllerBootId: bootId, controllerGeneration: 1, leaseOwner, leaseExpiresAt: new Date(Date.now() + 3_600_000) }); + const write = nativeRunnerCheckpointWriter(checkpointDb, { companyId: company!.id, issueId: issue!.id, runId, + authorityRunId: runId, bootId, generation: 1, leaseOwner }); + const checkpoint = (records: Map, changed: Map, reset: boolean) => ({ + runId, records, changedRecords: changed, removedRecords: [], reset, workCommandIds: [], predecessor: null, + runnerVersion: "bench", runnerDigest: `sha256:${"a".repeat(64)}`, + recoveryEvidence: { runnerInstanceId: "bench", normalizedSessionId: "bench", authenticatedConnectionCount: 1, highestRecordedSourceSeq: 0 }, + restoreSessionDirectory() {}, + }) as unknown as Parameters[0]; + const records = new Map([["header", JSON.stringify({ padding: randomBytes(6 * 1024).toString("base64") })], + ["cursor", JSON.stringify({ ackedSourceSeq: 0 })]]); + await write(checkpoint(records, new Map(records), true)); + const saves: number[] = []; + let tail = Promise.resolve(); + const started = performance.now(); + for (let seq = 1; seq <= EVENTS; seq++) { + const event = { sourceSeq: seq, eventType: "item.completed", text: randomBytes(EVENT_KIB * 768).toString("base64") }; + // The append commits before the controller ACKs, like onCommittedEvent. + await appendDb.transaction(async tx => { + await append(tx as unknown as typeof appendDb, { companyId: company!.id, agentId: agent!.id, runId, + eventType: "item.completed", stream: "system", level: "info", payload: { prpEvent: event }, + nativeSource: { sourceInstanceId: "bench-runner", sourceEventId: `event-${seq}`, sourceSeq: seq, + protocolSchemaVersion: 1, canonicalPayload: event } }); + if (HOLD_MS) await tx.execute(drizzleSql`select pg_sleep(${HOLD_MS / 1000})`); + }); + const changed = new Map([[`event:${seq}`, JSON.stringify(event)], ["cursor", JSON.stringify({ ackedSourceSeq: seq })]]); + for (const [key, value] of changed) records.set(key, value); + // Checkpoint saves are serialized per controller and overlap the next append. + tail = tail.then(async () => { + const start = performance.now(); + await write(checkpoint(new Map(records), changed, false)); + saves.push(performance.now() - start); + }); + if (SERIAL) await tail; + } + await tail; + const elapsed = performance.now() - started; + const [stored] = await checkpointDb.select({ count: drizzleSql`count(*)::int` }).from(nativeRunnerCheckpoints) + .where(eq(nativeRunnerCheckpoints.runId, runId)); + const events = await appendDb.select({ seq: heartbeatRunEvents.seq }).from(heartbeatRunEvents) + .where(eq(heartbeatRunEvents.runId, runId)).orderBy(heartbeatRunEvents.seq); + assert.equal(stored!.count, EVENTS + 2, "Every event row, the header and the cursor are stored"); + assert.deepEqual(events.map(row => row.seq), Array.from({ length: EVENTS }, (_, index) => index + 1)); + return { label, saves: saves.length, saveMedianMs: round(quantile(saves, 0.5)), saveP90Ms: round(quantile(saves, 0.9)), + saveMaxMs: round(Math.max(...saves)), savesOver2Ms: saves.filter(value => value > 2).length, + saveTotalMs: round(saves.reduce((sum, value) => sum + value, 0)), elapsedMs: round(elapsed) }; + } + + await pass("candidate"); // Warm caches and the prepared-statement path once. + const results = []; + for (let index = 0; index < PASSES; index++) results.push(await pass(index % 2 === 0 ? "baseline" : "candidate")); + console.log(JSON.stringify({ events: EVENTS, serial: SERIAL, holdMs: HOLD_MS, eventKiB: EVENT_KIB, results }, null, 2)); +} finally { + clearInterval(keepAlive); + await temporary.cleanup(); +} diff --git a/server/src/__tests__/heartbeat-run-event-atomic-append.test.ts b/server/src/__tests__/heartbeat-run-event-atomic-append.test.ts index 3d4b55b76cb..69d9b9e688c 100644 --- a/server/src/__tests__/heartbeat-run-event-atomic-append.test.ts +++ b/server/src/__tests__/heartbeat-run-event-atomic-append.test.ts @@ -3,7 +3,7 @@ import { randomBytes } from "node:crypto"; import { afterAll, beforeAll, describe, expect, it } from "vitest"; import { eq, inArray, sql } from "drizzle-orm"; import { drizzle } from "drizzle-orm/postgres-js"; -import { agents, companies, createDb, heartbeatRunEvents, heartbeatRuns } from "@paperclipai/db"; +import { agents, companies, createDb, heartbeatRunEvents, heartbeatRuns, issues, nativeRunnerCheckpoints } from "@paperclipai/db"; import { appendHeartbeatRunEvent } from "../services/heartbeat-run-events.js"; import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "./helpers/embedded-postgres.js"; @@ -51,6 +51,7 @@ describeEmbeddedPostgres("atomic ordinary run events", () => { const fixtureCompanies = [companyId, otherCompanyId]; await db.delete(heartbeatRunEvents).where(inArray(heartbeatRunEvents.companyId, fixtureCompanies)); await db.delete(heartbeatRuns).where(inArray(heartbeatRuns.companyId, fixtureCompanies)); + await db.delete(issues).where(inArray(issues.companyId, fixtureCompanies)); await db.delete(agents).where(inArray(agents.companyId, fixtureCompanies)); await db.delete(companies).where(inArray(companies.id, fixtureCompanies)); } @@ -145,7 +146,7 @@ describeEmbeddedPostgres("atomic ordinary run events", () => { expect((await readRun(runId)).nextEventSeq).toBe(2); expect(await readEvents(runId)).toEqual([committed.row]); const lockReads = statements.filter(query => query.startsWith("select ") && - query.includes('from "heartbeat_runs"') && query.endsWith("for update")); + query.includes('from "heartbeat_runs"') && query.endsWith("for no key update")); expect(lockReads).toHaveLength(2); for (const query of lockReads) { expect(query).toMatch(/^select "company_id", "agent_id" from "heartbeat_runs" /); @@ -153,6 +154,53 @@ describeEmbeddedPostgres("atomic ordinary run events", () => { } }); + it("serializes appenders without blocking child foreign key checks on the run", async () => { + const runId = await makeRun(); + const [issue] = await db.insert(issues).values({ companyId, title: "Checkpointed task" }).returning(); + await db.update(heartbeatRuns).set({ nativeIssueId: issue!.id, runtimeMode: "native" }) + .where(eq(heartbeatRuns.id, runId)); + const native = (ordinal: number) => ({ + companyId, agentId, runId, eventType: "item.completed", + nativeSource: { sourceInstanceId: "native-writer", sourceEventId: `held-${ordinal}`, + sourceSeq: ordinal, protocolSchemaVersion: 1, canonicalPayload: { ordinal } }, + }); + let releaseHolder!: () => void; + const held = new Promise(resolve => { releaseHolder = resolve; }); + let holderLocked!: () => void; + const locked = new Promise(resolve => { holderLocked = resolve; }); + // An append whose transaction is still open holds the run lock, as an + // event append does while its remaining statements run. + const holder = otherDb.transaction(async tx => { + await appendHeartbeatRunEvent(tx as unknown as typeof db, native(1)); + holderLocked(); + await held; + }); + await locked; + let secondSettled = false; + const second = appendHeartbeatRunEvent(db, native(2)).finally(() => { secondSettled = true; }); + try { + // A controller checkpoint row references the run row by foreign key. + // Its insert must not wait for the open append. + await db.transaction(async tx => { + await tx.execute(sql`set local lock_timeout = '2s'`); + await tx.insert(nativeRunnerCheckpoints).values({ companyId, issueId: issue!.id, runId, + part: "event:1", material: { fixture: true } }); + }); + // Another appender still waits for the open append. + await new Promise(resolve => setTimeout(resolve, 200)); + expect(secondSettled).toBe(false); + } finally { + releaseHolder(); + await holder; + } + expect((await second).row.seq).toBe(2); + expect((await readEvents(runId)).map(row => [row.seq, row.sourceSeq])).toEqual([[1, 1], [2, 2]]); + expect((await readRun(runId)).nextEventSeq).toBe(3); + expect(await db.select({ part: nativeRunnerCheckpoints.part }).from(nativeRunnerCheckpoints) + .where(eq(nativeRunnerCheckpoints.runId, runId))).toEqual([{ part: "event:1" }]); + await db.delete(nativeRunnerCheckpoints).where(eq(nativeRunnerCheckpoints.runId, runId)); + }); + it("rejects mismatched authority without changing either run", async () => { const runId = await makeRun(); const otherRunId = "30000000-0000-4000-8000-000000000099"; diff --git a/server/src/services/heartbeat-run-events.ts b/server/src/services/heartbeat-run-events.ts index 1c3325ebd05..695e740f658 100644 --- a/server/src/services/heartbeat-run-events.ts +++ b/server/src/services/heartbeat-run-events.ts @@ -134,7 +134,11 @@ export async function appendHeartbeatRunEvent( ...(projectsProgress ? { issueId: boundRunIssueIdSql() } : {}) }) .from(heartbeatRuns) .where(eq(heartbeatRuns.id, input.runId)) - .for("update") + // Appenders, Stop, recovery and deletion still serialize on this row. + // The append never changes key columns, so it must not block foreign + // key checks (FOR KEY SHARE) of child rows, such as controller + // checkpoints, that are written for the same run while events stream. + .for("no key update") .limit(1) .then((rows) => rows[0] ?? null); if (!run || run.companyId !== input.companyId || run.agentId !== input.agentId) {