diff --git a/server/src/services/native-runtime/native-provider-pack-admission.test.ts b/server/src/services/native-runtime/native-provider-pack-admission.test.ts index 5d5af36841e..8ac60525f12 100644 --- a/server/src/services/native-runtime/native-provider-pack-admission.test.ts +++ b/server/src/services/native-runtime/native-provider-pack-admission.test.ts @@ -226,6 +226,37 @@ describe("provider pack admission PostgreSQL fencing", () => { } finally { release(); await owner; } await expect(recordNativeProviderPackAdmission(db, f)).resolves.toBe(true); }); + it.each([false, true])("waits for a concurrent run-event append's run-row lock instead of losing the admission (waitForClaim: %s)", async waitForClaim => { + // Staging failure: a run-event append locks the run row + // (FOR NO KEY UPDATE, as appendHeartbeatRunEvent does) for milliseconds. A + // NOWAIT admission write failed, left the state pending, and artifact + // retention then failed the run as native_provider_artifacts_admission_unverified. + const f = await fixture(); + let release!: () => void, acquired!: () => void; + const holding = new Promise(resolve => { acquired = resolve; }); + const gate = new Promise(resolve => { release = resolve; }); + const append = db.transaction(async tx => { + await tx.select({ id: heartbeatRuns.id }).from(heartbeatRuns).where(eq(heartbeatRuns.id, f.identity.runId)).for("no key update"); + acquired(); await gate; + }); + await holding; + let settled = false; + const recording = recordNativeProviderPackAdmission(db, { ...f, waitForClaim } as Parameters[1]) + .then(value => { settled = true; return { value }; }, error => { settled = true; return { error }; }); + try { + // Release only after the admission is proven to wait on the run row, so + // the old NOWAIT code cannot pass by running after the append commits. + await vi.waitFor(async () => { + const [row] = await db.execute<{ waiting: boolean }>(sql`select exists(select 1 from pg_stat_activity + where datname = current_database() and wait_event_type = 'Lock' and state = 'active' + and query like '%heartbeat_runs%') as waiting`); + expect(settled || row.waiting).toBe(true); + }, { timeout: 800 }); + } finally { release(); await append; } + expect(await recording).toEqual({ value: true }); + expect((await read(f.identity.runId)).runnerProfileJson?.nativeProviderPackAdmissionState) + .toEqual({ ...f.profile.nativeProviderPackAdmissionState, state: "verified" }); + }); it.each([false, true])("waits within the durable-capture bound and rechecks the claim after its lock (revoked: %s)", async revoked => { const f = await fixture(); let release!: () => void, acquired!: () => void; diff --git a/server/src/services/native-runtime/native-provider-pack-admission.ts b/server/src/services/native-runtime/native-provider-pack-admission.ts index 8476f852b10..d815e427f63 100644 --- a/server/src/services/native-runtime/native-provider-pack-admission.ts +++ b/server/src/services/native-runtime/native-provider-pack-admission.ts @@ -132,8 +132,11 @@ export async function recordNativeProviderPackAdmission(db: Db, input: { }): Promise { const receipt = admittedNativeProviderPack(input.session, input.identity); return db.transaction(async tx => { + // Bounds every lock wait below, including the run-row lock. A run-event + // append holds that row FOR UPDATE for milliseconds; NOWAIT there lost the + // admission and then failed the run as admission_unverified. + await tx.execute(sql`select set_config('lock_timeout', '1000ms', true)`); if (input.waitForClaim) { - await tx.execute(sql`select set_config('lock_timeout', '1000ms', true)`); await assertNativeRunnerCheckpointOwner(tx, { companyId: input.identity.companyId, issueId: input.identity.issueId, runId: input.identity.runId, authorityRunId: input.identity.runId, bootId: input.claim.controllerBootId, leaseOwner: input.claim.leaseOwner, @@ -151,7 +154,7 @@ export async function recordNativeProviderPackAdmission(db: Db, input: { eq(heartbeatRuns.id, input.identity.runId), eq(heartbeatRuns.companyId, input.identity.companyId), eq(heartbeatRuns.agentId, input.identity.agentId), eq(heartbeatRuns.nativeIssueId, input.identity.issueId), eq(heartbeatRuns.runtimeMode, "native"), eq(heartbeatRuns.status, "running"), - )).for("update", { noWait: true }).limit(1); + )).for("update").limit(1); if (!run) return false; const current = run.profile?.nativeProviderPackAdmissionState as Record | undefined; if (current?.schema !== "paperclip.native-provider-pack-admission-state.v1" || current.admissionEpoch !== input.admissionEpoch