mirror of
https://github.com/paperclipai/paperclip.git
synced 2026-11-12 22:56:58 +01:00
perf(db): index pending budget Stops instead of scanning all run results (#20391)
## Thinking Path > - Paperclip is the open source app people use to manage AI agents for work > - Budget enforcement stops agent runs when a budget hard limit is reached. The stop intent is saved in each run's `result_json` until run control delivers it. > - The budget worker, run control, and the idle-sleep check find these pending stops with a predicate on `result_json`. No index supports that predicate. > - Each check therefore reads every heartbeat run and its large result JSON from TOAST. On a hosted database, this statement family was among the most expensive: about 8,000 calls at about 1 s each, about 9,800 s in total. > - Pending stops are rare. The predicate is constant, so a partial index can contain only those rows. > - This pull request adds that partial index. The SQL that the server sends does not change. > - The benefit is that each check reads a few index entries instead of the full table. The database has more capacity for request work. ## Linked Issues or Issue Description **What existing behavior does this improve?** The check for committed budget Stops that are not delivered yet. The budget worker runs it each time it looks for work (`hasPendingBudgetEnforcement`), and again for each delivery pass (`listPendingBudgetCancellationScopes`). Run control and the idle-sleep check use the same predicate. **Subsystem affected** Budget enforcement in the server (`server/src/services/heartbeat/budget-cancellation.ts`, `server/src/services/budgets.ts`, `server/src/services/heartbeat/run-control.ts`, `server/src/services/idle-sleep-safety.ts`) and the `heartbeat_runs` schema in `packages/db`. **Current behavior** The query groups all rows of `heartbeat_runs` and reads `result_json->'startupCancellation'->'budgetEnforcement'` for each row. A limit of 1 does not help, because the group and sort occur before the limit. Cost grows with the run history, not with the number of pending stops. **Proposed behavior** A partial index, `heartbeat_runs_pending_budget_cancellation_idx`, contains only runs with a pending budget Stop. The planner uses it for all four callers. The results are the same. **Reason and benefit** Local SQL measurements show a 300x to 1,000x decrease in time and buffer reads (see Verification). **Breaking changes** None. ## What Changed - Added the partial index `heartbeat_runs_pending_budget_cancellation_idx` on `heartbeat_runs (company_id)`. Its predicate is identical to `pendingBudgetRunCancellationPredicate`. - Added migration `0455_pending_budget_cancellation_index.sql`. It uses `CREATE INDEX CONCURRENTLY IF NOT EXISTS`, and it is in the list of migrations that run outside a transaction, as `0436` is. - Moved the scoped run-control read and the idle-sleep check into exported builders in `budget-cancellation.ts`, next to the scope list query. The SQL does not change. The callers and the test now share the same statements. - Added a real PostgreSQL regression test. It checks that delivered, malformed, and missing intents are not pending. It also EXPLAINs the statements that the callers send (scope list with limit 1 and 100, scoped read, `EXISTS` check) and requires the new index in each plan. If a person changes a query shape, the server predicate, or the index predicate, the test fails. - Added `scripts/benchmark-budget-cancellation-scan.ts`. It captures the SQL that the real service sends, then compares the plans, timings, and complete rows with and without the index. - Updated `doc/DATABASE.md`. ## Verification Red/green regression test: - Before the migration, the new test failed: the plan was `Seq Scan` with the predicate as a filter. - After the migration: `npx vitest run src/services/heartbeat/budget-cancellation.test.ts` passes (17 tests). - Also passed: `src/__tests__/budgets-service.test.ts`, `src/__tests__/idle-sleep-diagnostics-postgres.test.ts`, `src/services/budget-resume-continuation.test.ts`. - `pnpm --filter @paperclipai/db check:migrations`, and typecheck of `@paperclipai/db` and `@paperclipai/server` pass. Performance and result equality (local SQL measurements only): ```sh pnpm --filter @paperclipai/server exec tsx ../scripts/benchmark-budget-cancellation-scan.ts --runs=20000 --payload-kib=16 [--reverse] ``` - Fixture: embedded PostgreSQL 18.1. 20,241 runs in 2 companies, with incompressible result JSON of about 16 KiB (table 379 MB). 40 pending stops in 6 scopes, 200 delivered or malformed decoys, and 1 cross-company intent. - The SQL is the same in both treatments. The script toggles only the new index. 2 warmups and 7 samples per query. Rows are compared outside the timed part. Ordered queries must match in order. The scoped run-control read has no `ORDER BY`, so it is compared as a multiset. | Query (caller) | Baseline median | With index median | Shared buffers | | --- | --- | --- | --- | | `listPendingBudgetCancellationScopes(db, 1)` (worker `hasPending`) | 688 ms | 2.1 ms | 243,500 → 1,158 | | `listPendingBudgetCancellationScopes(db)` (delivery) | 799 ms | 1.4 ms | 243,500 → 1,158 | | Scoped pending runs (run control) | 357 ms | 1.1 ms | 121,835 → 322 | | `EXISTS` pending (idle-sleep check) | 440 ms | 0.37 ms | 240,560 → 2 | - Reverse order (index first) gave the same result: 524 ms → 1.6 ms, 464 ms → 1.4 ms, 306 ms → 1.0 ms, 306 ms → 0.42 ms. - Plans change from `Seq Scan` (or a broad bitmap scan) to `Index Scan` or `Index Only Scan` on the new index. - All result rows are equal in every run. Migration duration: - 20,241 runs (379 MB): index build 2.0 s to 3.9 s. - 200,241 runs with 8 KiB result JSON (2.1 GB): index build 5.6 s. With this table size, the four queries change from 1,148 / 879 / 514 / 4,392 ms to 1.1 / 1.2 / 0.85 / 8.3 ms. - The build is concurrent, so it does not block run writes. ## Risks - Migration: the build reads every run's result JSON one time. It is concurrent and resumable (`IF NOT EXISTS`). The executor repairs an invalid concurrent index on retry. - Write cost: run writes evaluate the index predicate. The index contains very few rows, so the extra cost is small. - Predicate drift: if a later change edits only one of the two predicates, the planner cannot use the index. The new regression test detects this. - The hosted timings are not measured yet. The local numbers come from a synthetic fixture. ## Model Used - Claude Opus 5.5 (`claude-opus-5-5`), Anthropic, used through Claude Code with tool use (shell, file edits, local PostgreSQL). ## Checklist - [x] I have included a thinking path that traces from project context to this change - [x] I have specified the model used (with version and capability details) - [x] I have checked ROADMAP.md and confirmed this PR does not duplicate planned core work - [x] I have searched GitHub for duplicate or related PRs and linked them above - [x] I have either (a) linked existing issues with `Fixes: #` / `Closes #` / `Refs #` OR (b) described the issue in-PR following the relevant issue template - [x] I have not referenced internal/instance-local Paperclip issues or links (only public GitHub `#NNN` / `github.com/paperclipai/paperclip` URLs) - [x] My branch name describes the change (e.g. `docs/...`, `fix/...`) and contains no internal Paperclip ticket id or instance-derived details - [x] I have run tests locally and they pass - [x] I have added or updated tests where applicable - [x] I have updated relevant documentation to reflect my changes - [x] I have considered and documented any risks above - [ ] All Paperclip CI gates are green - [ ] Greptile is 5/5 with no open P2s, recommendations, or follow-ups - [x] I will address all Greptile and reviewer comments before requesting merge 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
DottaandClaude Opus 5.5 authored and GitHub committed
1 parent
ab80fea62f
commit
0116e45958
11 files changed
+466
-13
No files matched your search
@@ -0,0 +1,154 @@
|
||||
// Measures the committed budget Stop scans with and without the partial index
|
||||
// `heartbeat_runs_pending_budget_cancellation_idx`. The SQL is the same in both
|
||||
// treatments; only that index is toggled, in a disposable database.
|
||||
//
|
||||
// pnpm --filter @paperclipai/server exec tsx ../scripts/benchmark-budget-cancellation-scan.ts [--runs=20000] [--payload-kib=16] [--reverse]
|
||||
import assert from "node:assert/strict";
|
||||
import { randomBytes, randomUUID } from "node:crypto";
|
||||
import { readFileSync, readdirSync } from "node:fs";
|
||||
import { createRequire } from "node:module";
|
||||
import { fileURLToPath } from "node:url";
|
||||
import { agents, companies, createDb, heartbeatRuns, startEmbeddedPostgresTestDatabase, type Db } from "../packages/db/src/index.js";
|
||||
import {
|
||||
listPendingBudgetCancellationScopes,
|
||||
pendingBudgetRunCancellationExistsQuery,
|
||||
pendingBudgetRunIdsForScopeQuery,
|
||||
} from "../server/src/services/heartbeat/budget-cancellation.js";
|
||||
|
||||
const { sql } = createRequire(new URL("../server/package.json", import.meta.url))("drizzle-orm") as typeof import("drizzle-orm");
|
||||
const arg = (name: string, fallback: number) => Number(process.argv.find(value => value.startsWith(`--${name}=`))?.split("=")[1] ?? fallback);
|
||||
const runCount = arg("runs", 20_000);
|
||||
const payloadKib = arg("payload-kib", 16);
|
||||
const samples = arg("samples", 7);
|
||||
const indexName = "heartbeat_runs_pending_budget_cancellation_idx";
|
||||
const migrationDir = fileURLToPath(new URL("../packages/db/src/migrations/", import.meta.url));
|
||||
const migrationFile = readdirSync(migrationDir).find(file => file.endsWith(".sql")
|
||||
&& readFileSync(`${migrationDir}${file}`, "utf8").includes(`"${indexName}"`));
|
||||
assert.ok(migrationFile, "The index migration was not found");
|
||||
const createIndexSql = readFileSync(`${migrationDir}${migrationFile}`, "utf8").trim();
|
||||
|
||||
type Query = { sql: string; params: unknown[] };
|
||||
const keepAlive = setInterval(() => {}, 1000);
|
||||
let database: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>> | undefined;
|
||||
let db: Db | undefined;
|
||||
try {
|
||||
database = await startEmbeddedPostgresTestDatabase("budget-cancellation-scan-benchmark-");
|
||||
db = createDb(database.connectionString);
|
||||
const [version] = await db.$client.unsafe("select version()");
|
||||
|
||||
// Two tenants. Most runs carry an incompressible result JSON, so a scan must
|
||||
// read TOAST. A few runs carry committed Stops: pending, delivered, decoys.
|
||||
const companyIds = [randomUUID(), randomUUID()];
|
||||
const agentIds = [randomUUID(), randomUUID(), randomUUID()];
|
||||
await db.insert(companies).values(companyIds.map((id, n) => ({ id, name: `Budget scan ${n}`, issuePrefix: `BS${n}` })));
|
||||
await db.insert(agents).values(agentIds.map((id, n) => ({ id, companyId: companyIds[n % 2]!, name: `Agent ${n}`, adapterType: "process" })));
|
||||
const payload = randomBytes(payloadKib * 768).toString("base64");
|
||||
await db.$client.unsafe(`insert into heartbeat_runs (company_id, agent_id, status, runtime_mode, created_at, updated_at, result_json)
|
||||
select case when n % 2 = 0 then $1::uuid else $2::uuid end, case when n % 2 = 0 then $3::uuid else $4::uuid end,
|
||||
(array['succeeded','failed','cancelled','timed_out'])[1 + n % 4], 'legacy',
|
||||
timestamptz '2026-10-01' + n * interval '1 second', timestamptz '2026-10-01' + n * interval '1 second',
|
||||
jsonb_build_object('summary', $5::text, 'stdoutExcerpt', left($5::text, 512))
|
||||
from generate_series(1, $6::int) n`, [companyIds[0], companyIds[1], agentIds[0], agentIds[1], payload, runCount]);
|
||||
const policyId = randomUUID();
|
||||
const intent = (companyId: string, scopeType: string, scopeId: string, extra: Record<string, unknown> = {}) =>
|
||||
({ companyId, scopeType, scopeId, policyId, version: 3, ...extra });
|
||||
const special: Array<{ companyId: string; agentId: string; status: string; updatedAt: Date; resultJson: Record<string, unknown> }> = [];
|
||||
for (let n = 0; n < 40; n++) {
|
||||
const companyId = companyIds[n % 2]!, agentId = n % 2 === 0 ? agentIds[0]! : agentIds[1]!;
|
||||
const scopes: Array<[string, string]> = [["company", companyId], ["agent", agentId], ["agent", agentIds[2]!]];
|
||||
const [scopeType, scopeId] = scopes[n % 3]!;
|
||||
const extra = n % 5 === 0 ? { lastAttemptAt: new Date(Date.UTC(2026, 10, 1, 0, n)).toISOString() } : {};
|
||||
special.push({ companyId, agentId, status: n % 4 === 0 ? "cancelled" : "running", updatedAt: new Date(Date.UTC(2026, 10, 2, 0, n)),
|
||||
resultJson: { summary: payload, startupCancellation: {
|
||||
...(n % 7 === 0 ? {} : { requestedAt: new Date(Date.UTC(2026, 10, 1, 1, 40 - n)).toISOString() }),
|
||||
budgetEnforcement: intent(companyId, scopeType, scopeId, extra) } } });
|
||||
}
|
||||
// Delivered Stops and malformed intents are not pending.
|
||||
for (let n = 0; n < 200; n++) special.push({ companyId: companyIds[n % 2]!, agentId: agentIds[n % 2]!, status: "cancelled",
|
||||
updatedAt: new Date(Date.UTC(2026, 10, 3, 0, 0, n)), resultJson: { summary: payload, startupCancellation: {
|
||||
budgetEnforcement: n % 10 === 0 ? "legacy-string" : intent(companyIds[n % 2]!, "company", companyIds[n % 2]!,
|
||||
{ deliveredAt: new Date(Date.UTC(2026, 10, 3)).toISOString() }) } } });
|
||||
// An intent whose company does not match its run is returned by SQL and
|
||||
// filtered by the service; it must stay visible to both treatments.
|
||||
special.push({ companyId: companyIds[0]!, agentId: agentIds[0]!, status: "running", updatedAt: new Date(Date.UTC(2026, 10, 4)),
|
||||
resultJson: { startupCancellation: { budgetEnforcement: intent(companyIds[1]!, "company", companyIds[1]!) } } });
|
||||
await db.insert(heartbeatRuns).values(special.map(row => ({ ...row, runtimeMode: "legacy" })));
|
||||
await db.$client.unsafe("vacuum analyze heartbeat_runs");
|
||||
const [{ total_bytes: tableBytes }] = await db.$client.unsafe(`select pg_total_relation_size('heartbeat_runs')::bigint as total_bytes`);
|
||||
|
||||
// Capture the SQL that the real service and its callers send.
|
||||
const captured: Query[] = [];
|
||||
let capturing = true;
|
||||
const session = Reflect.get(db, "session");
|
||||
const logger = { logQuery(query: string, params: unknown[]) { if (capturing) captured.push({ sql: query, params }); } };
|
||||
session.logger = logger; session.options.logger = logger;
|
||||
const queries: Array<{ name: string; query: Query }> = [];
|
||||
const captureOne = async (name: string, invoke: () => Promise<unknown>) => {
|
||||
captured.length = 0; await invoke();
|
||||
assert.equal(captured.length, 1, `${name} sent ${captured.length} statements`);
|
||||
queries.push({ name, query: captured[0]! });
|
||||
};
|
||||
await captureOne("hasPendingBudgetEnforcement(limit 1)", () => listPendingBudgetCancellationScopes(db!, 1));
|
||||
await captureOne("deliverBudgetEnforcement(limit 100)", () => listPendingBudgetCancellationScopes(db!));
|
||||
await captureOne("runControl.cancelBudgetScopeWork", async () => pendingBudgetRunIdsForScopeQuery(db!,
|
||||
{ companyId: companyIds[0]!, scopeType: "agent", scopeId: agentIds[2]! }));
|
||||
await captureOne("idleSleep.budget_cancellation", () => db!.execute(sql`select exists (${pendingBudgetRunCancellationExistsQuery()}) as pending`));
|
||||
capturing = false;
|
||||
|
||||
async function measure(query: Query) {
|
||||
const timings: number[] = []; let rows: unknown[] = [];
|
||||
for (let n = 0; n < samples + 2; n++) {
|
||||
const start = performance.now();
|
||||
rows = await db!.$client.unsafe(query.sql, query.params as never[]).values();
|
||||
if (n > 1) timings.push(performance.now() - start);
|
||||
}
|
||||
const [plan] = await db!.$client.unsafe(`EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) ${query.sql}`, query.params as never[]);
|
||||
const root = plan["QUERY PLAN"][0];
|
||||
const nodes: string[] = [];
|
||||
(function walk(node: Record<string, unknown>) { nodes.push(String(node["Node Type"]) + (node["Index Name"] ? `(${node["Index Name"]})` : ""));
|
||||
for (const child of (node.Plans as Record<string, unknown>[] | undefined) ?? []) walk(child); })(root.Plan);
|
||||
timings.sort((a, b) => a - b);
|
||||
return { medianMs: timings[Math.floor(timings.length / 2)]!, minMs: timings[0]!, maxMs: timings.at(-1)!, rows,
|
||||
sharedBuffers: Number(root.Plan["Shared Hit Blocks"] ?? 0) + Number(root.Plan["Shared Read Blocks"] ?? 0), nodes };
|
||||
}
|
||||
const indexed = async (present: boolean) => {
|
||||
const existing = await db!.$client.unsafe(`select 1 from pg_indexes where indexname = $1`, [indexName]);
|
||||
if (existing.length) await db!.$client.unsafe(`drop index "${indexName}"`);
|
||||
if (present) {
|
||||
const start = performance.now();
|
||||
await db!.$client.unsafe(createIndexSql);
|
||||
await db!.$client.unsafe("analyze heartbeat_runs");
|
||||
return performance.now() - start;
|
||||
}
|
||||
return null;
|
||||
};
|
||||
const order: Array<"baseline" | "treatment"> = process.argv.includes("--reverse") ? ["treatment", "baseline"] : ["baseline", "treatment"];
|
||||
const results: Record<string, Record<string, Awaited<ReturnType<typeof measure>>>> = {};
|
||||
let indexBuildMs: number | null = null;
|
||||
for (const variant of order) {
|
||||
const built = await indexed(variant === "treatment");
|
||||
if (built !== null) indexBuildMs = built;
|
||||
for (const { name, query } of queries) (results[name] ??= {})[variant] = await measure(query);
|
||||
}
|
||||
const report = Object.entries(results).map(([name, { baseline, treatment }]) => {
|
||||
// Run control's scoped read has no ORDER BY; its caller treats the rows as a
|
||||
// set. Every other read promises its order, so compare it exactly.
|
||||
const comparable = (rows: unknown[]) => name === "runControl.cancelBudgetScopeWork"
|
||||
? rows.map(row => JSON.stringify(row)).sort() : rows;
|
||||
assert.deepEqual(comparable(treatment!.rows), comparable(baseline!.rows), `${name} returned different rows`);
|
||||
if (name !== "idleSleep.budget_cancellation") assert.ok(treatment!.nodes.some(node => node.includes(indexName)), `${name} did not use the index`);
|
||||
return { name, equalRows: baseline!.rows.length, baselineMs: baseline!.medianMs, baselineRangeMs: [baseline!.minMs, baseline!.maxMs],
|
||||
treatmentMs: treatment!.medianMs, treatmentRangeMs: [treatment!.minMs, treatment!.maxMs],
|
||||
speedup: Number((baseline!.medianMs / treatment!.medianMs).toFixed(1)),
|
||||
baselineBuffers: baseline!.sharedBuffers, treatmentBuffers: treatment!.sharedBuffers,
|
||||
baselinePlan: baseline!.nodes.join(" > "), treatmentPlan: treatment!.nodes.join(" > ") };
|
||||
});
|
||||
// The service's own post-processing must agree too.
|
||||
const scopes = await listPendingBudgetCancellationScopes(db);
|
||||
assert.ok(scopes.length > 0);
|
||||
console.log(JSON.stringify({ postgres: version.version, runs: runCount + special.length, payloadKib, tableBytes: Number(tableBytes),
|
||||
order, samples, indexBuildMs, createIndexSql, servicesScopes: scopes.length, report }, null, 2));
|
||||
} finally {
|
||||
try { await db?.$client.end({ timeout: 0 }); await database?.cleanup(); }
|
||||
finally { clearInterval(keepAlive); }
|
||||
}
|
||||
Reference in new issue
Block a user