diff --git a/server/src/__tests__/environment-runtime.test.ts b/server/src/__tests__/environment-runtime.test.ts index 03abe6b72cf..d92eb32aaa4 100644 --- a/server/src/__tests__/environment-runtime.test.ts +++ b/server/src/__tests__/environment-runtime.test.ts @@ -5,6 +5,7 @@ import { createSandboxOrphanCleanupSpool } from "../services/sandbox-orphan-clea import { idleWorkSnapshot } from "../services/task-admission.js"; import { remoteTerminationReceipt } from "../services/remote-execution-termination.js"; import { createHash, randomUUID } from "node:crypto"; +import { spawn } from "node:child_process"; import { mkdir, mkdtemp, readdir, rm, writeFile } from "node:fs/promises"; import os from "node:os"; import path from "node:path"; @@ -5306,6 +5307,58 @@ describeEmbeddedPostgres("environmentRuntimeService", () => { }, ); + it("retries a run-release teardown at the next sweep after the releasing server dies", async () => { + // Shape of a chaos incident: the run finished, its release claimed the lease + // and called the provider, and the server was SIGKILLed before the reply. + const seeded = await seedReusablePluginSandboxLease(); + await db.update(heartbeatRuns).set({ status: "succeeded" }).where(eq(heartbeatRuns.id, seeded.runId)); + let enterProvider!: () => void; + let finishProvider!: () => void; + const entered = new Promise((resolve) => { enterProvider = resolve; }); + const finish = new Promise((resolve) => { finishProvider = resolve; }); + let serverAlive = true; + const call = vi.fn(async (_id: string, method: string) => { + if (method !== "environmentDestroyLease") throw new Error(`Unexpected method ${method}`); + if (serverAlive) { + enterProvider(); + await finish; + throw new Error("Server stopped"); + } + return { providerLeaseId: seeded.reusableLease.providerLeaseId, state: "destroyed" }; + }); + const makeRuntime = () => environmentRuntimeService(db, { pluginWorkerManager: { + isRunning: () => true, call, + getWorker: () => ({ supportedMethods: ["environmentDestroyLease"] }), + } as unknown as PluginWorkerManager }); + const closing = makeRuntime().releaseRunLeases(seeded.runId, "released", undefined, "destroy"); + try { + await entered; + const claimed = (await environmentService(db).getLeaseById(seeded.reusableLease.id))!; + expect(claimed).toMatchObject({ status: "pending_cleanup", metadata: { pendingCleanupInFlight: true } }); + const attemptId = claimed.metadata?.pendingCleanupAttemptId; + // The releasing server exits: its recorded process is gone. Mark it as an + // earlier boot whose PID has exited, as a restarted server observes it. + const child = spawn(process.execPath, ["-e", ""], { stdio: "ignore" }); + await new Promise(resolve => child.once("exit", resolve)); + const recorded = claimed.metadata?.pendingCleanupController as Record | undefined; + expect(recorded).toMatchObject({ attemptId }); + await db.update(environmentLeases).set({ metadata: { ...claimed.metadata, + pendingCleanupController: { ...recorded, bootId: randomUUID(), pid: child.pid } } }) + .where(eq(environmentLeases.id, claimed.id)); + serverAlive = false; + + const restarted = heartbeatService(db, { environmentRuntime: makeRuntime() }); + expect(await restarted.sweepPendingCleanupLeases({ backoffMs: 0 })).toEqual({ swept: 1, destroyed: 1, capped: 0 }); + const settled = (await environmentService(db).getLeaseById(claimed.id))!; + expect(settled).toMatchObject({ status: "expired", cleanupStatus: "success" }); + // The interrupted release was not a failed teardown attempt. + expect(settled.metadata?.pendingCleanupRetryAttempts).toBe(1); + } finally { + finishProvider(); + await closing.catch(() => undefined); + } + }); + it.each([false, true])("preserves a newer destroy receipt when an old run cleanup returns (failure=%s)", async failure => { const seeded = await seedReusablePluginSandboxLease(); const service = environmentService(db); diff --git a/server/src/__tests__/heartbeat-pending-cleanup-sweep.test.ts b/server/src/__tests__/heartbeat-pending-cleanup-sweep.test.ts index a637dc89bec..fd83dac4a70 100644 --- a/server/src/__tests__/heartbeat-pending-cleanup-sweep.test.ts +++ b/server/src/__tests__/heartbeat-pending-cleanup-sweep.test.ts @@ -21,6 +21,25 @@ import { const mockTelemetryClient = vi.hoisted(() => ({ track: vi.fn() })); vi.mock("../telemetry.ts", () => ({ getTelemetryClient: () => mockTelemetryClient })); +// A test can hold every caller of the exited-process proof until N callers +// arrive, so overlapping sweeps all read a claim before any of them updates it. +const exitProofBarrier = vi.hoisted(() => ({ parties: 0, waiting: [] as Array<() => void> })); +vi.mock("../services/controller-process-identity.js", async importOriginal => { + const actual = await importOriginal(); + return { ...actual, controllerProcessExited: async (...args: Parameters) => { + if (exitProofBarrier.parties > 1) { + await new Promise(resolve => { + exitProofBarrier.waiting.push(resolve); + if (exitProofBarrier.waiting.length >= exitProofBarrier.parties) { + for (const release of exitProofBarrier.waiting.splice(0)) release(); + exitProofBarrier.parties = 0; + } + }); + } + return actual.controllerProcessExited(...args); + } }; +}); + vi.mock("../middleware/logger.js", () => ({ logger: { child: vi.fn(function child(this: unknown) { @@ -380,6 +399,34 @@ describeEmbeddedPostgres("heartbeat sweepPendingCleanupLeases", () => { expect(teardown).toHaveBeenCalledTimes(1); expect(result.destroyed).toBe(1); }); + it("does not count a sweep attempt lost to a dead server as a failed attempt", async () => { + // The lost sweep claim had already counted attempt 3; the retry is attempt 3 again. + const scope = await seedCompanyAndEnvironment(); + const leaseId = await insertOrphanEphemeralLease({ ...scope, updatedAt: new Date(Date.now() - 60_000), + metadata: { ...claimMetadata({ ...await exitedClaimer(), countedAttempt: true }), [ATTEMPTS_KEY]: 3 } }); + const teardown = vi.fn(async () => { throw new Error("Provider unavailable"); }); + const service = heartbeatService(db, { environmentRuntime: { + retryPendingSandboxTeardown: teardown, + } as unknown as HeartbeatEnvironmentRuntime }); + await service.sweepPendingCleanupLeases(); + expect(teardown).toHaveBeenCalledOnce(); + expect(await readAttempts(leaseId)).toBe(3); + }); + it("refunds a lost sweep attempt only once when recovery sweeps overlap", async () => { + const scope = await seedCompanyAndEnvironment(); + const leaseId = await insertOrphanEphemeralLease({ ...scope, updatedAt: new Date(Date.now() - 60_000), + metadata: { ...claimMetadata({ ...await exitedClaimer(), countedAttempt: true }), [ATTEMPTS_KEY]: 3 } }); + const teardown = vi.fn(async () => { throw new Error("Provider unavailable"); }); + const runtime = { retryPendingSandboxTeardown: teardown } as unknown as HeartbeatEnvironmentRuntime; + exitProofBarrier.parties = 2; + await Promise.all([ + heartbeatService(db, { environmentRuntime: runtime }).sweepPendingCleanupLeases(), + heartbeatService(db, { environmentRuntime: runtime }).sweepPendingCleanupLeases(), + ]); + // One refund (3 -> 2) and one new failed attempt (2 -> 3); never a second refund. + expect(teardown).toHaveBeenCalledOnce(); + expect(await readAttempts(leaseId)).toBe(3); + }); it("keeps the fence of a claim held by this live server", async () => { const { leaseId, teardown } = await sweepClaimedBy({ ...await currentControllerProcess(), attemptId: "lost-claim" }); expect(teardown).not.toHaveBeenCalled(); diff --git a/server/src/services/environment-runtime.ts b/server/src/services/environment-runtime.ts index dfd42af573c..149923a2ca1 100644 --- a/server/src/services/environment-runtime.ts +++ b/server/src/services/environment-runtime.ts @@ -21,6 +21,7 @@ import { remoteTerminationReceipt } from "./remote-execution-termination.js"; import { inspectCloudSandboxCleanupAbsence } from "./sandbox-absence-proof.js"; import { hasNativeWorkspaceExportResume, releaseCompletedNativeWorkspaceExportRetention } from "./native-runtime/native-workspace-export-resume.js"; import { createHash, randomUUID } from "node:crypto"; +import { currentControllerProcess } from "./controller-process-identity.js"; import { and, eq, inArray, ne, or, sql } from "drizzle-orm"; import type { Db } from "@paperclipai/db"; import { companySecrets, companySecretVersions, environmentLeases, heartbeatRuns } from "@paperclipai/db"; @@ -3685,6 +3686,10 @@ export function environmentRuntimeService( if (!(await allowLegacyShutdownWorkspaceCleanup(db, input.leaseRow))) return null; const now = new Date(); const attemptId = randomUUID(); + // The claimer is bound to this attempt. If this server dies mid-teardown, + // the startup sweep proves it exited and retries now instead of waiting + // out the 15-minute ownership fence. + const controller = { ...await currentControllerProcess(), attemptId }; // Persist recovery ownership before any provider work. Re-check scope and // run liveness in this write: the earlier selection cannot fence a resume. // The in-flight lease also excludes concurrent closures and cleanup sweeps. @@ -3698,7 +3703,7 @@ export function environmentRuntimeService( metadata: sql`(case when jsonb_typeof(${environmentLeases.metadata}) = 'object' then ${environmentLeases.metadata} else '{}'::jsonb end - 'remoteExecutionTermination') || ${JSON.stringify({ pendingCleanupAttemptId: attemptId, pendingCleanupInFlight: true, - pendingCleanupLeaseExpiresAtMs: Date.now() + 15 * 60_000 })}::jsonb`, + pendingCleanupLeaseExpiresAtMs: Date.now() + 15 * 60_000, pendingCleanupController: controller })}::jsonb`, }).where(and( eq(environmentLeases.id, input.leaseRow.id), eq(environmentLeases.companyId, input.leaseRow.companyId), @@ -3931,10 +3936,12 @@ export function environmentRuntimeService( if (hasStopOnlyCleanup(leaseSnapshot)) { const intentKey = stopOnlyCleanupKey(leaseSnapshot); const attemptId = randomUUID(); + const controller = { ...await currentControllerProcess(), attemptId }; const [claimed] = await writeEnvironmentCleanup(db, tx => tx.update(environmentLeases).set({ status: "pending_cleanup", cleanupStatus: "failed", releasedAt: new Date(), metadata: sql`${environmentLeases.metadata} || ${JSON.stringify({ pendingCleanupAttemptId: attemptId, - pendingCleanupInFlight: true, pendingCleanupLeaseExpiresAtMs: Date.now() + 15 * 60_000 })}::jsonb`, + pendingCleanupInFlight: true, pendingCleanupLeaseExpiresAtMs: Date.now() + 15 * 60_000, + pendingCleanupController: controller })}::jsonb`, }).where(and(eq(environmentLeases.id, leaseRow.id), eq(environmentLeases.heartbeatRunId, heartbeatRunId), inArray(environmentLeases.status, ["active", "pending_cleanup"]), sql`coalesce(${environmentLeases.metadata}->>'pendingCleanupInFlight', 'false') = 'false'`, diff --git a/server/src/services/heartbeat/recovery.ts b/server/src/services/heartbeat/recovery.ts index 83f533df2c8..f52151df613 100644 --- a/server/src/services/heartbeat/recovery.ts +++ b/server/src/services/heartbeat/recovery.ts @@ -1463,13 +1463,26 @@ export function createHeartbeatRecovery(db: Db, dependencies: HeartbeatRecoveryD const controller = readControllerProcess(saved); const attemptId = saved?.attemptId; if (!controller || typeof attemptId !== "string" || !await controllerProcessExited(controller)) continue; + // A dead server is not a failed teardown. A sweep claim already counted + // its attempt, so give it back; claims from other paths never counted. + // The update is single-use: it matches only while the fence is still in + // the future and it consumes `countedAttempt`, so an overlapping sweep + // that read the same claim cannot refund the attempt again. + const returnAttempt = saved?.countedAttempt === true; await db.update(environmentLeases) - .set({ metadata: sql`jsonb_set(${pendingCleanupMetadataObjectSql()}, '{pendingCleanupLeaseExpiresAtMs}', to_jsonb(${nowMs}::bigint))` }) + .set({ metadata: sql`jsonb_set(jsonb_set(jsonb_set(${pendingCleanupMetadataObjectSql()}, + '{pendingCleanupLeaseExpiresAtMs}', to_jsonb(${nowMs}::bigint)), + array[${PENDING_CLEANUP_ATTEMPTS_METADATA_KEY}], + to_jsonb(greatest(${pendingCleanupAttemptsSql()} - ${returnAttempt ? 1 : 0}, 0)::int), true), + '{pendingCleanupController,countedAttempt}', 'false'::jsonb, true)` }) .where(and( eq(environmentLeases.id, row.id), eq(environmentLeases.status, "pending_cleanup"), + sql`${environmentLeases.metadata}->>'pendingCleanupInFlight' = 'true'`, + sql`(${environmentLeases.metadata}->>'pendingCleanupLeaseExpiresAtMs')::numeric > ${nowMs}`, sql`${environmentLeases.metadata}->>'pendingCleanupAttemptId' = ${attemptId}`, sql`${environmentLeases.metadata}->'pendingCleanupController'->>'attemptId' = ${attemptId}`, + sql`coalesce(${environmentLeases.metadata}->'pendingCleanupController'->>'countedAttempt', 'false') = ${returnAttempt ? "true" : "false"}`, )); } if (claimed.length < PENDING_CLEANUP_SWEEP_PAGE_SIZE) return; @@ -1503,7 +1516,8 @@ export function createHeartbeatRecovery(db: Db, dependencies: HeartbeatRecoveryD metadata: sql`jsonb_set(${pendingCleanupMetadataObjectSql()}, array[${PENDING_CLEANUP_ATTEMPTS_METADATA_KEY}], to_jsonb(${expectedAttempts + 1}::int), true) || ${JSON.stringify({ ...(manualAttempt ? { pendingCleanupManualAttemptId: randomUUID() } : {}), pendingCleanupAttemptId: attemptId, pendingCleanupInFlight: true, - pendingCleanupLeaseExpiresAtMs: Date.now() + 15 * 60_000, pendingCleanupController: { ...controller, attemptId }, + pendingCleanupLeaseExpiresAtMs: Date.now() + 15 * 60_000, + pendingCleanupController: { ...controller, attemptId, countedAttempt: true }, })}::jsonb`, lastUsedAt: now, updatedAt: now,