diff --git a/server/src/services/heartbeat/recovery.test.ts b/server/src/services/heartbeat/recovery.test.ts index f031b88878c..e26bb112ccf 100644 --- a/server/src/services/heartbeat/recovery.test.ts +++ b/server/src/services/heartbeat/recovery.test.ts @@ -16,6 +16,7 @@ import type { NativeRestartRecoveryClaim, NativeRestartRecoveryDisposition } fro import { NATIVE_OWNERSHIP_UNVERIFIED_ERROR_CODE } from "../native-runtime/native-runner-ownership.js"; import { claimNativeRestartRecoveries, currentNativeControllerIdentity, NATIVE_RESTART_RECOVERY_FAILED } from "../native-runtime/native-restart-recovery.js"; import { startNativeSessionExecutionLeaseRenewal } from "../native-runtime/native-execution-lease.js"; +import { createRunLogWriterAuthority } from "./run-log-authority.js"; import { allowLegacyShutdownWorkspaceCleanup, beginLegacyShutdownWorkspaceSettlement, finishLegacyShutdownWorkspaceSettlement, maintainLegacyShutdownWorkspaceSettlement, legacyShutdownWorkspaceResourceProtected, sweepLegacyShutdownWorkspaceSettlements } from "../legacy-shutdown-workspace-settlement.js"; import { adapterExecutionControls, createAdapterExecutionControl } from "../adapter-execution-control.js"; import { currentControllerProcessStamp, currentProcessSpaceHash, currentProcessStartedAtMs } from "../controller-process-identity.js"; @@ -1319,6 +1320,67 @@ describePostgres("heartbeat recovery module ownership", () => { } }, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS); + it("keeps a remote bootstrap restart claim through preparation longer than its first lease, then hands it to the next attempt", async () => { + const run = await seedRun(); + const deps = callbacks(db); + const controller = await currentNativeControllerIdentity(); + const claim: NativeRestartRecoveryClaim = { kind: "bootstrap_incomplete", runId: run.id, leaseOwner: randomUUID(), + controllerGeneration: 3, providerAttempt: 1, restartKind: "hard", recoveryRequestId: null }; + const scope = { companyId: run.companyId, issueId: run.issueId!, controller, remote: true }; + await db.update(heartbeatRuns).set({ runtimeMode: "native", nativeIssueId: run.issueId, status: "running" }).where(eq(heartbeatRuns.id, run.id)); + const firstLeaseEndsAt = new Date(Date.now() + 4_000); + await db.insert(nativeRunFinalizations).values({ runId: run.id, companyId: run.companyId, issueId: run.issueId!, phase: "observed", attempt: 1, + controllerBootId: controller.bootId, controllerPid: controller.pid, controllerProcessStartedAt: controller.processStartedAt, + leaseOwner: claim.leaseOwner, leaseExpiresAt: firstLeaseEndsAt }); + vi.useFakeTimers({ toFake: ["setInterval", "clearInterval"] }); + native.claims.mockImplementationOnce(async input => { + await input.onClaim?.(claim, scope); + return [claim]; + }); + const [admitted] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, run.id)); + const authority = createRunLogWriterAuthority(db, admitted, claim.leaseOwner); + let finish!: () => void, report!: (outcome: unknown) => void; + const outcome = new Promise(resolve => { report = resolve; }); + const finished = new Promise(resolve => { finish = resolve; }); + deps.executeRun.mockImplementationOnce(async (_runId, options) => { try { + await vi.advanceTimersByTimeAsync(21_000); + await vi.waitFor(async () => { + const [coordinator] = await db.select().from(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, run.id)); + expect(coordinator.leaseExpiresAt!.getTime()).toBeGreaterThan(firstLeaseEndsAt.getTime() + 10_000); + }); + await new Promise(resolve => setTimeout(resolve, Math.max(0, firstLeaseEndsAt.getTime() - Date.now() + 1_000))); + await authority(); + await db.update(nativeRunFinalizations).set({ attempt: 2 }).where(eq(nativeRunFinalizations.runId, run.id)); + const execution = startNativeSessionExecutionLeaseRenewal({ db, ...scope, runId: run.id, + leaseOwner: claim.leaseOwner, attempt: 2, onLost: () => native.detach() }); + try { + await options?.onNativeExecutionLeaseAcquired?.(); + const [before] = await db.select().from(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, run.id)); + await vi.advanceTimersByTimeAsync(21_000); + await vi.waitFor(async () => { + const [after] = await db.select().from(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, run.id)); + expect(after.leaseExpiresAt!.getTime()).toBeGreaterThan(before.leaseExpiresAt!.getTime()); + }); + report("handed off"); + await finished; + } finally { await execution.stop(); } + } catch (error) { report(error); } }); + try { + native.detach.mockClear(); + await createHeartbeatRecovery(db, deps).recoverNativeRunsAfterRestart(new Date(), [run.id]); + expect(await outcome).toBe("handed off"); + expect(deps.executeRun).toHaveBeenCalledWith(run.id, expect.objectContaining({ nativeLeaseOwner: claim.leaseOwner, + onNativeExecutionLeaseAcquired: expect.any(Function) })); + await vi.waitFor(() => expect(native.detach).not.toHaveBeenCalled()); + } finally { + finish(); + await Promise.all([...deps.activeRunExecutionPromises]); + vi.useRealTimers(); + await db.delete(nativeRunFinalizations).where(eq(nativeRunFinalizations.runId, run.id)); + await db.delete(heartbeatRuns).where(eq(heartbeatRuns.id, run.id)); + } + }, EMBEDDED_POSTGRES_TEST_TIMEOUT_MS); + it.each(["bootstrap_incomplete", "blocked"] as const)("retries durable remote takeover and preserves its %s disposition", async kind => { const run = await seedRun(); const deps = callbacks(db); diff --git a/server/src/services/heartbeat/recovery.ts b/server/src/services/heartbeat/recovery.ts index 18798f49eeb..01cc437ce74 100644 --- a/server/src/services/heartbeat/recovery.ts +++ b/server/src/services/heartbeat/recovery.ts @@ -999,7 +999,7 @@ export function createHeartbeatRecovery(db: Db, dependencies: HeartbeatRecoveryD try { const dispositions = await claimNativeRestartRecoveries({ onClaim: (claim, scope) => { - if (!scope.remote || claim.kind !== "reattach_remote_runner") return; + if (!scope.remote) return; // Start before the rest of the batch or its event writes can delay dispatch. preparationRenewals.set(claim.runId, startNativeSessionExecutionLeaseRenewal({ db, ...scope, runId: claim.runId, leaseOwner: claim.leaseOwner, diff --git a/server/src/services/native-runtime/native-execution-lease.test.ts b/server/src/services/native-runtime/native-execution-lease.test.ts index 4159771e327..9b5f97c5caf 100644 --- a/server/src/services/native-runtime/native-execution-lease.test.ts +++ b/server/src/services/native-runtime/native-execution-lease.test.ts @@ -4,10 +4,13 @@ import { startNativeSessionExecutionLeaseRenewal } from "./native-execution-leas function fixture() { const returning = vi.fn(async () => [{ runId: "run" }]); - const update = vi.fn(() => ({ set: () => ({ where: () => ({ returning }) }) })); - const transaction = vi.fn(async (run: (tx: unknown) => Promise) => run({ update, execute: vi.fn() })); + const coordinator = vi.fn(async (): Promise> => []); + const renewed = vi.fn((_values: unknown) => undefined); + const update = vi.fn(() => ({ set: (values: unknown) => { renewed(values); return { where: (condition: unknown) => { renewed(condition); return { returning }; } }; } })); + const select = () => ({ from: () => ({ where: () => ({ limit: coordinator }) }) }); + const transaction = vi.fn(async (run: (tx: unknown) => Promise) => run({ update, select, execute: vi.fn() })); const db = { update, transaction }; - return { update, returning, transaction, input: { db: db as unknown as Db, runId: "run", companyId: "company", issueId: "issue", + return { update, returning, transaction, coordinator, input: { db: db as unknown as Db, runId: "run", companyId: "company", issueId: "issue", leaseOwner: "owner", attempt: 1, controller: { bootId: "boot", pid: 1, processStartedAt: new Date() }, remote: true, onLost: vi.fn(async () => undefined) } }; } @@ -88,6 +91,100 @@ describe("native execution renewal ownership", () => { expect(returning).toHaveBeenCalledTimes(1); expect(vi.getTimerCount()).toBe(0); }); + it("hands preparation renewal to the same owner's executor at the next attempt", async () => { + vi.useFakeTimers(); + const { input, returning } = fixture(); + const preparation = startNativeSessionExecutionLeaseRenewal(input); + expect(() => startNativeSessionExecutionLeaseRenewal({ ...input, attempt: 0 })).toThrow("native_session_renewal_binding_changed"); + const executorLost = vi.fn(async () => undefined); + const execution = startNativeSessionExecutionLeaseRenewal({ ...input, attempt: 2, onLost: executorLost }); + await preparation.stop(); + await vi.advanceTimersByTimeAsync(20_000); + expect(returning).toHaveBeenCalledTimes(1); + returning.mockResolvedValueOnce([]); + await vi.advanceTimersByTimeAsync(20_000); + expect(executorLost).toHaveBeenCalledTimes(1); + expect(input.onLost).not.toHaveBeenCalled(); + await expect(execution.stop()).rejects.toThrow("native_session_lease_lost"); + }); + it("gives the executor that takes over at the next attempt its own later deadline", async () => { + vi.useFakeTimers(); + const { input, returning } = fixture(); + let complete!: (value: { runId: string }[]) => void; + const preparation = startNativeSessionExecutionLeaseRenewal({ ...input, confirmedDeadline: performance.now() + 25_000 }); + returning.mockImplementationOnce(() => new Promise(resolve => { complete = resolve; })); + await vi.advanceTimersByTimeAsync(20_000); + const executorLost = vi.fn(async () => undefined); + const execution = startNativeSessionExecutionLeaseRenewal({ ...input, attempt: 2, onLost: executorLost, + confirmedDeadline: performance.now() + 60_000 }); + expect(execution.remainingMs()).toBe(60_000); + const released = preparation.stop().catch(() => undefined); + await vi.advanceTimersByTimeAsync(30_000); + expect(executorLost).not.toHaveBeenCalled(); + complete([{ runId: "run" }]); + await vi.advanceTimersByTimeAsync(0); + expect(execution.remainingMs()).toBeGreaterThanOrEqual(30_000); + await released; + await execution.stop(); + }); + it("reacquires a lapsed local lease with the generation of the attempt that took over", async () => { + vi.useFakeTimers(); + const { input, update, returning } = fixture(); + const conditions: unknown[] = []; + update.mockImplementation(() => ({ set: () => ({ where: (condition: unknown) => { conditions.push(condition); return { returning }; } }) })); + const local = { ...input, remote: false, controllerGeneration: 3 }; + const preparation = startNativeSessionExecutionLeaseRenewal(local); + const execution = startNativeSessionExecutionLeaseRenewal({ ...local, attempt: 2, controllerGeneration: 4 }); + await preparation.stop(); + await vi.advanceTimersByTimeAsync(5 * 60_000); + const values = (root: unknown) => { + const found = new Set(), seen = new WeakSet(); + const walk = (node: unknown) => { + if (node === null || typeof node !== "object") { found.add(node); return; } + if (seen.has(node)) return; seen.add(node); + for (const child of Object.values(node)) walk(child); + }; + walk(root); + return found; + }; + expect(conditions.length).toBeGreaterThan(0); + expect(values(conditions.at(-1)).has(4)).toBe(true); + expect(values(conditions.at(-1)).has(3)).toBe(false); + await execution.stop(); + }); + it("does not count a renewal that failed because the owner's executor advanced the attempt", async () => { + vi.useFakeTimers(); + const { input, returning } = fixture(); + let complete!: (value: { runId: string }[]) => void; + returning.mockImplementationOnce(() => new Promise(resolve => { complete = resolve; })); + const preparation = startNativeSessionExecutionLeaseRenewal(input); + await vi.advanceTimersByTimeAsync(20_000); + expect(returning).toHaveBeenCalledTimes(1); + const execution = startNativeSessionExecutionLeaseRenewal({ ...input, attempt: 2 }); + complete([]); + await vi.advanceTimersByTimeAsync(20_000); + expect(returning).toHaveBeenCalledTimes(2); + expect(input.onLost).not.toHaveBeenCalled(); + await preparation.stop(); + await execution.stop(); + expect(vi.getTimerCount()).toBe(0); + }); + it("keeps renewing when the coordinator shows its own owner already claimed the next attempt", async () => { + vi.useFakeTimers(); + const { input, returning, coordinator } = fixture(); + coordinator.mockResolvedValue([{ attempt: 2, leaseOwner: "owner" }]); + returning.mockResolvedValueOnce([]); + const preparation = startNativeSessionExecutionLeaseRenewal(input); + await vi.advanceTimersByTimeAsync(20_000); + expect(coordinator).toHaveBeenCalledTimes(1); + expect(input.onLost).not.toHaveBeenCalled(); + const execution = startNativeSessionExecutionLeaseRenewal({ ...input, attempt: 2 }); + await preparation.stop(); + await vi.advanceTimersByTimeAsync(20_000); + expect(returning).toHaveBeenCalledTimes(2); + expect(input.onLost).not.toHaveBeenCalled(); + await execution.stop(); + }); it("does not queue SQL behind a slow renewal and fences a lost owner", async () => { vi.useFakeTimers(); const { input, update, returning } = fixture(); diff --git a/server/src/services/native-runtime/native-execution-lease.ts b/server/src/services/native-runtime/native-execution-lease.ts index 1460efa42e4..851c608574a 100644 --- a/server/src/services/native-runtime/native-execution-lease.ts +++ b/server/src/services/native-runtime/native-execution-lease.ts @@ -253,20 +253,37 @@ const renewals = new WeakMap; pending: Promise; busy: boolean; lost: Error | null; deadline: ReturnType; loss: Promise; - confirmedDeadline: number; + confirmedDeadline: number; adopt(confirmedDeadline: number): void; }>>(); +async function attemptAdvancedUnderOwner(binding: RenewalInput): Promise { + const [row] = await livenessDb(binding.db).transaction(async tx => { + await tx.execute(sql`select set_config('statement_timeout', ${`${RENEWAL_STATEMENT_TIMEOUT_MS}ms`}, true)`); + return tx.select({ attempt: nativeRunFinalizations.attempt, leaseOwner: nativeRunFinalizations.leaseOwner }) + .from(nativeRunFinalizations) + .where(and(eq(nativeRunFinalizations.runId, binding.runId), eq(nativeRunFinalizations.companyId, binding.companyId))) + .limit(1); + }); + return row?.leaseOwner === binding.leaseOwner && row.attempt > binding.attempt; +} + export function startNativeSessionExecutionLeaseRenewal(input: RenewalInput): NativeExecutionLeaseRenewal { let owners = renewals.get(input.db); if (!owners) { owners = new Map(); renewals.set(input.db, owners); } let state = owners.get(input.leaseOwner); if (state) { if (state.input.runId !== input.runId || state.input.companyId !== input.companyId - || state.input.issueId !== input.issueId || state.input.attempt !== input.attempt + || state.input.issueId !== input.issueId || state.input.attempt > input.attempt || state.input.controller.bootId !== input.controller.bootId || state.input.remote !== input.remote) { throw new Error("native_session_renewal_binding_changed"); } if (state.lost) throw state.lost; + if (input.attempt > state.input.attempt) { + state.input = { ...state.input, attempt: input.attempt, onLost: input.onLost, controller: input.controller, + controllerGeneration: input.controllerGeneration ?? state.input.controllerGeneration }; + state.adopt(input.confirmedDeadline + ?? performance.now() + (input.remote ? REMOTE_CONTROLLER_LEASE_MS : NATIVE_SESSION_EXECUTION_LEASE_TTL_MS)); + } state.references += 1; } else { const ttlMs = input.remote ? REMOTE_CONTROLLER_LEASE_MS : NATIVE_SESSION_EXECUTION_LEASE_TTL_MS; @@ -275,7 +292,7 @@ export function startNativeSessionExecutionLeaseRenewal(input: RenewalInput): Na state!.lost = error instanceof Error ? error : new Error("native_session_lease_lost"); clearInterval(state!.timer); clearTimeout(state!.deadline); - state!.loss = Promise.resolve().then(input.onLost).catch(() => undefined); + state!.loss = Promise.resolve().then(state!.input.onLost).catch(() => undefined); }; const deadline = (delay: number) => { const timer = setTimeout(() => lose(new Error("native_session_lease_renewal_deadline")), Math.max(0, delay)); @@ -285,18 +302,18 @@ export function startNativeSessionExecutionLeaseRenewal(input: RenewalInput): Na const renew = () => { if (state!.lost || state!.busy) return; state!.busy = true; - state!.pending = renewNativeSessionExecutionLeaseDeadline({ ...input, leaseTtlMs: ttlMs, - reacquire: !input.remote && input.controllerGeneration !== undefined - ? { controllerGeneration: input.controllerGeneration } : undefined }) + const binding = state!.input; + state!.pending = renewNativeSessionExecutionLeaseDeadline({ ...binding, leaseTtlMs: ttlMs, + reacquire: !binding.remote && binding.controllerGeneration !== undefined + ? { controllerGeneration: binding.controllerGeneration } : undefined }) .then(confirmedDeadline => { // A late database response cannot restore local authority after loss // or recreate a timer after the final caller has stopped. - if (state!.lost || state!.references === 0) return; - clearTimeout(state!.deadline); - state!.confirmedDeadline = confirmedDeadline; - state!.deadline = deadline(confirmedDeadline - performance.now()); + if (state!.references === 0) return; + state!.adopt(confirmedDeadline); }) - .catch((error: unknown) => { + .catch(async (error: unknown) => { + if (state!.input !== binding) return; // A cancelled statement can retry at the next tick, but does not // extend the ownership deadline. Never queue overlapping attempts. for (let cause = error, depth = 0; cause && depth < 4; cause = (cause as { cause?: unknown }).cause, depth++) { @@ -305,14 +322,22 @@ export function startNativeSessionExecutionLeaseRenewal(input: RenewalInput): Na const code = (cause as { code?: unknown }).code; if (code === "57014" || code === "40P01") return; } + const advanced = await attemptAdvancedUnderOwner(binding).catch(() => null); + if (advanced !== false || state!.input !== binding) return; lose(error); }).finally(() => { state!.busy = false; }); }; const timer = setInterval(renew, input.remote ? REMOTE_CONTROLLER_LEASE_MS / 3 : NATIVE_SESSION_EXECUTION_LEASE_RENEW_INTERVAL_MS); timer.unref?.(); const confirmedDeadline = input.confirmedDeadline ?? performance.now() + ttlMs; + const adopt = (confirmed: number) => { + if (state!.lost || confirmed <= state!.confirmedDeadline) return; + clearTimeout(state!.deadline); + state!.confirmedDeadline = confirmed; + state!.deadline = deadline(confirmed - performance.now()); + }; state = { input, references: 1, timer, pending: Promise.resolve(), busy: false, lost: null, - deadline: deadline(confirmedDeadline - performance.now()), confirmedDeadline, loss: Promise.resolve() }; + deadline: deadline(confirmedDeadline - performance.now()), confirmedDeadline, loss: Promise.resolve(), adopt }; owners.set(input.leaseOwner, state); } const owned = state;