diff --git a/apps/hub/README.md b/apps/hub/README.md index 02e7151b9..a1fd820fd 100644 --- a/apps/hub/README.md +++ b/apps/hub/README.md @@ -27,6 +27,12 @@ native tenant middleware. - `credential-expiry-sweep.ts`, `cron-due.ts`, `routine-launcher.ts`, `routine-scheduler.ts`, `tenant-create-guard.ts` are host-level background jobs and guards wired at boot, alongside the route mounts. +- `in-flight-requests.ts` and `shutdown.ts` bound SIGINT/SIGTERM: the hub + waits for Hono handlers that have not yet returned a Response (a + request still in a Postgres transaction, a git write), then + `server.stop(true)` so a live SSE bridge or sidecar websocket cannot + hang the drain. The sequence is capped at 10s; a lingering stream is + not a shutdown fault. ## Running diff --git a/apps/hub/src/in-flight-requests.test.ts b/apps/hub/src/in-flight-requests.test.ts new file mode 100644 index 000000000..528279da3 --- /dev/null +++ b/apps/hub/src/in-flight-requests.test.ts @@ -0,0 +1,59 @@ +import { expect, test } from "bun:test"; +import { Hono } from "hono"; + +import { createInFlightRequestTracker } from "./in-flight-requests"; + +test("pending starts at zero and whenIdle resolves immediately", async () => { + const tracker = createInFlightRequestTracker(); + expect(tracker.pending).toBe(0); + await tracker.whenIdle(); + expect(tracker.pending).toBe(0); +}); + +test("middleware counts a request until the handler returns", async () => { + const tracker = createInFlightRequestTracker(); + const app = new Hono(); + app.use(tracker.middleware); + let pendingDuringHandler: number | undefined; + app.get("/work", async () => { + pendingDuringHandler = tracker.pending; + return new Response("ok"); + }); + + const response = await app.request("/work"); + expect(response.status).toBe(200); + expect(pendingDuringHandler).toBe(1); + expect(tracker.pending).toBe(0); +}); + +test("whenIdle waits for an in-flight handler, including one that throws", async () => { + const tracker = createInFlightRequestTracker(); + const app = new Hono(); + app.use(tracker.middleware); + app.onError(() => new Response("error", { status: 500 })); + + let release!: () => void; + const held = new Promise((resolve) => { + release = resolve; + }); + app.get("/held", async () => { + await held; + throw new Error("handler fault"); + }); + + const request = app.request("/held"); + const idle = tracker.whenIdle(); + let idleSettled = false; + void idle.then(() => { + idleSettled = true; + }); + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(tracker.pending).toBe(1); + expect(idleSettled).toBe(false); + + release(); + await request; + await idle; + expect(idleSettled).toBe(true); + expect(tracker.pending).toBe(0); +}); diff --git a/apps/hub/src/in-flight-requests.ts b/apps/hub/src/in-flight-requests.ts new file mode 100644 index 000000000..bbf0fbf67 --- /dev/null +++ b/apps/hub/src/in-flight-requests.ts @@ -0,0 +1,57 @@ +import { Hono, type MiddlewareHandler } from "hono"; + +export type InFlightRequestTracker = { + readonly middleware: MiddlewareHandler; + readonly pending: number; + whenIdle: () => Promise; +}; + +/** + * Counts Hono handlers that have not yet returned a Response. Streaming + * bodies (SSE) and upgraded websockets stay open after that return, so they + * do not keep `pending` above zero — those are lingering connections the + * drain force-closes, not in-flight work. + */ +export function createInFlightRequestTracker(): InFlightRequestTracker { + let pending = 0; + const waiters = new Set<() => void>(); + + const notifyIfIdle = () => { + if (pending !== 0) return; + for (const waiter of waiters) waiter(); + waiters.clear(); + }; + + const middleware: MiddlewareHandler = async (_c, next) => { + pending += 1; + try { + await next(); + } finally { + pending -= 1; + notifyIfIdle(); + } + }; + + return { + middleware, + get pending() { + return pending; + }, + whenIdle() { + if (pending === 0) return Promise.resolve(); + return new Promise((resolve) => { + waiters.add(resolve); + }); + }, + }; +} + +export function withInFlightRequestTracking( + app: Hono, + tracker: InFlightRequestTracker, +): Hono { + const outer = new Hono(); + outer.use(tracker.middleware); + outer.route("/", app); + return outer; +} diff --git a/apps/hub/src/index.ts b/apps/hub/src/index.ts index bf17db671..6b0d5c2b3 100644 --- a/apps/hub/src/index.ts +++ b/apps/hub/src/index.ts @@ -367,7 +367,11 @@ import { createToolGrantsForPins } from "./tool-grants"; import { createMcpCredentialBindingsFor } from "./mcp-credential-bindings"; import { reconcilePinnedToolPackagesAfterConnect } from "./connection-live-reconcile"; import { createPinnedPackageCredentialBindingsFor } from "./pinned-package-credential-bindings"; -import { shutdownHub } from "./shutdown"; +import { drainHubServer, shutdownHub } from "./shutdown"; +import { + createInFlightRequestTracker, + withInFlightRequestTracking, +} from "./in-flight-requests"; // Host policy constants, not configuration. const MAX_TARBALL_BYTES = 10 * 1024 * 1024; @@ -3492,6 +3496,8 @@ export async function createHub(config: HubConfig) { guardDeps.operatorTenantId = config.operatorTenantId; } const guardedApp = guardedHubApp(app, guardDeps); + const inFlight = createInFlightRequestTracker(); + const servingApp = withInFlightRequestTracking(guardedApp, inFlight); // Env-key auto-plant (CL-6101): runs in-process against the app this // function is about to return, so it needs nothing more than that @@ -3502,11 +3508,12 @@ export async function createHub(config: HubConfig) { envProviderKeys: config.envProviderKeys, envProviderBaseUrls: config.envProviderBaseUrls, admin: config.envCredentialPlantAdmin, - fetch: (request) => Promise.resolve(guardedApp.fetch(request)), + fetch: (request) => Promise.resolve(servingApp.fetch(request)), }); return { - app: guardedApp, + app: servingApp, + whenRequestsIdle: () => inFlight.whenIdle(), db, close: async () => { sidecarAllocationReconciliationStopped = true; @@ -3552,14 +3559,20 @@ if (import.meta.main) { const log = getLogger(["hub"]); log.info`Hub serving on port ${port}`; const SHUTDOWN_DRAIN_MS = 10_000; - // `server.stop()` waits for open connections and websockets by default, - // so it sits inside the same bound as the hub's own closes. + // In-flight Hono handlers (a request mid-Postgres-transaction, a git + // write, anything that has not returned a Response yet) must finish + // before connections are torn down. `server.stop()` with no argument + // also waits for SSE bridges and idle sidecar websockets, which never + // close on their own — so once handlers are idle, force-close what's + // left. A live stream must not turn this drain into a timeout fault. const shutdown = () => shutdownHub({ - drain: async () => { - await server.stop(); - await hub.close(); - }, + drain: () => + drainHubServer({ + whenRequestsIdle: hub.whenRequestsIdle, + stop: (force) => server.stop(force), + close: hub.close, + }), timeoutMs: SHUTDOWN_DRAIN_MS, exit: (code) => process.exit(code), }); diff --git a/apps/hub/src/shutdown.test.ts b/apps/hub/src/shutdown.test.ts index fb70a27c1..4f6640e32 100644 --- a/apps/hub/src/shutdown.test.ts +++ b/apps/hub/src/shutdown.test.ts @@ -1,5 +1,5 @@ import { expect, test } from "bun:test"; -import { drainWithTimeout, shutdownHub } from "./shutdown"; +import { drainHubServer, drainWithTimeout, shutdownHub } from "./shutdown"; test("drainWithTimeout resolves drained when the drain completes inside the bound", async () => { const outcome = await drainWithTimeout(() => Promise.resolve(), 1_000); @@ -76,3 +76,19 @@ test("shutdownHub exits non-zero and reports the cause when the drain rejects", expect(exitCode).toBe(1); expect(reported).toEqual([error]); }); + +test("drainHubServer waits for idle handlers then force-stops", async () => { + const order: string[] = []; + await drainHubServer({ + whenRequestsIdle: async () => { + order.push("idle"); + }, + stop: (force) => { + order.push(force ? "stop-force" : "stop"); + }, + close: async () => { + order.push("close"); + }, + }); + expect(order).toEqual(["idle", "stop-force", "close"]); +}); diff --git a/apps/hub/src/shutdown.ts b/apps/hub/src/shutdown.ts index b97be00fa..6fa08dee0 100644 --- a/apps/hub/src/shutdown.ts +++ b/apps/hub/src/shutdown.ts @@ -39,6 +39,27 @@ export async function drainWithTimeout( return outcome; } +export type DrainHubServerArgs = { + whenRequestsIdle: () => Promise; + stop: (force: boolean) => void | Promise; + close: () => void | Promise; +}; + +/** + * Wait for every in-flight Hono handler to return, then force-stop the + * listener so lingering SSE/websocket connections cannot hang `server.stop()`, + * then close hub resources. + */ +export async function drainHubServer({ + whenRequestsIdle, + stop, + close, +}: DrainHubServerArgs): Promise { + await whenRequestsIdle(); + await stop(true); + await close(); +} + export type ShutdownHubDeps = { drain: () => Promise; timeoutMs: number; diff --git a/apps/hub/test/shutdown-drain.test.ts b/apps/hub/test/shutdown-drain.test.ts new file mode 100644 index 000000000..79d9159bd --- /dev/null +++ b/apps/hub/test/shutdown-drain.test.ts @@ -0,0 +1,168 @@ +// A real Bun.serve() with a live websocket or SSE stream reproduces the +// drain hang that unit tests on drainWithTimeout (src/shutdown.test.ts) +// cannot see: `server.stop()` with no argument waits for those connections +// to close on their own, and they never do. The hub drain waits for the +// Hono in-flight request counter instead, then force-stops so lingering +// streams cannot turn a deploy into a timeout fault. +import { afterEach, describe, expect, test } from "bun:test"; +import { Hono } from "hono"; +import { upgradeWebSocket, websocket } from "hono/bun"; +import { streamSSE } from "hono/streaming"; + +import { + createInFlightRequestTracker, + withInFlightRequestTracking, +} from "../src/in-flight-requests"; +import { drainHubServer, drainWithTimeout, shutdownHub } from "../src/shutdown"; + +type Serving = { + port: number; + stop: (force?: boolean) => Promise; +}; + +const cleanups: (() => void)[] = []; + +afterEach(() => { + for (const cleanup of cleanups.splice(0)) cleanup(); +}); + +function serve(app: Hono): Serving { + const server = Bun.serve({ + port: 0, + idleTimeout: 0, + fetch: app.fetch, + websocket, + }); + const port = server.port; + if (port === undefined) { + throw new Error("Bun.serve did not bind a port"); + } + cleanups.push(() => { + void server.stop(true); + }); + return { + port, + stop: (force = false) => server.stop(force), + }; +} + +function startTracked(): { + inner: Hono; + tracker: ReturnType; + serve: () => Serving; +} { + const tracker = createInFlightRequestTracker(); + const inner = new Hono(); + return { + inner, + tracker, + serve: () => serve(withInFlightRequestTracking(inner, tracker)), + }; +} + +async function openWebSocket(port: number, path = "/ws"): Promise { + const client = new WebSocket(`ws://localhost:${String(port)}${path}`); + cleanups.push(() => client.close()); + await new Promise((resolve, reject) => { + client.addEventListener("open", () => resolve(), { once: true }); + client.addEventListener("error", reject, { once: true }); + }); + return client; +} + +describe("hub shutdown drain against a real server", () => { + test("server.stop() with no argument never drains while a websocket stays open", async () => { + const app = new Hono(); + app.get( + "/ws", + upgradeWebSocket(() => ({ + onMessage() { + // The socket just stays open. + }, + })), + ); + const server = serve(app); + await openWebSocket(server.port); + + const outcome = await drainWithTimeout(() => server.stop(), 200); + expect(outcome).toEqual({ kind: "timed-out" }); + }); + + test("drainHubServer force-stops after in-flight handlers finish, even with a live websocket", async () => { + const { inner, tracker, serve: start } = startTracked(); + let handlerFinished = false; + inner.get("/slow", async () => { + await new Promise((resolve) => setTimeout(resolve, 80)); + handlerFinished = true; + return new Response("ok"); + }); + inner.get( + "/ws", + upgradeWebSocket(() => ({ + onMessage() { + // The socket just stays open. + }, + })), + ); + const server = start(); + await openWebSocket(server.port); + + const inFlight = fetch(`http://localhost:${String(server.port)}/slow`); + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(tracker.pending).toBe(1); + + const outcome = await drainWithTimeout( + () => + drainHubServer({ + whenRequestsIdle: () => tracker.whenIdle(), + stop: (force) => server.stop(force), + close: () => Promise.resolve(), + }), + 1_000, + ); + expect(outcome).toEqual({ kind: "drained" }); + expect(handlerFinished).toBe(true); + expect(tracker.pending).toBe(0); + + await inFlight.catch(() => undefined); + }); + + test("a live SSE stream does not hold the in-flight count or fail drain", async () => { + const { inner, tracker, serve: start } = startTracked(); + inner.get("/stream", (c) => + streamSSE(c, async (stream) => { + await stream.writeSSE({ data: "hello" }); + await new Promise(() => undefined); + }), + ); + const server = start(); + const response = await fetch( + `http://localhost:${String(server.port)}/stream`, + ); + expect(response.ok).toBe(true); + expect(tracker.pending).toBe(0); + + let exitCode: number | undefined; + const reported: unknown[] = []; + await shutdownHub({ + drain: () => + drainHubServer({ + whenRequestsIdle: () => tracker.whenIdle(), + stop: (force) => server.stop(force), + close: () => Promise.resolve(), + }), + timeoutMs: 1_000, + exit: (code) => { + exitCode = code; + }, + report: (error) => { + reported.push(error); + return "unused-ref-id"; + }, + }); + expect(exitCode).toBe(0); + expect(reported).toEqual([]); + + await response.body?.cancel().catch(() => undefined); + }); +});