mirror of
https://github.com/paperclipai/paperclip.git
synced 2026-11-11 05:11:49 +01:00
perf(native): index the backup lease lookup by scope digest (#20571)
## Thinking Path > - Paperclip is the open source app people use to manage AI agents for work > - Native harness backup expiration removes old session backups. Before it removes a backup, it finds the environment leases that reference the backup root > - This lookup reads leases by run ID, by the `metadata.nativeHarnessBackup.sessionScopeSha256` digest, or by the identity's own lease ID, in one `OR` query > - The metadata arm has no index. So PostgreSQL reads every lease of the company, and the metadata of each lease, for each lookup > - On a hosted test database, this statement ran 21,736 times at a mean of 158 ms in one window. After other fixes, it still ran 2,364 times at a mean of 121 ms > - This pull request adds a small partial expression index on the digest > - The benefit is that each lookup reads a few index entries instead of all company leases. The query text and the result do not change ## Linked Issues or Issue Description **What existing behavior does this improve?** The lease lookup in `backupDependencies` in `server/src/services/native-runtime/native-harness-backup-expiration.ts`. `reconcileBackup` calls it for each backup root when the expiration worker loads or checks a root. **Subsystem affected** Native harness backup expiration and the `environment_leases` table schema. **Current behavior** The query is `company_id = $1 and (heartbeat_run_id in (...) or metadata->'nativeHarnessBackup'->>'sessionScopeSha256' = $n or id = $m)`. The run arm and the ID arm have indexes. The digest arm has no index. A `BitmapOr` plan needs an index for every arm, so the planner reads all company leases and evaluates the JSON expression on each row. Lease metadata holds sandbox details, so these reads are large. **Proposed behavior** The partial index `environment_leases_native_backup_scope_idx` keys leases by the digest. It holds only leases that have a backup stamp. Every arm of the `OR` now has an index, and the planner combines three index probes. The SQL and the returned lease IDs do not change. **Reason and benefit** This statement was in the top ten by total database time on a loaded hosted database. A local reproduction shows a query that is more than 6,000x faster and reads 9 buffers instead of about 436,000. **Breaking changes** None. The API, the lease-matching rules and the results do not change. ## What Changed - Add the partial expression index `environment_leases_native_backup_scope_idx` on `(metadata->'nativeHarnessBackup'->>'sessionScopeSha256')`, where that value is not null. - Add migration `0459`. It uses `CREATE INDEX CONCURRENTLY IF NOT EXISTS`. The migration runner runs it as a top-level statement, the same as migrations 0454–0458. - Add a regression test. It checks the lease-matching rules: run leases, a lease with only the backup stamp, and the identity's own lease match. Another digest, a non-string stamp, a top-level key with the same name, and the same digest in another company do not match. The test also checks that the custom plan and a forced generic plan of the generated SQL use `BitmapOr` with the new index and no sequential scan. - Add `scripts/benchmark-environment-lease-backup-lookup.ts`. It captures the SQL that the real service generates and measures it with and without the index. - Document the index in `doc/DATABASE.md`. ## Verification Cause and thesis. The cost comes from the one `OR` arm without an index, which forces a scan of all company leases and their metadata. Thesis: an index on the digest lets the planner use a `BitmapOr` of three index probes, so the cost no longer depends on the company's lease count. The results stay identical. The benchmark confirms the thesis. I did not change the SQL. With the index, the existing `OR` query already gets the index plan. A `UNION` rewrite would not help without the index, because the digest arm would still scan the company. Local SQL benchmark: ```sh pnpm --filter @paperclipai/server exec tsx ../scripts/benchmark-environment-lease-backup-lookup.ts ``` Fixture (synthetic): one company with 100,000 leases with 4 KiB of incompressible metadata each. 30% of the leases have a backup stamp for another session. A second company has 20,000 leases. The table is 678 MB. The backup root has 3 run leases, 1 lease with only the stamp, and 1 identity lease. The same digest is also in the other company. Embedded PostgreSQL 18 on an Apple laptop, which other work loaded heavily. The hosted database is Neon with different compute, so absolute times are local SQL measurements only. The script measures in the order: no index, index, no index, index. Each measurement has 2 warmups and 11 samples of `EXPLAIN ANALYZE` execution time and of the full `reconcileBackup` call. | Measurement | Without index | With index | | --- | --- | --- | | Lease query execution, run 1 | median 196 ms (138–562 ms) | median 0.010 ms (0.007–0.017 ms) | | Lease query execution, run 2 | median 185 ms (144–295 ms) | median 0.030 ms (0.020–0.045 ms) | | `reconcileBackup` wall time (includes file reads), runs 1 and 2 | median 188 ms and 216 ms | median 1.4 ms and 4.2 ms | | Plan | `Gather > Seq Scan` | `Bitmap Heap Scan > BitmapOr` of `heartbeat_run_idx`, the new index and `pkey` | | Buffers | 365,880 hit + 70,050 read | 9 hit | These numbers are from head `65bd9d7ac6`. The final head only renumbers the migration to 0459 after rebases; the SQL, the index and the service code are the same. Two earlier runs on the same fixture, SQL and index gave 236–425 ms without the index and 0.023–0.033 ms with the index. The machine had heavy load from other work, so the spread without the index is wide. Output check: in all four measurements, the service returned the same 5 lease IDs, and the complete query rows were equal. Index build (`CREATE INDEX CONCURRENTLY` from the migration file): 0.8–2.4 s for the 678 MB table over all runs. Index size: 3.4 MB with 30% of leases stamped. Tests: - `cd server && npx vitest run src/__tests__/native-harness-backup-expiration.test.ts`: 109 passed. The new test failed when I dropped the index before the plan check (confirmed red), and it passes with the index. - `pnpm --filter @paperclipai/db check:migrations`: passed. - `pnpm --filter @paperclipai/db typecheck` and `pnpm --filter @paperclipai/server typecheck`: passed. - `pnpm verify`: it asks for broader coverage (exit 2) because of the new script. See the CI results for this head. ## Risks - Migration time. A concurrent index build reads the full `environment_leases` table and does not block writes. Locally, the build took 0.8–2.4 s for 678 MB. The time grows linearly with the table size. I did not time it on a production-size clone. I expect a few seconds to tens of seconds on a large instance, which is in the 1–2 minute budget. This estimate is not proven. - If the build stops, PostgreSQL can leave an invalid index. The migration runner repairs invalid concurrent indexes when it runs again, and `IF NOT EXISTS` makes a rerun safe. - The index expression must stay textually equal to the query expression. If someone changes the JSON path, the planner cannot use the index. The new test catches this for custom and generic plans. - The index is partial, so ANALYZE does not collect statistics for the expression. The planner uses the default equality estimate. On the large fixture and in the small test, the planner still selects the index plan. A company with very few leases can still get a sequential scan, which is cheap. - Each lease write with a backup stamp writes one more small index entry. The write cost is small. - This change does not reduce how often the worker calls the lookup. It makes each call cheap. The other digest lookups in backup retention (`native-harness-backup-retention.ts`) and session retention checks can also use the index. - Previous-release story: no persisted state changes. The index is derived from existing lease metadata. - The hosted improvement is not verified until deployment. After deployment, compare a new query-statistics window for this statement. ## Model Used - Claude Opus 5.5 (`claude-opus-5-5`), Anthropic, used through Claude Code with tool use (shell, file edits, local PostgreSQL benchmarks). ## 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:
1 parent
7368a0b15c
commit
d5c1f48180
8 files changed
+437
-2
No files matched your search
@@ -461,6 +461,21 @@ pnpm --filter @paperclipai/server exec tsx ../scripts/benchmark-chat-admission-i
|
||||
pnpm --filter @paperclipai/server exec tsx ../scripts/benchmark-chat-admission-index-migration.ts
|
||||
```
|
||||
|
||||
Native harness backup expiration finds the leases that reference one backup root
|
||||
by run, by `metadata.nativeHarnessBackup.sessionScopeSha256`, and by the
|
||||
identity's own lease ID, in one `OR` query. The partial expression index
|
||||
`environment_leases_native_backup_scope_idx` from migration 0459 keys leases by
|
||||
that scope digest. With it, every arm of the `OR` has an index, so PostgreSQL
|
||||
combines three index probes instead of reading every lease of the company and its
|
||||
metadata. The expression must stay identical to the one in
|
||||
`server/src/services/native-runtime/native-harness-backup-expiration.ts`; a
|
||||
regression test checks that the planner uses the index. The migration builds the
|
||||
index concurrently. To compare both plans on a disposable local fixture:
|
||||
|
||||
```sh
|
||||
pnpm --filter @paperclipai/server exec tsx ../scripts/benchmark-environment-lease-backup-lookup.ts
|
||||
```
|
||||
|
||||
The chat wake-notice sweep compares UUID columns directly against guarded,
|
||||
canonical lowercase UUID payload values. Malformed and noncanonical payloads
|
||||
still match nothing. The query can use existing primary keys without adding a
|
||||
|
||||
@@ -621,6 +621,7 @@ const MIGRATIONS_WITH_TOP_LEVEL_STATEMENTS = new Set([
|
||||
"0456_chat_review_lookup_indexes.sql",
|
||||
"0457_native_session_run_index.sql",
|
||||
"0458_chat_admission_issue_index.sql",
|
||||
"0459_environment_lease_backup_scope_index.sql",
|
||||
|
||||
|
||||
"0384_funny_adam_warlock.sql",
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
CREATE INDEX CONCURRENTLY IF NOT EXISTS "environment_leases_native_backup_scope_idx" ON "environment_leases" USING btree (("metadata"->'nativeHarnessBackup'->>'sessionScopeSha256')) WHERE ("environment_leases"."metadata"->'nativeHarnessBackup'->>'sessionScopeSha256') is not null;
|
||||
+171
-2
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"id": "f98af245-4a2e-4c3d-b89e-6ed3e81a074c",
|
||||
"prevId": "d54bfeb4-dedb-4f8b-9195-a4d4aa67d71f",
|
||||
"id": "e070677d-cfb8-4438-847c-29c4da7ca8c5",
|
||||
"prevId": "3347b37d-b35b-4357-b08c-a78c017086ac",
|
||||
"version": "7",
|
||||
"dialect": "postgresql",
|
||||
"tables": {
|
||||
@@ -10188,6 +10188,40 @@
|
||||
"method": "btree",
|
||||
"with": {}
|
||||
},
|
||||
"chat_actions_company_issue_kind_idx": {
|
||||
"name": "chat_actions_company_issue_kind_idx",
|
||||
"columns": [
|
||||
{
|
||||
"expression": "company_id",
|
||||
"isExpression": false,
|
||||
"asc": true,
|
||||
"nulls": "last"
|
||||
},
|
||||
{
|
||||
"expression": "(\"payload\"->>'issueId')",
|
||||
"asc": true,
|
||||
"isExpression": true,
|
||||
"nulls": "last"
|
||||
},
|
||||
{
|
||||
"expression": "kind",
|
||||
"isExpression": false,
|
||||
"asc": true,
|
||||
"nulls": "last"
|
||||
},
|
||||
{
|
||||
"expression": "status",
|
||||
"isExpression": false,
|
||||
"asc": true,
|
||||
"nulls": "last"
|
||||
}
|
||||
],
|
||||
"isUnique": false,
|
||||
"where": "\"chat_actions\".\"kind\" IN ('inbound_wakeup', 'failed_run_retry')",
|
||||
"concurrently": true,
|
||||
"method": "btree",
|
||||
"with": {}
|
||||
},
|
||||
"chat_actions_github_stale_head_idx": {
|
||||
"name": "chat_actions_github_stale_head_idx",
|
||||
"columns": [
|
||||
@@ -10955,6 +10989,28 @@
|
||||
"concurrently": false,
|
||||
"method": "btree",
|
||||
"with": {}
|
||||
},
|
||||
"chat_deliveries_github_automatic_delivery_idx": {
|
||||
"name": "chat_deliveries_github_automatic_delivery_idx",
|
||||
"columns": [
|
||||
{
|
||||
"expression": "endpoint_id",
|
||||
"isExpression": false,
|
||||
"asc": true,
|
||||
"nulls": "last"
|
||||
},
|
||||
{
|
||||
"expression": "(\"normalized_event\"->'githubAutomatic'->'context'->>'deliveryId')",
|
||||
"asc": true,
|
||||
"isExpression": true,
|
||||
"nulls": "last"
|
||||
}
|
||||
],
|
||||
"isUnique": false,
|
||||
"where": "(\"chat_deliveries\".\"normalized_event\"->'githubAutomatic'->'context'->>'deliveryId') is not null",
|
||||
"concurrently": true,
|
||||
"method": "btree",
|
||||
"with": {}
|
||||
}
|
||||
},
|
||||
"foreignKeys": {
|
||||
@@ -12393,6 +12449,44 @@
|
||||
"concurrently": false,
|
||||
"method": "btree",
|
||||
"with": {}
|
||||
},
|
||||
"chat_message_links_company_comment_idx": {
|
||||
"name": "chat_message_links_company_comment_idx",
|
||||
"columns": [
|
||||
{
|
||||
"expression": "company_id",
|
||||
"isExpression": false,
|
||||
"asc": true,
|
||||
"nulls": "last"
|
||||
},
|
||||
{
|
||||
"expression": "comment_id",
|
||||
"isExpression": false,
|
||||
"asc": true,
|
||||
"nulls": "last"
|
||||
}
|
||||
],
|
||||
"isUnique": false,
|
||||
"where": "\"chat_message_links\".\"comment_id\" is not null",
|
||||
"concurrently": true,
|
||||
"method": "btree",
|
||||
"with": {}
|
||||
},
|
||||
"chat_message_links_delivery_idx": {
|
||||
"name": "chat_message_links_delivery_idx",
|
||||
"columns": [
|
||||
{
|
||||
"expression": "delivery_id",
|
||||
"isExpression": false,
|
||||
"asc": true,
|
||||
"nulls": "last"
|
||||
}
|
||||
],
|
||||
"isUnique": false,
|
||||
"where": "\"chat_message_links\".\"delivery_id\" is not null",
|
||||
"concurrently": true,
|
||||
"method": "btree",
|
||||
"with": {}
|
||||
}
|
||||
},
|
||||
"foreignKeys": {
|
||||
@@ -29073,6 +29167,22 @@
|
||||
"concurrently": false,
|
||||
"method": "btree",
|
||||
"with": {}
|
||||
},
|
||||
"environment_leases_native_backup_scope_idx": {
|
||||
"name": "environment_leases_native_backup_scope_idx",
|
||||
"columns": [
|
||||
{
|
||||
"expression": "(\"metadata\"->'nativeHarnessBackup'->>'sessionScopeSha256')",
|
||||
"asc": true,
|
||||
"isExpression": true,
|
||||
"nulls": "last"
|
||||
}
|
||||
],
|
||||
"isUnique": false,
|
||||
"where": "(\"environment_leases\".\"metadata\"->'nativeHarnessBackup'->>'sessionScopeSha256') is not null",
|
||||
"concurrently": true,
|
||||
"method": "btree",
|
||||
"with": {}
|
||||
}
|
||||
},
|
||||
"foreignKeys": {
|
||||
@@ -35421,6 +35531,22 @@
|
||||
"method": "btree",
|
||||
"with": {}
|
||||
},
|
||||
"heartbeat_runs_pending_budget_cancellation_idx": {
|
||||
"name": "heartbeat_runs_pending_budget_cancellation_idx",
|
||||
"columns": [
|
||||
{
|
||||
"expression": "company_id",
|
||||
"isExpression": false,
|
||||
"asc": true,
|
||||
"nulls": "last"
|
||||
}
|
||||
],
|
||||
"isUnique": false,
|
||||
"where": "jsonb_typeof(\"heartbeat_runs\".\"result_json\"->'startupCancellation'->'budgetEnforcement') = 'object'\n and \"heartbeat_runs\".\"result_json\"->'startupCancellation'->'budgetEnforcement'->>'deliveredAt' is null",
|
||||
"concurrently": true,
|
||||
"method": "btree",
|
||||
"with": {}
|
||||
},
|
||||
"heartbeat_runs_tree_hold_suppression_idx": {
|
||||
"name": "heartbeat_runs_tree_hold_suppression_idx",
|
||||
"columns": [
|
||||
@@ -35459,6 +35585,27 @@
|
||||
"method": "btree",
|
||||
"with": {}
|
||||
},
|
||||
"heartbeat_runs_company_wake_comment_idx": {
|
||||
"name": "heartbeat_runs_company_wake_comment_idx",
|
||||
"columns": [
|
||||
{
|
||||
"expression": "company_id",
|
||||
"isExpression": false,
|
||||
"asc": true,
|
||||
"nulls": "last"
|
||||
},
|
||||
{
|
||||
"expression": "(\"context_snapshot\"->>'wakeCommentId')",
|
||||
"asc": true,
|
||||
"isExpression": true,
|
||||
"nulls": "last"
|
||||
}
|
||||
],
|
||||
"isUnique": false,
|
||||
"concurrently": true,
|
||||
"method": "btree",
|
||||
"with": {}
|
||||
},
|
||||
"heartbeat_runs_cost_accounting_pending_idx": {
|
||||
"name": "heartbeat_runs_cost_accounting_pending_idx",
|
||||
"columns": [
|
||||
@@ -35611,6 +35758,28 @@
|
||||
"method": "btree",
|
||||
"with": {}
|
||||
},
|
||||
"heartbeat_runs_company_native_session_idx": {
|
||||
"name": "heartbeat_runs_company_native_session_idx",
|
||||
"columns": [
|
||||
{
|
||||
"expression": "company_id",
|
||||
"isExpression": false,
|
||||
"asc": true,
|
||||
"nulls": "last"
|
||||
},
|
||||
{
|
||||
"expression": "native_session_id",
|
||||
"isExpression": false,
|
||||
"asc": true,
|
||||
"nulls": "last"
|
||||
}
|
||||
],
|
||||
"isUnique": false,
|
||||
"where": "\"heartbeat_runs\".\"native_session_id\" is not null",
|
||||
"concurrently": true,
|
||||
"method": "btree",
|
||||
"with": {}
|
||||
},
|
||||
"heartbeat_runs_company_issue_created_idx": {
|
||||
"name": "heartbeat_runs_company_issue_created_idx",
|
||||
"columns": [
|
||||
+7
@@ -3200,6 +3200,13 @@
|
||||
"when": 1794326875077,
|
||||
"tag": "0458_chat_admission_issue_index",
|
||||
"breakpoints": true
|
||||
},
|
||||
{
|
||||
"idx": 459,
|
||||
"version": "7",
|
||||
"when": 1794329161946,
|
||||
"tag": "0459_environment_lease_backup_scope_index",
|
||||
"breakpoints": true
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -54,5 +54,12 @@ export const environmentLeases = pgTable(
|
||||
heartbeatRunIdx: index("environment_leases_heartbeat_run_idx").on(table.heartbeatRunId),
|
||||
companyLastUsedIdx: index("environment_leases_company_last_used_idx").on(table.companyId, table.lastUsedAt),
|
||||
providerLeaseIdx: index("environment_leases_provider_lease_idx").on(table.providerLeaseId),
|
||||
// Leases that reference one native harness backup root by its scope digest
|
||||
// (backup expiration, retention and session retention checks). Partial:
|
||||
// most leases carry no backup stamp.
|
||||
nativeBackupScopeIdx: index("environment_leases_native_backup_scope_idx")
|
||||
.on(sql`(${table.metadata}->'nativeHarnessBackup'->>'sessionScopeSha256')`)
|
||||
.where(sql`(${table.metadata}->'nativeHarnessBackup'->>'sessionScopeSha256') is not null`)
|
||||
.concurrently(),
|
||||
}),
|
||||
);
|
||||
@@ -0,0 +1,181 @@
|
||||
/**
|
||||
* Local proof for the native backup lease lookup (migration 0459).
|
||||
*
|
||||
* `reconcileBackup` reads the leases that reference one backup root: the
|
||||
* session's run leases, leases whose `metadata.nativeHarnessBackup` stamp has
|
||||
* the root's scope digest, and the identity's own lease. The metadata arm had
|
||||
* no index, so PostgreSQL read every lease of the company and its metadata.
|
||||
*
|
||||
* The script captures the SQL that the real service generates, then measures
|
||||
* it without and with the new index (no index, index, no index again). It
|
||||
* fails unless the service's lease IDs and the complete query rows match.
|
||||
*
|
||||
* pnpm --filter @paperclipai/server exec tsx ../scripts/benchmark-environment-lease-backup-lookup.ts
|
||||
*
|
||||
* Optional sizes: LEASES (default 100000), STAMPED_PERCENT (30),
|
||||
* METADATA_KIB (4), OTHER_COMPANY_LEASES (20000).
|
||||
*/
|
||||
import assert from "node:assert/strict";
|
||||
import { createHash, randomUUID } from "node:crypto";
|
||||
import { promises as fs, readFileSync } from "node:fs";
|
||||
import os from "node:os";
|
||||
import path from "node:path";
|
||||
import { agents, companies, createDb, environmentLeases, heartbeatRuns, startEmbeddedPostgresTestDatabase, type Db } from "../packages/db/src/index.js";
|
||||
import { nativeHarnessBackupExpirationService } from "../server/src/services/native-runtime/native-harness-backup-expiration.js";
|
||||
|
||||
const size = (name: string, fallback: number) => Number(process.env[name] ?? fallback);
|
||||
const LEASES = size("LEASES", 100_000), STAMPED_PERCENT = size("STAMPED_PERCENT", 30);
|
||||
const METADATA_KIB = size("METADATA_KIB", 4), OTHER_COMPANY_LEASES = size("OTHER_COMPANY_LEASES", 20_000);
|
||||
const migration = readFileSync(process.env.INDEX_SQL_FILE ?? new URL(
|
||||
"../packages/db/src/migrations/0459_environment_lease_backup_scope_index.sql", import.meta.url), "utf8");
|
||||
const indexStatements = migration.split("--> statement-breakpoint").map(statement => statement.trim()).filter(Boolean);
|
||||
const indexNames = indexStatements.map(statement => statement.match(/IF NOT EXISTS "([a-z_]+)"/)![1]!);
|
||||
|
||||
type Captured = { sql: string; params: unknown[] };
|
||||
/** Records the SQL of each `db.select(...).from(...)` the service awaits. */
|
||||
function capture(db: Db, queries: Captured[]): Db {
|
||||
return new Proxy(db, { get(target, key) {
|
||||
if (key !== "select") return Reflect.get(target, key);
|
||||
return (...fields: Parameters<Db["select"]>) => {
|
||||
const builder = target.select(...fields), from = builder.from.bind(builder);
|
||||
builder.from = ((...args: Parameters<typeof from>) => {
|
||||
const selection = from(...args), then = selection.then.bind(selection);
|
||||
selection.then = ((...handlers: Parameters<typeof then>) => {
|
||||
queries.push(selection.toSQL());
|
||||
return then(...handlers);
|
||||
}) as typeof selection.then;
|
||||
return selection;
|
||||
}) as typeof builder.from;
|
||||
return builder;
|
||||
};
|
||||
}});
|
||||
}
|
||||
|
||||
const keepAlive = setInterval(() => {}, 1_000);
|
||||
let temporary: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>> | undefined;
|
||||
let db: ReturnType<typeof createDb> | undefined;
|
||||
let stateBase: string | undefined;
|
||||
try {
|
||||
temporary = await startEmbeddedPostgresTestDatabase("environment-lease-backup-lookup-");
|
||||
db = createDb(temporary.connectionString);
|
||||
const sql = db.$client;
|
||||
for (const name of indexNames) await sql.unsafe(`DROP INDEX IF EXISTS "${name}"`);
|
||||
await sql.unsafe("create extension if not exists pgcrypto");
|
||||
|
||||
const companyId = randomUUID(), otherCompanyId = randomUUID();
|
||||
await db.insert(companies).values([{ id: companyId, name: "Leases", issuePrefix: "LEA" },
|
||||
{ id: otherCompanyId, name: "Other", issuePrefix: "OTH" }]);
|
||||
const [agent] = await db.insert(agents).values({ companyId, name: "Coder" }).returning();
|
||||
const sessionId = randomUUID();
|
||||
// Three runs of the backed-up session; the root's identity names the first.
|
||||
const runs = await db.insert(heartbeatRuns).values([0, 1, 2].map(() => ({
|
||||
companyId, agentId: agent!.id, status: "succeeded", nativeSessionId: sessionId }))).returning();
|
||||
const scope = createHash("sha256").update(randomUUID()).digest("hex");
|
||||
|
||||
const loadStart = performance.now();
|
||||
// History: incompressible sandbox metadata; some leases carry another
|
||||
// session's backup stamp, most carry none.
|
||||
const seed = (company: string, count: number) => sql.unsafe(`insert into environment_leases
|
||||
(company_id, status, lease_policy, provider, provider_lease_id, metadata, created_at, last_used_at)
|
||||
select $1::uuid, case when g % 10 = 0 then 'active' else 'released' end, 'retain_on_failure', 'sandbox',
|
||||
'sbx-' || g, jsonb_build_object(
|
||||
'nativeHarnessBackup', case when g % 100 < $2::int then jsonb_build_object(
|
||||
'sessionScopeSha256', encode(digest(gen_random_uuid()::text, 'sha256'), 'hex'),
|
||||
'manifestSha256', encode(digest(g::text, 'sha256'), 'hex')) end,
|
||||
'sandbox', (select string_agg(encode(gen_random_bytes(768), 'base64'), '') from generate_series(1, $3::int) s(i) where g > 0)),
|
||||
now() - make_interval(secs => g), now() - make_interval(secs => g)
|
||||
from generate_series(1, $4::int) g`, [company, STAMPED_PERCENT, METADATA_KIB, count]);
|
||||
await seed(companyId, LEASES);
|
||||
await seed(otherCompanyId, OTHER_COMPANY_LEASES);
|
||||
// References to the root: a lease of each run; one lease known only by its
|
||||
// backup stamp (its run row was deleted); the identity's own lease; and the
|
||||
// same digest in another company, which must stay out.
|
||||
const runLeases = await db.insert(environmentLeases).values(runs.map(run => ({
|
||||
companyId, heartbeatRunId: run.id, status: "released", metadata: { sandbox: "run" } }))).returning();
|
||||
const [stamped] = await db.insert(environmentLeases).values({ companyId, status: "retained",
|
||||
metadata: { nativeHarnessBackup: { sessionScopeSha256: scope } } }).returning();
|
||||
const [own] = await db.insert(environmentLeases).values({ companyId, status: "active", metadata: {} }).returning();
|
||||
await db.insert(environmentLeases).values({ companyId: otherCompanyId, status: "retained",
|
||||
metadata: { nativeHarnessBackup: { sessionScopeSha256: scope } } });
|
||||
await sql.unsafe("vacuum analyze");
|
||||
const loadMs = performance.now() - loadStart;
|
||||
const expectedLeaseIds = [...runLeases.map(lease => lease.id), stamped!.id, own!.id].sort();
|
||||
|
||||
stateBase = await fs.mkdtemp(path.join(os.tmpdir(), "environment-lease-backup-lookup-"));
|
||||
const root = path.join(stateBase, scope);
|
||||
await fs.mkdir(path.join(root, "control-plane"), { recursive: true });
|
||||
await fs.mkdir(path.join(root, "failover-backups", "current"), { recursive: true });
|
||||
await fs.writeFile(path.join(root, "control-plane", "control-plane-state.json"), JSON.stringify({
|
||||
identity: { runId: runs[0]!.id, normalizedSessionId: sessionId, environmentLeaseId: own!.id } }));
|
||||
const service = (target: Db) => nativeHarnessBackupExpirationService(target, {
|
||||
stateBase, closedForDays: 7, isSessionBusy: () => false });
|
||||
|
||||
const queries: Captured[] = [];
|
||||
const reconciled = await service(capture(db, queries)).reconcileBackup(scope);
|
||||
assert.deepEqual([...reconciled!.dependencies.leaseIds].sort(), expectedLeaseIds);
|
||||
const leaseQuery = queries.find(query => query.sql.includes('from "environment_leases"') && query.sql.includes("sessionScopeSha256"));
|
||||
assert.ok(leaseQuery, "The service must read the root's leases");
|
||||
|
||||
const median = (values: number[]) => [...values].sort((a, b) => a - b)[Math.floor(values.length / 2)]!;
|
||||
const round = (value: number) => Number(value.toFixed(3));
|
||||
const summary = (samples: number[]) => ({ samplesMs: samples.map(round), medianMs: round(median(samples)),
|
||||
minMs: round(Math.min(...samples)), maxMs: round(Math.max(...samples)) });
|
||||
async function explain() {
|
||||
const [result] = await sql.unsafe(`EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) ${leaseQuery!.sql}`, leaseQuery!.params as never[]);
|
||||
return (result!["QUERY PLAN"] as Array<{ Plan: Record<string, unknown>; "Execution Time": number }>)[0]!;
|
||||
}
|
||||
function nodes(plan: Record<string, unknown>, out: string[] = []) {
|
||||
out.push(`${plan["Node Type"]}${plan["Index Name"] ? ` using ${plan["Index Name"]}` : ""}`);
|
||||
for (const child of (plan.Plans as Array<Record<string, unknown>> | undefined) ?? []) nodes(child, out);
|
||||
return out;
|
||||
}
|
||||
async function measure(indexed: boolean) {
|
||||
const builds: Array<{ index: string; ms: number }> = [];
|
||||
if (!indexed) for (const name of indexNames) await sql.unsafe(`DROP INDEX CONCURRENTLY IF EXISTS "${name}"`);
|
||||
else for (const [n, statement] of indexStatements.entries()) {
|
||||
const start = performance.now();
|
||||
await sql.unsafe(statement);
|
||||
builds.push({ index: indexNames[n]!, ms: Math.round(performance.now() - start) });
|
||||
}
|
||||
await sql.unsafe("analyze environment_leases");
|
||||
const rows = (await sql.unsafe(leaseQuery!.sql, leaseQuery!.params as never[]).values())
|
||||
.map(row => row.map(value => value?.toString() ?? null)).sort();
|
||||
const ids = [...(await service(db!).reconcileBackup(scope))!.dependencies.leaseIds].sort();
|
||||
for (let n = 0; n < 2; n++) await explain();
|
||||
const execution: number[] = [], wall: number[] = [];
|
||||
for (let n = 0; n < 11; n++) {
|
||||
execution.push((await explain())["Execution Time"]);
|
||||
const start = performance.now();
|
||||
await service(db!).reconcileBackup(scope);
|
||||
wall.push(performance.now() - start);
|
||||
}
|
||||
const plan = (await explain()).Plan;
|
||||
return { indexed, builds, rows, ids, serverExecution: summary(execution), reconcileBackupWall: summary(wall),
|
||||
plan: { nodes: nodes(plan), sharedHitBlocks: plan["Shared Hit Blocks"], sharedReadBlocks: plan["Shared Read Blocks"] } };
|
||||
}
|
||||
|
||||
const results = [await measure(false), await measure(true), await measure(false), await measure(true)];
|
||||
for (const result of results) {
|
||||
assert.deepEqual(result.ids, expectedLeaseIds, "Service lease IDs must match");
|
||||
assert.deepEqual(result.rows, results[0]!.rows, "Complete query rows must match");
|
||||
}
|
||||
console.log(JSON.stringify({
|
||||
postgresVersion: (await sql.unsafe("select version()"))[0]!.version,
|
||||
settings: await sql.unsafe("select name, setting from pg_settings where name in ('shared_buffers','work_mem','jit') order by name"),
|
||||
fixture: { synthetic: true, leases: LEASES, stampedPercent: STAMPED_PERCENT, metadataKiB: METADATA_KIB,
|
||||
otherCompanyLeases: OTHER_COMPANY_LEASES, loadMs: Math.round(loadMs),
|
||||
tableSize: (await sql.unsafe("select pg_size_pretty(pg_total_relation_size('environment_leases')) as size"))[0]!.size,
|
||||
indexSizes: await sql.unsafe(`select c.relname, pg_size_pretty(pg_relation_size(c.oid)) as size from pg_class c
|
||||
where c.relname in (${indexNames.map((_, n) => `$${n + 1}`).join(", ")})`, indexNames) },
|
||||
statement: leaseQuery.sql,
|
||||
correctness: { serviceLeaseIdsEqual: true, completeRowsEqual: true, matchedLeases: expectedLeaseIds.length,
|
||||
cases: ["run leases", "stamp-only lease", "identity lease", "same digest in another company excluded"] },
|
||||
results: results.map(({ rows: _rows, ids: _ids, ...result }) => result),
|
||||
}, null, 2));
|
||||
} finally {
|
||||
try {
|
||||
await db?.$client.end({ timeout: 0 });
|
||||
await temporary?.cleanup();
|
||||
if (stateBase) await fs.rm(stateBase, { recursive: true, force: true });
|
||||
} finally { clearInterval(keepAlive); }
|
||||
}
|
||||
@@ -985,6 +985,60 @@ describeEmbeddedPostgres("native harness backup expiration", () => {
|
||||
await quiet();
|
||||
});
|
||||
|
||||
it("finds a backup root's leases by run, backup stamp and identity through indexes, within one company", async () => {
|
||||
const session = await seedSession({ leaseStatus: "retained" });
|
||||
const [stamped] = await db.select({ id: environmentLeases.id }).from(environmentLeases)
|
||||
.where(eq(environmentLeases.companyId, session.companyId));
|
||||
const [runLease] = await db.insert(environmentLeases).values({ companyId: session.companyId,
|
||||
heartbeatRunId: session.runId, status: "released", metadata: {} }).returning();
|
||||
const [own] = await db.insert(environmentLeases).values({ companyId: session.companyId, status: "active",
|
||||
metadata: { nativeHarnessBackup: { sessionScopeSha256: "0".repeat(64) } } }).returning();
|
||||
// Not references: another digest, a stamp that is not a string match, and
|
||||
// the same digest in another company.
|
||||
const otherCompanyId = randomUUID();
|
||||
await db.insert(companies).values({ id: otherCompanyId, name: "Other", issuePrefix: `O${otherCompanyId.slice(0, 6)}` });
|
||||
await db.insert(environmentLeases).values([
|
||||
{ companyId: session.companyId, status: "retained", metadata: { nativeHarnessBackup: { sessionScopeSha256: "f".repeat(64) } } },
|
||||
{ companyId: session.companyId, status: "retained", metadata: { nativeHarnessBackup: { sessionScopeSha256: [session.scope] } } },
|
||||
{ companyId: session.companyId, status: "retained", metadata: { sessionScopeSha256: session.scope } },
|
||||
{ companyId: otherCompanyId, status: "retained", metadata: { nativeHarnessBackup: { sessionScopeSha256: session.scope } } },
|
||||
]);
|
||||
// Unrelated lease history, so the planner prefers indexes to a company scan.
|
||||
await db.$client.unsafe(`insert into environment_leases (company_id, status, metadata)
|
||||
select $1::uuid, 'released', jsonb_build_object('sandbox', repeat('x', 512)) from generate_series(1, 2000)`, [session.companyId]);
|
||||
await db.$client.unsafe("analyze environment_leases");
|
||||
const statePath = path.join(session.root, "control-plane", "control-plane-state.json");
|
||||
await fs.writeFile(statePath, JSON.stringify({ identity: {
|
||||
runId: session.runId, normalizedSessionId: session.sessionId, environmentLeaseId: own!.id } }));
|
||||
|
||||
const queries: Array<{ text: string; params: unknown[] }> = [];
|
||||
const previous = db.$client.options.debug;
|
||||
db.$client.options.debug = (_connection, text, params) => { queries.push({ text, params: params as unknown[] }); };
|
||||
let reconciled: Awaited<ReturnType<ReturnType<typeof service>["reconcileBackup"]>>;
|
||||
try { reconciled = await service().reconcileBackup(session.scope); } finally { db.$client.options.debug = previous; }
|
||||
expect([...reconciled!.dependencies.leaseIds].sort()).toEqual([stamped!.id, runLease!.id, own!.id].sort());
|
||||
|
||||
// Every arm of the company's lease lookup must be index-backed, or the
|
||||
// planner reads each lease's metadata (environment_leases_native_backup_scope_idx).
|
||||
const lookup = queries.find(query => query.text.startsWith('select "id" from "environment_leases"')
|
||||
&& query.text.includes("sessionScopeSha256"));
|
||||
expect(lookup).toBeDefined();
|
||||
const custom = JSON.stringify(await db.$client.unsafe(`explain (format json) ${lookup!.text}`, lookup!.params as never[]));
|
||||
// A reused prepared statement gets a generic plan; the bound digest must
|
||||
// still prove the partial index predicate.
|
||||
const generic = JSON.stringify(await db.$client.begin(async tx => {
|
||||
await tx.unsafe("set local plan_cache_mode = force_generic_plan");
|
||||
await tx.unsafe(`prepare backup_lease_lookup as ${lookup!.text}`);
|
||||
const values = lookup!.params.map(value => `'${String(value).replaceAll("'", "''")}'`).join(", ");
|
||||
return tx.unsafe(`explain (format json) execute backup_lease_lookup(${values})`);
|
||||
}));
|
||||
for (const text of [custom, generic]) {
|
||||
expect(text).toContain("environment_leases_native_backup_scope_idx");
|
||||
expect(text).toContain("BitmapOr");
|
||||
expect(text).not.toContain("Seq Scan");
|
||||
}
|
||||
});
|
||||
|
||||
it("rejects destructive backup scopes outside the controller digest namespace", async () => {
|
||||
await emptyStateBase(); const remove = vi.spyOn(fs, "rm");
|
||||
await expect(service().expireBackup("../outside")).rejects.toThrow("Invalid native harness backup scope");
|
||||
|
||||
Reference in new issue
Block a user