diff --git a/server/src/__tests__/server-startup-feedback-export.test.ts b/server/src/__tests__/server-startup-feedback-export.test.ts index 6c367a02f5b..944eb1dcc49 100644 --- a/server/src/__tests__/server-startup-feedback-export.test.ts +++ b/server/src/__tests__/server-startup-feedback-export.test.ts @@ -12,6 +12,8 @@ import { getStartupRecoveryState } from "../startup-recovery-state.js"; import { idleWorkSnapshot } from "../services/task-admission.js"; import { readIdleSleepSafety } from "../services/idle-sleep-safety.js"; import * as idleLocalWork from "../services/idle-local-work.js"; +import * as workspaceRecovery from "../services/workspace-recovery.js"; +import * as workspaceRecoveryWork from "../services/workspace-recovery-work.js"; import { idleBackupWakeMarker, markIdleBackupWakeRequired } from "../services/idle-database-backup.js"; import { DELIVERY_QUEUES } from "../services/delivery-work-notifications.js"; import { decisionService, decisionRetentionService, toolAccessService, githubConnectionEventService, reconcileBuiltInAgentsOnStartup, reconcileCodexLocalManagedHomesOnStartup, reconcilePersistedRuntimeServicesOnStartup } from "../services/index.js"; @@ -1705,6 +1707,71 @@ describe("startServer feedback export wiring", () => { } finally { release(); await starting; timer.mockRestore(); } }); + it.each([false, true])("joins timed-out workspace repairs before teardown with hot-restart skipDrain=%s", async skipDrain => { + loadConfigMock.mockReturnValue(buildTestConfig({ heartbeatSchedulerEnabled: true })); + let releaseFirst!: () => void; + let releaseSecond!: () => void; + const first = new Promise(resolve => { releaseFirst = resolve; }); + const second = new Promise(resolve => { releaseSecond = resolve; }); + const scan = vi.spyOn(workspaceRecovery, "reconcileRetainedWorkspaces").mockResolvedValue([first, second]); + const register = vi.spyOn(workspaceRecoveryWork, "registerWorkspaceRecoveryWork"); + const drain = vi.fn(async () => undefined); + const heartbeat = { ...heartbeatServiceMock, + drainRunningRunsForShutdown: drain, + prepareHotRestartShutdown: vi.fn(async () => ({ skipDrain })), + }; + heartbeatServiceFactoryMock.mockReturnValueOnce(heartbeat); + const interval = vi.spyOn(globalThis, "setInterval").mockImplementation(() => 1 as unknown as ReturnType); + let stopping: Promise | undefined; + let stop: ReturnType | undefined; + try { + const server = await startServer(); + await settleStartupHousekeeping(); + expect(scan).toHaveBeenCalledOnce(); + stop = vi.spyOn(register.mock.results[0]!.value, "stop"); + const app = await createAppMock.mock.results.at(-1)!.value as { + locals: { paperclipShutdown?: () => Promise; paperclipDatabaseClosing?: boolean }; + }; + const appShutdown = vi.fn(async () => undefined); + app.locals.paperclipShutdown = appShutdown; + vi.useFakeTimers({ toFake: ["Date", "setTimeout", "clearTimeout"] }); + let stopped = false; + stopping = server.shutdown().then(() => { stopped = true; }); + await settleStartupHousekeeping(); + expect(stop).toHaveBeenCalledOnce(); + await vi.advanceTimersByTimeAsync(20_000); + await settleStartupHousekeeping(); + expect(logger.info).toHaveBeenCalledWith( + expect.objectContaining({ step: "workspace_recovery_stop" }), + "shutdown step timed out; continuing shutdown without it", + ); + expect(drain).toHaveBeenCalledTimes(skipDrain ? 0 : 1); + expect(heartbeatServiceMock.drainActiveRunExecutions).toHaveBeenCalledTimes(skipDrain ? 0 : 1); + expect(appShutdown).not.toHaveBeenCalled(); + expect(app.locals.paperclipDatabaseClosing).toBeUndefined(); + expect(stopped).toBe(false); + releaseFirst(); + await settleStartupHousekeeping(); + expect(appShutdown).not.toHaveBeenCalled(); + expect(stopped).toBe(false); + releaseSecond(); + await stopping; + expect(appShutdown).toHaveBeenCalledOnce(); + expect(app.locals.paperclipDatabaseClosing).toBe(true); + expect(stopped).toBe(true); + expect(stop).toHaveBeenCalledOnce(); + } finally { + releaseFirst(); + releaseSecond(); + await stopping; + vi.useRealTimers(); + stop?.mockRestore(); + interval.mockRestore(); + register.mockRestore(); + scan.mockRestore(); + } + }); + it.each(["retained cleanup", "workspace finalization", "environment settlement", "native run projection", "native result commit", "native recovery notices", "external conversation settlement", "native response presentation"] as const)("becomes ready while %s is held but joins it at shutdown", async owner => { loadConfigMock.mockReturnValue(buildTestConfig({ heartbeatSchedulerEnabled: true })); let releaseCleanup!: () => void; diff --git a/server/src/index.ts b/server/src/index.ts index 17466c5168c..74818db2dd3 100644 --- a/server/src/index.ts +++ b/server/src/index.ts @@ -200,9 +200,11 @@ import { createDecisionRetentionNotifyOriginAgent, createDecisionWakeOriginAgent import { closeHttpListenerForShutdown, coordinateHeartbeatSchedulerShutdown, + createPreDrainShutdownBudget, drainRunExecutionFinalizersForShutdown, finalizeServerShutdown, loadWithoutCoordinatedShutdownSignalHooks, + runBoundedShutdownStep, } from "./shutdown.js"; import { initializeCloudRuntimeIdentity } from "./services/cloud-runtime-identity.js"; import { initializeAdapterSettings } from "./services/adapter-settings.js"; @@ -2019,6 +2021,8 @@ async function startServerWithDatabaseTeardown( signal: "SIGINT" | "SIGTERM", exitProcess: boolean, ) => { + const shutdownStartedAt = Date.now(); + logger.info({ signal }, "shutdown signal received; starting ordered shutdown"); setStartupRecoveryPhase("draining"); for (const client of liveEventsWebSocketServer.clients) client.close(1012, "Server restarting"); liveEventsWebSocketServer.close(); @@ -2026,33 +2030,51 @@ async function startServerWithDatabaseTeardown( // are still available. Railway must switch routing before signaling A. await closeHttpListenerForShutdown({ server, signal, timeoutMs: 30_000, log: logger, drainRequests: abort => drainHttpRequestWork(app, abort) }); - await systemdNotify(["--stopping", `--status=Stopping after ${signal}`]); + const preDrain = createPreDrainShutdownBudget({ signal, log: logger }); + await preDrain.run("systemd_notify", () => + systemdNotify(["--stopping", `--status=Stopping after ${signal}`])); heartbeatSchedulerStopped = true; skillVersionRetentionWorker?.unsubscribe(); - await taskDrainExpiryWork.stop(); + await preDrain.run("task_drain_expiry_stop", () => taskDrainExpiryWork.stop()); nativeRunnerEvidenceRetentionWorker?.unsubscribe(); - await skillSourceSync.stop(); + await preDrain.run("skill_source_sync_stop", () => skillSourceSync.stop()); unsubscribeChatCompletions(); unsubscribeNativeResponsePresentation?.(); unsubscribeNativeWorkspaceFinalization?.(); unsubscribeRetainedNativeSessionCleanup?.(); // B can accept Stop/answers while A still owns an execution. Preserve the // owner command workers and their transport until the handoff below ends. - await deliveryWork.stop(deploymentHints ? new Set([ + await preDrain.run("delivery_work_stop", () => deliveryWork.stop(deploymentHints ? new Set([ DELIVERY_QUEUES.nativeCancellation, DELIVERY_QUEUES.nativeRuntimeResponses, DELIVERY_QUEUES.question, DELIVERY_QUEUES.deploymentNotifications, DELIVERY_QUEUES.liveInvalidations, - ]) : undefined); - await runEvidenceMigrationWork?.stop(); - await providerTraceRetentionWork.stop(); - await skillVersionRetentionWorker?.dispose(); - await workspaceRecoveryWork.stop(); + ]) : undefined)); + await preDrain.run("run_evidence_migration_stop", async () => { + await runEvidenceMigrationWork?.stop(); + }); + await preDrain.run("provider_trace_retention_stop", () => providerTraceRetentionWork.stop()); + await preDrain.run("skill_version_retention_dispose", async () => { + await skillVersionRetentionWorker?.dispose(); + }); + const workspaceRecoveryStop = workspaceRecoveryWork.stop(); + await preDrain.run("workspace_recovery_stop", () => workspaceRecoveryStop); unregisterTerminalWorkspaceProbe?.(); unsubscribeNativeWorkspaceSyncRetention?.(); + const prepareHotRestart = prepareHotRestartShutdown; const heartbeatShutdown = await coordinateHeartbeatSchedulerShutdown({ signal, - prepareHotRestartShutdown, - waitForHeartbeatSchedulerIdle, + prepareHotRestartShutdown: prepareHotRestart + ? async (preparationSignal) => { + const preparation = await preDrain.run("hot_restart_prepare", () => + prepareHotRestart(preparationSignal)); + if (preparation.outcome === "completed") return preparation.value; + if (preparation.outcome === "failed") throw preparation.error; + throw new Error("hot-restart shutdown preparation timed out"); + } + : null, + waitForHeartbeatSchedulerIdle: async () => { + await preDrain.run("scheduler_idle", waitForHeartbeatSchedulerIdle); + }, }); const skipHeartbeatDrain = heartbeatShutdown.hotRestart?.skipDrain === true; const selectiveDrainRunIds = heartbeatShutdown.hotRestart?.drainRunIds ?? null; @@ -2068,16 +2090,18 @@ async function startServerWithDatabaseTeardown( ); } - const telemetryClient = getTelemetryClient(); - if (telemetryClient) { - telemetryClient.stop(); - await telemetryClient.flush(); - } - if (!skipHeartbeatDrain && drainHeartbeatRunsForShutdown) { + const drainStartedAt = Date.now(); + logger.info( + { signal, sinceSignalMs: drainStartedAt - shutdownStartedAt }, + "starting graceful heartbeat run drain", + ); try { const drain = await drainHeartbeatRunsForShutdown(signal, selectiveDrainRunIds); - logger.info({ signal, drain }, "graceful heartbeat run drain complete"); + logger.info( + { signal, drain, durationMs: Date.now() - drainStartedAt }, + "graceful heartbeat run drain complete", + ); } catch (err) { logger.error({ err, signal }, "graceful heartbeat run drain failed"); } @@ -2092,6 +2116,20 @@ async function startServerWithDatabaseTeardown( }); } + await workspaceRecoveryStop; + + const telemetryClient = getTelemetryClient(); + if (telemetryClient) { + telemetryClient.stop(); + await runBoundedShutdownStep({ + step: "telemetry_flush", + signal, + timeoutMs: 5_000, + run: () => telemetryClient.flush(), + log: logger, + }); + } + await cancellationDelivery?.stop(); await runtimeResponseDelivery?.stop(); await deploymentHints?.stop(); diff --git a/server/src/shutdown.test.ts b/server/src/shutdown.test.ts index 754902cc3ba..83176e35210 100644 --- a/server/src/shutdown.test.ts +++ b/server/src/shutdown.test.ts @@ -3,10 +3,12 @@ import { describe, expect, it, vi } from "vitest"; import { closeHttpListenerForShutdown, coordinateHeartbeatSchedulerShutdown, + createPreDrainShutdownBudget, drainRunExecutionFinalizersForShutdown, finalizeServerShutdown, finishServiceShutdown, loadWithoutCoordinatedShutdownSignalHooks, + runBoundedShutdownStep, } from "./shutdown.js"; function deferred() { @@ -357,6 +359,131 @@ describe("drainRunExecutionFinalizersForShutdown", () => { }); }); +describe("runBoundedShutdownStep", () => { + it("returns the step value and logs its duration", async () => { + const log = stubLogger(); + await expect(runBoundedShutdownStep({ + step: "delivery_work_stop", + signal: "SIGTERM", + timeoutMs: 1_000, + run: async () => "stopped", + log, + })).resolves.toMatchObject({ outcome: "completed", value: "stopped" }); + expect(log.info).toHaveBeenCalledWith( + expect.objectContaining({ step: "delivery_work_stop", signal: "SIGTERM", durationMs: expect.any(Number) }), + "shutdown step completed", + ); + }); + + it("stops waiting for a step that outlives its bound", async () => { + vi.useFakeTimers(); + try { + const log = stubLogger(); + const pending = runBoundedShutdownStep({ + step: "scheduler_idle", + signal: "SIGTERM", + timeoutMs: 250, + run: () => new Promise(() => undefined), + log, + }); + await vi.advanceTimersByTimeAsync(250); + await expect(pending).resolves.toMatchObject({ outcome: "timed_out" }); + expect(log.info).toHaveBeenCalledWith( + expect.objectContaining({ step: "scheduler_idle", timeoutMs: 250 }), + expect.stringContaining("timed out"), + ); + } finally { + vi.useRealTimers(); + } + }); + + it("logs a step that rejects and does not reject the shutdown", async () => { + const log = stubLogger(); + const error = new Error("skill sync stop failed"); + await expect(runBoundedShutdownStep({ + step: "skill_source_sync_stop", + signal: "SIGINT", + timeoutMs: 1_000, + run: async () => { throw error; }, + log, + })).resolves.toMatchObject({ outcome: "failed", error }); + expect(log.error).toHaveBeenCalledWith( + expect.objectContaining({ step: "skill_source_sync_stop", err: error }), + expect.stringContaining("failed"), + ); + }); + + it("contains a step that throws synchronously", async () => { + const result = await runBoundedShutdownStep({ + step: "systemd_notify", + signal: "SIGTERM", + timeoutMs: 1_000, + run: () => { throw new Error("spawn failed"); }, + log: stubLogger(), + }); + expect(result.outcome).toBe("failed"); + }); +}); + +describe("createPreDrainShutdownBudget", () => { + it("gives each step only the time that is left in the shared budget", async () => { + vi.useFakeTimers(); + try { + const log = stubLogger(); + const budget = createPreDrainShutdownBudget({ signal: "SIGTERM", budgetMs: 1_000, log }); + const first = budget.run("skill_source_sync_stop", () => + new Promise((resolve) => setTimeout(resolve, 700))); + await vi.advanceTimersByTimeAsync(700); + await expect(first).resolves.toMatchObject({ outcome: "completed" }); + + const second = budget.run("scheduler_idle", () => new Promise(() => undefined)); + await vi.advanceTimersByTimeAsync(300); + await expect(second).resolves.toMatchObject({ outcome: "timed_out" }); + expect(log.info).toHaveBeenCalledWith( + expect.objectContaining({ step: "scheduler_idle", timeoutMs: 300 }), + expect.stringContaining("timed out"), + ); + } finally { + vi.useRealTimers(); + } + }); + + it("reaches the run drain within the budget when every pre-drain step hangs", async () => { + vi.useFakeTimers(); + try { + const budget = createPreDrainShutdownBudget({ signal: "SIGTERM", budgetMs: 1_000, log: stubLogger() }); + const hang = () => new Promise(() => undefined); + const drain = vi.fn(async () => ({ interrupted: 1 })); + const shutdown = (async () => { + await budget.run("skill_source_sync_stop", hang); + await budget.run("delivery_work_stop", hang); + await coordinateHeartbeatSchedulerShutdown({ + signal: "SIGTERM", + prepareHotRestartShutdown: async () => { + const preparation = await budget.run("hot_restart_prepare", () => + new Promise<{ skipDrain: boolean }>(() => undefined)); + if (preparation.outcome === "completed") return preparation.value; + throw new Error("hot-restart shutdown preparation timed out"); + }, + waitForHeartbeatSchedulerIdle: async () => { + await budget.run("scheduler_idle", hang); + }, + }); + return drain(); + })(); + + const startedAt = Date.now(); + await vi.advanceTimersByTimeAsync(1_000); + await vi.runAllTimersAsync(); + await expect(shutdown).resolves.toEqual({ interrupted: 1 }); + expect(drain).toHaveBeenCalledOnce(); + expect(Date.now() - startedAt).toBeLessThan(1_010); + } finally { + vi.useRealTimers(); + } + }); +}); + describe("loadWithoutCoordinatedShutdownSignalHooks", () => { it("removes the eager signal handlers from the real embedded-postgres import", async () => { const before = { diff --git a/server/src/shutdown.ts b/server/src/shutdown.ts index 45721a5d4a3..ebc9e590b8f 100644 --- a/server/src/shutdown.ts +++ b/server/src/shutdown.ts @@ -7,6 +7,77 @@ type ShutdownLogger = { error(obj: object, msg: string): void; }; +export const PRE_DRAIN_SHUTDOWN_BUDGET_MS = 20_000; + +export type ShutdownStepResult = + | { outcome: "completed"; value: T; durationMs: number } + | { outcome: "timed_out"; durationMs: number } + | { outcome: "failed"; error: unknown; durationMs: number }; + +export async function runBoundedShutdownStep(input: { + step: string; + signal: "SIGINT" | "SIGTERM"; + timeoutMs: number; + run: () => Promise; + log: ShutdownLogger; + now?: () => number; +}): Promise> { + const now = input.now ?? Date.now; + const startedAt = now(); + const timeoutMs = Math.max(0, input.timeoutMs); + let timer: NodeJS.Timeout | null = null; + try { + const result = await Promise.race([ + Promise.resolve() + .then(input.run) + .then( + (value) => ({ outcome: "completed" as const, value }), + (error: unknown) => ({ outcome: "failed" as const, error }), + ), + new Promise<{ outcome: "timed_out" }>((resolve) => { + timer = setTimeout(() => resolve({ outcome: "timed_out" }), timeoutMs); + timer.unref?.(); + }), + ]); + const durationMs = now() - startedAt; + const fields = { signal: input.signal, step: input.step, durationMs, timeoutMs }; + if (result.outcome === "completed") { + input.log.info(fields, "shutdown step completed"); + return { outcome: "completed", value: result.value, durationMs }; + } + if (result.outcome === "failed") { + input.log.error({ ...fields, err: result.error }, "shutdown step failed; continuing shutdown"); + return { outcome: "failed", error: result.error, durationMs }; + } + input.log.info(fields, "shutdown step timed out; continuing shutdown without it"); + return { outcome: "timed_out", durationMs }; + } finally { + if (timer) clearTimeout(timer); + } +} + +export function createPreDrainShutdownBudget(input: { + signal: "SIGINT" | "SIGTERM"; + budgetMs?: number; + log: ShutdownLogger; + now?: () => number; +}) { + const now = input.now ?? Date.now; + const deadline = now() + (input.budgetMs ?? PRE_DRAIN_SHUTDOWN_BUDGET_MS); + return { + run(step: string, run: () => Promise) { + return runBoundedShutdownStep({ + step, + signal: input.signal, + timeoutMs: deadline - now(), + run, + log: input.log, + now, + }); + }, + }; +} + export async function drainRunExecutionFinalizersForShutdown(input: { signal: "SIGINT" | "SIGTERM"; drain: (() => Promise) | null;