diff --git a/apps/server/src/graceful-shutdown.integration.test.ts b/apps/server/src/graceful-shutdown.integration.test.ts index 620a1bb65..6f5c9183b 100644 --- a/apps/server/src/graceful-shutdown.integration.test.ts +++ b/apps/server/src/graceful-shutdown.integration.test.ts @@ -1,6 +1,11 @@ import { describe, expect, it } from "bun:test"; +import { REPO_ROOT } from "./test-helpers/workspace"; -const REPO_ROOT = "/Users/brain/Coding/snipeship/ccflare"; +// REPO_ROOT was previously hardcoded to the original author's local machine path +// ("/Users/brain/Coding/snipeship/ccflare"), which does not exist on any other checkout and made +// this test non-portable (Bun.spawn's cwd pointing at a nonexistent directory fails with a +// misleading ENOENT attributed to the spawned binary, not the cwd). Reusing the existing portable +// helper (computed via import.meta.dir, not hardcoded) fixes this for any checkout location. describe("graceful shutdown integration", () => { it("awaits programmatic stop until pending request writes are flushed", async () => { diff --git a/packages/proxy/src/usage-worker.test.ts b/packages/proxy/src/usage-worker.test.ts index c748489d6..5c3a07b86 100644 --- a/packages/proxy/src/usage-worker.test.ts +++ b/packages/proxy/src/usage-worker.test.ts @@ -287,6 +287,181 @@ describe("UsageWorkerController", () => { expect(controller.getHealthSnapshot().state).toBe("stopped"); }); + it("flushes a message queued just before shutdown once the worker becomes ready during the wait, instead of dropping it", async () => { + const workers: FakeWorker[] = []; + const logger = new TestLogger(); + const controller = new UsageWorkerController({ + createWorker() { + const worker = new FakeWorker(); + workers.push(worker); + return worker; + }, + readyTimeoutMs: 200, + ackTimeoutMs: 10_000, + shutdownDelayMs: 0, + logger, + }); + controllers.push(controller); + + // Worker is not ready yet, so this queues instead of sending immediately -- + // this is the exact scenario that used to lose the message on shutdown. + controller.postMessage(createStartMessage()); + expect(workers[0]?.postedMessages).toHaveLength(0); + + const shutdownPromise = controller.terminateGracefully(); + + // Ready arrives DURING the shutdown wait window (before readyTimeoutMs elapses). + await Bun.sleep(5); + workers[0]?.emitMessage({ type: "ready" }); + + await waitFor( + () => + workers[0]?.postedMessages.some( + (message) => (message as { type?: string }).type === "start", + ) ?? false, + ); + + workers[0]?.emitMessage({ + type: "shutdown-complete", + asyncWriter: { healthy: true, failureCount: 0, queuedJobs: 0 }, + }); + + await shutdownPromise; + expect( + logger.warnings.some((warning) => warning.includes("Dropping")), + ).toBe(false); + }); + + it("resolves shutdown within readyTimeoutMs + shutdownDelayMs (never hangs) and drops queued messages with a warning when the worker never becomes ready", async () => { + const workers: FakeWorker[] = []; + const logger = new TestLogger(); + const readyTimeoutMs = 30; + const shutdownDelayMs = 20; + const controller = new UsageWorkerController({ + createWorker() { + const worker = new FakeWorker(); + workers.push(worker); + return worker; + }, + readyTimeoutMs, + ackTimeoutMs: 10_000, + shutdownDelayMs, + logger, + }); + controllers.push(controller); + + // Worker never emits "ready" in this test. + controller.postMessage(createStartMessage()); + expect(workers[0]?.postedMessages).toHaveLength(0); + + const startedAt = Date.now(); + await expect(controller.terminateGracefully()).rejects.toThrow(); + const elapsedMs = Date.now() - startedAt; + + // Generous slack for test-runner scheduling jitter -- the point is this is + // BOUNDED (readyTimeoutMs + shutdownDelayMs), not that it hangs indefinitely. + expect(elapsedMs).toBeLessThan(readyTimeoutMs + shutdownDelayMs + 1_000); + expect( + logger.warnings.some((warning) => + warning.includes("Dropping 1 queued usage worker message"), + ), + ).toBe(true); + expect(workers[0]?.terminateCalls).toBe(1); + }); + + it("force termination flushes a queued message only if the worker is already ready, without waiting", async () => { + const workers: FakeWorker[] = []; + const logger = new TestLogger(); + const controller = new UsageWorkerController({ + createWorker() { + const worker = new FakeWorker(); + workers.push(worker); + return worker; + }, + readyTimeoutMs: 10_000, + ackTimeoutMs: 10_000, + shutdownDelayMs: 0, + logger, + }); + controllers.push(controller); + + workers[0]?.emitMessage({ type: "ready" }); + controller.postMessage(createStartMessage()); + expect(workers[0]?.postedMessages).toHaveLength(1); + + controller.forceTerminate(); + + expect(workers[0]?.terminateCalls).toBe(1); + expect( + logger.warnings.some((warning) => warning.includes("Dropping")), + ).toBe(false); + }); + + it("promptly rejects an in-flight terminateGracefully() wait when forceTerminate() races it, instead of idling out the full readyTimeoutMs", async () => { + const workers: FakeWorker[] = []; + const logger = new TestLogger(); + const readyTimeoutMs = 300; + const controller = new UsageWorkerController({ + createWorker() { + const worker = new FakeWorker(); + workers.push(worker); + return worker; + }, + readyTimeoutMs, + ackTimeoutMs: 10_000, + shutdownDelayMs: 0, + logger, + }); + controllers.push(controller); + + // Worker never becomes ready, so terminateGracefully() below enters its + // wait-for-ready phase and suspends for up to readyTimeoutMs. + controller.postMessage(createStartMessage()); + const gracefulShutdown = controller.terminateGracefully(); + + // Race it with a forceTerminate() call (mirrors getUsageWorker()'s self-heal + // path in proxy.ts, which calls forceTerminate() on an instance for which + // isShuttingDown() is already true). + await Bun.sleep(5); + const startedAt = Date.now(); + controller.forceTerminate(); + + await expect(gracefulShutdown).rejects.toThrow( + "Usage worker was force terminated", + ); + const elapsedMs = Date.now() - startedAt; + + // Must settle promptly (driven by forceTerminate unblocking the wait), not by + // idling out the full readyTimeoutMs. + expect(elapsedMs).toBeLessThan(readyTimeoutMs / 2); + }); + + it("does not let a standalone forceTerminate() poison a later, unrelated terminateGracefully() call on the same instance", async () => { + const workers: FakeWorker[] = []; + const logger = new TestLogger(); + const controller = new UsageWorkerController({ + createWorker() { + const worker = new FakeWorker(); + workers.push(worker); + return worker; + }, + readyTimeoutMs: 10_000, + ackTimeoutMs: 10_000, + shutdownDelayMs: 0, + logger, + }); + controllers.push(controller); + + // Force-terminate with nothing racing it -- this used to leave a stale + // earlyTerminationError sitting on the instance. + controller.forceTerminate(); + + // A later, wholly independent terminateGracefully() call must resolve + // normally (there is no worker left and nothing queued), not reject with the + // stale "Usage worker was force terminated" error from the earlier call. + await expect(controller.terminateGracefully()).resolves.toBeUndefined(); + }); + it("rejects graceful shutdown when the worker never confirms completion", async () => { const workers: FakeWorker[] = []; const controller = new UsageWorkerController({ diff --git a/packages/proxy/src/usage-worker.ts b/packages/proxy/src/usage-worker.ts index 22f6839e0..c5b922065 100644 --- a/packages/proxy/src/usage-worker.ts +++ b/packages/proxy/src/usage-worker.ts @@ -159,6 +159,24 @@ export class UsageWorkerController implements UsageWorkerTransport { private resolveShutdown: (() => void) | null = null; private rejectShutdown: ((error: Error) => void) | null = null; private lastError: string | null = null; + // Guards terminateGracefully() re-entrancy for the FULL async flow (including the + // wait-for-ready phase below, which happens before shutdownPromise exists), so a + // second concurrent call always reuses the first call's in-flight promise instead + // of starting a duplicate termination sequence. + private terminatePromise: Promise | null = null; + // Resolved (with whether the worker became ready) by handleWorkerMessage's ready + // branch, or forced to `false` if the worker crashes/errors while shutting down, so + // terminateGracefully()'s bounded wait below never has to rely solely on its own + // timeout to make progress. + private readyWaiters: Array<(isReady: boolean) => void> = []; + // Set by finishShutdown() when it is called WITH an error while a + // terminateGracefully() call is still suspended in its wait-for-ready phase (e.g. a + // concurrent forceTerminate(), or a worker crash) -- at that point shutdownPromise + // does not exist yet, so resolveShutdown/rejectShutdown are still null and the + // error would otherwise be silently swallowed. runGracefulTermination() checks this + // after its wait and re-throws it, so the caller's promise correctly rejects + // instead of resolving as if shutdown had completed normally. + private earlyTerminationError: Error | null = null; constructor(options: UsageWorkerControllerOptions = {}) { this.createWorkerImpl = options.createWorker ?? createDefaultWorker; @@ -213,20 +231,59 @@ export class UsageWorkerController implements UsageWorkerTransport { } terminateGracefully(): Promise { - if (this.shutdownPromise) { - return this.shutdownPromise; + if (this.terminatePromise) { + return this.terminatePromise; } + this.terminatePromise = this.runGracefulTermination(); + return this.terminatePromise; + } + + // Previously this cleared queuedMessages synchronously and unconditionally, which + // silently dropped any usage/analytics message that was queued (worker not yet + // ready) at the moment shutdown was triggered -- losing that request's log entry + // forever. Now, if there is something queued and the worker isn't ready yet, we + // give it up to readyTimeoutMs to become ready and flush before giving up (and + // only THEN do we drop + warn, instead of dropping silently up front). + private async runGracefulTermination(): Promise { this.shuttingDown = true; - this.ready = false; this.clearReadyTimer(); + // Discard any stale error left over from an earlier, unrelated termination + // cycle on this same controller instance (e.g. a standalone forceTerminate() + // call that fired with nothing racing it) -- only an error produced DURING + // this call's own wait below should ever be observed at the check further down. + this.earlyTerminationError = null; + + if (this.worker && this.queuedMessages.length > 0 && !this.ready) { + const becameReady = await this.waitForReady(this.readyTimeoutMs); + if (becameReady) { + this.flushQueuedMessages(); + } + } + + if (this.earlyTerminationError) { + const error = this.earlyTerminationError; + this.earlyTerminationError = null; + // finishShutdown() already reset shuttingDown/worker/etc for whichever path + // (forceTerminate / crash) produced this error -- just propagate it so the + // original terminateGracefully() caller sees a rejection instead of a + // silent success. + throw error; + } + + this.ready = false; this.clearPendingAcks(); - this.queuedMessages.length = 0; + if (this.queuedMessages.length > 0) { + this.logger.warn( + `Dropping ${this.queuedMessages.length} queued usage worker message(s) that could not be flushed before shutdown`, + ); + this.queuedMessages.length = 0; + } const worker = this.worker; if (!worker) { this.shuttingDown = false; - return Promise.resolve(); + return; } this.shutdownPromise = new Promise((resolve, reject) => { @@ -269,10 +326,28 @@ export class UsageWorkerController implements UsageWorkerTransport { forceTerminate(): void { this.shuttingDown = true; - this.ready = false; this.clearReadyTimer(); + // Unblock any terminateGracefully() call currently suspended in its + // wait-for-ready phase (see runGracefulTermination/waitForReady below) -- + // otherwise that caller's promise would not settle until readyTimeoutMs + // elapses even though this worker is being torn down right now. Mirrors the + // same unblock done in the onerror/onmessageerror shutdown-crash branches. + this.resolveReadyWaiters(false); + + // Best-effort, zero-wait flush: forceTerminate() must stay non-blocking (no + // await), so only flush if the worker already happens to be ready right now -- + // never wait for readiness the way terminateGracefully() does. + if (this.ready) { + this.flushQueuedMessages(); + } + this.ready = false; this.clearPendingAcks(); - this.queuedMessages.length = 0; + if (this.queuedMessages.length > 0) { + this.logger.warn( + `Dropping ${this.queuedMessages.length} queued usage worker message(s) on force termination`, + ); + this.queuedMessages.length = 0; + } if (this.shutdownTimer) { clearTimeout(this.shutdownTimer); this.shutdownTimer = null; @@ -282,6 +357,47 @@ export class UsageWorkerController implements UsageWorkerTransport { this.finishShutdown(worker, new Error("Usage worker was force terminated")); } + // Resolves to true once a "ready" message arrives, or false if readiness is not + // reached within timeoutMs, or if the worker crashes/errors while shutting down. + // Always resolves (never rejects) and never hangs past timeoutMs, so this can't + // make terminateGracefully() hang indefinitely. + private waitForReady(timeoutMs: number): Promise { + if (this.ready) { + return Promise.resolve(true); + } + + return new Promise((resolve) => { + let settled = false; + let timer: ReturnType; + const settle = (isReady: boolean) => { + if (settled) { + return; + } + settled = true; + clearTimeout(timer); + const index = this.readyWaiters.indexOf(settle); + if (index !== -1) { + this.readyWaiters.splice(index, 1); + } + resolve(isReady); + }; + + this.readyWaiters.push(settle); + timer = setTimeout(() => settle(false), timeoutMs); + unrefTimer(timer); + }); + } + + private resolveReadyWaiters(isReady: boolean): void { + if (this.readyWaiters.length === 0) { + return; + } + const waiters = this.readyWaiters.splice(0, this.readyWaiters.length); + for (const settle of waiters) { + settle(isReady); + } + } + private startWorker(): void { if (this.shuttingDown) { return; @@ -310,6 +426,10 @@ export class UsageWorkerController implements UsageWorkerTransport { this.lastError = message; this.logger.error(`Usage worker crashed: ${message}`); if (this.shuttingDown) { + // Unblock any pending terminateGracefully() wait-for-ready immediately + // instead of leaving it to idle out its own timeout -- the worker is gone, + // it will never become ready. + this.resolveReadyWaiters(false); this.finishShutdown( worker, new Error(`Usage worker crashed during shutdown: ${message}`), @@ -327,6 +447,7 @@ export class UsageWorkerController implements UsageWorkerTransport { this.lastError = "Usage worker emitted an invalid message payload"; this.logger.error("Usage worker emitted an invalid message payload"); if (this.shuttingDown) { + this.resolveReadyWaiters(false); this.finishShutdown( worker, new Error("Usage worker emitted an invalid message during shutdown"), @@ -395,15 +516,26 @@ export class UsageWorkerController implements UsageWorkerTransport { return; } - if (this.shuttingDown) { - return; - } - if (isReadyMessage(message)) { this.ready = true; this.lastError = null; this.clearReadyTimer(); - this.flushQueuedMessages(); + // Unblock any terminateGracefully() wait-for-ready call unconditionally -- + // this must run even while shuttingDown, otherwise a ready-signal that + // arrives during the shutdown window would be silently ignored and the + // bounded wait below would have to idle out its full timeout instead of + // resolving immediately. + this.resolveReadyWaiters(true); + // The explicit wait-then-flush path in runGracefulTermination() already + // handles flushing for the shutdown case; only auto-flush here when we are + // NOT shutting down, to avoid a double-flush race with that path. + if (!this.shuttingDown) { + this.flushQueuedMessages(); + } + return; + } + + if (this.shuttingDown) { return; } @@ -547,10 +679,22 @@ export class UsageWorkerController implements UsageWorkerTransport { this.resolveShutdown = null; this.rejectShutdown = null; this.shutdownPromise = null; + // Reset so a controller that gets restarted/reused after a full shutdown + // (e.g. crash-during-shutdown, see onerror above) is not permanently stuck + // returning an already-settled promise from a future terminateGracefully() call. + this.terminatePromise = null; if (error) { this.lastError = error.message; - rejectShutdown?.(error); + if (rejectShutdown) { + rejectShutdown(error); + } else { + // No in-flight shutdownPromise to reject (finishShutdown fired while a + // terminateGracefully() call was still awaiting waitForReady, before + // shutdownPromise existed) -- stash it for runGracefulTermination to + // re-throw once its wait resolves. + this.earlyTerminationError = error; + } return; }