From 42bf8ca3a78f3a06c6e67d62c2a7891cbb68ea38 Mon Sep 17 00:00:00 2001 From: Bug-Hunter-X <148070949+Bug-Hunter-X@users.noreply.github.com> Date: Wed, 27 May 2026 02:00:44 +0100 Subject: [PATCH] task: drain jobs and proxy streams on graceful shutdown --- README.md | 3 +- src/index.test.ts | 99 ++++++++++++++++- src/index.ts | 141 ++++++++++++++++++++++-- src/webhooks/webhook.dispatcher.test.ts | 13 ++- src/webhooks/webhook.dispatcher.ts | 125 +++++++++++++-------- 5 files changed, 325 insertions(+), 56 deletions(-) diff --git a/README.md b/README.md index 036b2e0..78378a9 100644 --- a/README.md +++ b/README.md @@ -152,8 +152,9 @@ The app validates all environment variables at startup using [Zod](https://zod.d ## Production Shutdown Expectations - The server listens for `SIGTERM` and `SIGINT` and performs a graceful shutdown. -- On shutdown, it stops accepting new HTTP requests, waits for active connections to finish, and closes database resources. +- On shutdown, it stops accepting new HTTP requests, drains in-flight `/v1/call` proxy work, waits for active webhook deliveries to finish, and then closes database resources. - A 30 second timeout is enforced for in-flight connections; lingering sockets are destroyed to prevent hung termination. +- Background workers should stop scheduling new runs as soon as shutdown begins and finish any in-flight work inside the same drain window. - Shutdown hooks are registered with `process.once(...)` to avoid duplicate execution during restarts. - The dev workflow (`npm run dev` with `tsx watch`) is preserved. Restarts trigger the same graceful path instead of abrupt termination. diff --git a/src/index.test.ts b/src/index.test.ts index 1bb2b6c..b073907 100644 --- a/src/index.test.ts +++ b/src/index.test.ts @@ -1,7 +1,7 @@ /// import request from 'supertest'; import type { Server } from 'http'; -import app, { createGracefulShutdownHandler } from './index.js'; +import app, { createGracefulShutdownHandler, createInFlightDrainTracker } from './index.js'; jest.mock('./db/index.js', () => ({ db: {}, @@ -17,6 +17,10 @@ describe('Health API', () => { }); describe('graceful shutdown', () => { + afterEach(() => { + jest.useRealTimers(); + }); + it('closes server and database resources', async () => { const closeServer = jest.fn((callback: (err?: Error) => void) => callback()); const closeDatabase = jest.fn(async () => Promise.resolve()); @@ -35,6 +39,67 @@ describe('graceful shutdown', () => { expect(closeDatabase).toHaveBeenCalledTimes(1); }); + it('stops subsystems and waits for in-flight work before closing database resources', async () => { + let closeCallback: ((err?: Error) => void) | undefined; + let resolveDrain: (() => void) | undefined; + const closeServer = jest.fn((callback: (err?: Error) => void) => { + closeCallback = callback; + }); + const closeDatabase = jest.fn(async () => Promise.resolve()); + const beginShutdown = jest.fn(); + const awaitIdle = jest.fn( + () => + new Promise((resolve) => { + resolveDrain = resolve; + }), + ); + + const shutdown = createGracefulShutdownHandler({ + server: { close: closeServer } as unknown as Server, + activeConnections: new Set(), + closeDatabase, + timeoutMs: 50, + subsystems: [{ name: 'jobs', beginShutdown, awaitIdle }], + }); + + const promise = shutdown('SIGTERM'); + await Promise.resolve(); + expect(beginShutdown).toHaveBeenCalledTimes(1); + expect(awaitIdle).toHaveBeenCalledTimes(1); + expect(closeDatabase).not.toHaveBeenCalled(); + + closeCallback?.(); + expect(closeDatabase).not.toHaveBeenCalled(); + + resolveDrain?.(); + await expect(promise).resolves.toBe(0); + expect(closeDatabase).toHaveBeenCalledTimes(1); + }); + + it('destroys lingering sockets after the drain timeout', async () => { + jest.useFakeTimers(); + + const destroy = jest.fn(); + const socket = { destroy } as never; + const closeServer = jest.fn((_callback: (err?: Error) => void) => { + // Intentionally never closes to force timeout handling. + }); + const closeDatabase = jest.fn(async () => Promise.resolve()); + + const shutdown = createGracefulShutdownHandler({ + server: { close: closeServer } as unknown as Server, + activeConnections: new Set([socket]), + closeDatabase, + timeoutMs: 25, + }); + + void shutdown('SIGTERM'); + jest.advanceTimersByTime(25); + + expect(destroy).toHaveBeenCalledTimes(1); + expect(closeDatabase).not.toHaveBeenCalled(); + }); + it('reuses in-flight shutdown promise on repeated signals', async () => { let closeCallback: ((err?: Error) => void) | undefined; const closeServer = jest.fn((callback: (err?: Error) => void) => { @@ -60,3 +125,35 @@ describe('graceful shutdown', () => { expect(closeDatabase).toHaveBeenCalledTimes(1); }); }); + +describe('proxy drain tracker', () => { + it('waits for active proxy requests to finish before becoming idle', async () => { + const tracker = createInFlightDrainTracker('gateway-proxy'); + const next = jest.fn(); + const listeners = new Map void>(); + const res = { + setHeader: jest.fn(), + once: jest.fn((event: string, handler: () => void) => { + listeners.set(event, handler); + return res; + }), + } as any; + + tracker.middleware({} as any, res, next); + tracker.subsystem.beginShutdown(); + const idlePromise = tracker.subsystem.awaitIdle(); + + let settled = false; + void idlePromise.then(() => { + settled = true; + }); + await Promise.resolve(); + + expect(next).toHaveBeenCalledTimes(1); + expect(settled).toBe(false); + + listeners.get('finish')?.(); + await expect(idlePromise).resolves.toBeUndefined(); + expect(settled).toBe(true); + }); +}); diff --git a/src/index.ts b/src/index.ts index 9ffd7a8..b0a13c4 100644 --- a/src/index.ts +++ b/src/index.ts @@ -7,8 +7,10 @@ import { closeDbPool } from './config/health.js'; import { disconnectPrisma } from './lib/prisma.js'; import { errorHandler } from './middleware/errorHandler.js'; import { createGatewayIpAllowlist } from './middleware/ipAllowlist.js'; +import { awaitWebhookDispatcherIdle, stopWebhookDispatching } from './webhooks/webhook.dispatcher.js'; import type { Socket } from 'net'; import type { Server } from 'http'; +import type { RequestHandler } from 'express'; import { createDeveloperRouter } from './routes/developerRoutes.js'; import { createGatewayRouter } from './routes/gatewayRoutes.js'; @@ -30,6 +32,73 @@ interface GracefulShutdownOptions { closeDatabase: () => Promise; logger?: Pick; timeoutMs?: number; + subsystems?: DrainableSubsystem[]; +} + +export interface DrainableSubsystem { + name: string; + beginShutdown: () => void | Promise; + awaitIdle: () => Promise; +} + +export function createInFlightDrainTracker(name: string): { + middleware: RequestHandler; + subsystem: DrainableSubsystem; +} { + let active = 0; + let accepting = true; + const waiters = new Set<() => void>(); + + const notifyIfIdle = () => { + if (active === 0) { + for (const resolve of waiters) { + resolve(); + } + waiters.clear(); + } + }; + + const middleware: RequestHandler = (_req, res, next) => { + active += 1; + let settled = false; + + const finish = () => { + if (settled) { + return; + } + + settled = true; + active = Math.max(0, active - 1); + notifyIfIdle(); + }; + + if (!accepting) { + res.setHeader('Connection', 'close'); + } + + res.once('finish', finish); + res.once('close', finish); + next(); + }; + + return { + middleware, + subsystem: { + name, + beginShutdown() { + accepting = false; + }, + awaitIdle() { + if (active === 0) { + return Promise.resolve(); + } + + return new Promise((resolve) => { + waiters.add(resolve); + }); + }, + }, + }; } export function createGracefulShutdownHandler({ @@ -38,6 +107,7 @@ export function createGracefulShutdownHandler({ closeDatabase, logger = console, timeoutMs = 10_000, + subsystems = [], }: GracefulShutdownOptions) { let inFlight: Promise | null = null; @@ -55,18 +125,64 @@ export function createGracefulShutdownHandler({ } }, timeoutMs); - server.close(async (error?: Error) => { - clearTimeout(timeout); + const drainSubsystems = async (): Promise => { + for (const subsystem of subsystems) { + try { + await subsystem.beginShutdown(); + } catch (error) { + logger.error(`Error while stopping subsystem ${subsystem.name}`, error); + return false; + } + } - if (error) { - logger.error('Error while closing HTTP server', error); - resolve(1); - return; + const results = await Promise.race([ + Promise.allSettled( + subsystems.map(async (subsystem) => { + await subsystem.awaitIdle(); + }), + ), + new Promise<'timeout'>((timeoutResolve) => { + setTimeout(() => timeoutResolve('timeout'), timeoutMs); + }), + ]); + + if (results === 'timeout') { + logger.warn(`Timed out waiting for in-flight subsystem work after ${timeoutMs}ms`); + return true; + } + + let ok = true; + for (const [index, result] of results.entries()) { + if (result.status === 'rejected') { + ok = false; + logger.error( + `Error while draining subsystem ${subsystems[index]?.name ?? 'unknown'}`, + result.reason, + ); + } } + return ok; + }; + + const closeServer = new Promise((closeResolve) => { + server.close((error?: Error) => { + if (error) { + logger.error('Error while closing HTTP server', error); + closeResolve(false); + return; + } + + closeResolve(true); + }); + }); + + void Promise.all([closeServer, drainSubsystems()]).then(async ([serverClosed, drained]) => { + clearTimeout(timeout); + try { await closeDatabase(); - resolve(0); + resolve(serverClosed && drained ? 0 : 1); } catch (closeError) { logger.error('Error while closing data resources', closeError); resolve(1); @@ -154,6 +270,16 @@ if (isDirectExecution) { timeoutMs: config.proxy.timeoutMs, }, }); + const proxyDrainTracker = createInFlightDrainTracker('gateway-proxy'); + const shutdownSubsystems: DrainableSubsystem[] = [ + proxyDrainTracker.subsystem, + { + name: 'webhook-dispatcher', + beginShutdown: stopWebhookDispatching, + awaitIdle: awaitWebhookDispatcherIdle, + }, + ]; + app.use('/v1/call', proxyDrainTracker.middleware); app.use('/v1/call', proxyRouter); @@ -194,6 +320,7 @@ if (isDirectExecution) { server, activeConnections, closeDatabase: closeAllDataResources, + subsystems: shutdownSubsystems, }); const onSignal = (signal: NodeJS.Signals) => { diff --git a/src/webhooks/webhook.dispatcher.test.ts b/src/webhooks/webhook.dispatcher.test.ts index b935b9a..d2b8495 100644 --- a/src/webhooks/webhook.dispatcher.test.ts +++ b/src/webhooks/webhook.dispatcher.test.ts @@ -1,4 +1,4 @@ -import { dispatchWebhook } from './webhook.dispatcher.js'; +import { dispatchWebhook, resetWebhookDispatcherForTests, stopWebhookDispatching } from './webhook.dispatcher.js'; import type { WebhookConfig, WebhookPayload } from './webhook.types.js'; describe('Webhook Dispatcher', () => { @@ -6,6 +6,7 @@ describe('Webhook Dispatcher', () => { beforeEach(() => { originalFetch = global.fetch; + resetWebhookDispatcherForTests(); jest.useFakeTimers(); jest.spyOn(console, 'warn').mockImplementation(() => {}); jest.spyOn(console, 'log').mockImplementation(() => {}); @@ -117,4 +118,14 @@ describe('Webhook Dispatcher', () => { expect(fetchMock).toHaveBeenCalledTimes(5); }); + + it('does not start new deliveries after shutdown begins', async () => { + const fetchMock = jest.fn(); + global.fetch = fetchMock as any; + + stopWebhookDispatching(); + await dispatchWebhook(config, payload); + + expect(fetchMock).not.toHaveBeenCalled(); + }); }); diff --git a/src/webhooks/webhook.dispatcher.ts b/src/webhooks/webhook.dispatcher.ts index a8946c7..2c49833 100644 --- a/src/webhooks/webhook.dispatcher.ts +++ b/src/webhooks/webhook.dispatcher.ts @@ -5,6 +5,8 @@ import { logger } from '../logger.js'; const MAX_RETRIES = 5; const BASE_DELAY_MS = 1000; +let acceptingDispatches = true; +const inFlightDispatches = new Set>(); function sleep(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); @@ -14,6 +16,30 @@ function signPayload(secret: string, body: string): string { return crypto.createHmac('sha256', secret).update(body).digest('hex'); } +function trackDispatch(operation: Promise): Promise { + const tracked = operation.finally(() => { + inFlightDispatches.delete(tracked as Promise); + }); + + inFlightDispatches.add(tracked as Promise); + return tracked; +} + +export function stopWebhookDispatching(): void { + acceptingDispatches = false; +} + +export async function awaitWebhookDispatcherIdle(): Promise { + while (inFlightDispatches.size > 0) { + await Promise.allSettled([...inFlightDispatches]); + } +} + +export function resetWebhookDispatcherForTests(): void { + acceptingDispatches = true; + inFlightDispatches.clear(); +} + /** * Dispatches a webhook payload to the registered URL. * @@ -27,61 +53,68 @@ export async function dispatchWebhook( config: WebhookConfig, payload: WebhookPayload ): Promise { - const body = JSON.stringify(payload); - const deliveryId = crypto.randomUUID(); - const headers: Record = { - 'Content-Type': 'application/json', - 'User-Agent': 'Callora-Webhook/1.0', - 'X-Callora-Event': payload.event, - 'X-Callora-Timestamp': payload.timestamp, - 'X-Callora-Delivery': deliveryId, - }; - - if (config.secret) { - headers['X-Callora-Signature'] = `sha256=${signPayload(config.secret, body)}`; + if (!acceptingDispatches) { + logger.warn(`[webhook] Skipping ${payload.event} dispatch during shutdown for ${config.url}`); + return; } - let lastError: unknown; + return trackDispatch((async () => { + const body = JSON.stringify(payload); + const deliveryId = crypto.randomUUID(); + const headers: Record = { + 'Content-Type': 'application/json', + 'User-Agent': 'Callora-Webhook/1.0', + 'X-Callora-Event': payload.event, + 'X-Callora-Timestamp': payload.timestamp, + 'X-Callora-Delivery': deliveryId, + }; - for (let attempt = 0; attempt < MAX_RETRIES; attempt++) { - try { - const response = await fetch(config.url, { - method: 'POST', - body, - headers, - signal: AbortSignal.timeout(10_000), // 10s timeout per attempt - }); + if (config.secret) { + headers['X-Callora-Signature'] = `sha256=${signPayload(config.secret, body)}`; + } + + let lastError: unknown; + + for (let attempt = 0; attempt < MAX_RETRIES; attempt++) { + try { + const response = await fetch(config.url, { + method: 'POST', + body, + headers, + signal: AbortSignal.timeout(10_000), // 10s timeout per attempt + }); - if (response.ok) { - console.log( - `[webhook] ✓ Delivered ${payload.event} to ${config.url} (attempt ${attempt + 1})` + if (response.ok) { + console.log( + `[webhook] ✓ Delivered ${payload.event} to ${config.url} (attempt ${attempt + 1})` + ); + return; + } + + lastError = new Error(`HTTP ${response.status} ${response.statusText}`); + console.warn( + `[webhook] Non-2xx response (${response.status}) for ${config.url}, attempt ${attempt + 1}` + ); + } catch (err) { + lastError = err; + console.warn( + `[webhook] Error delivering to ${config.url}, attempt ${attempt + 1}:`, + (err as Error).message ); - return; // success — stop retrying } - lastError = new Error(`HTTP ${response.status} ${response.statusText}`); - console.warn( - `[webhook] Non-2xx response (${response.status}) for ${config.url}, attempt ${attempt + 1}` - ); - } catch (err) { - lastError = err; - console.warn( - `[webhook] Error delivering to ${config.url}, attempt ${attempt + 1}:`, - (err as Error).message - ); - } - - if (attempt < MAX_RETRIES - 1) { - const delay = BASE_DELAY_MS * Math.pow(2, attempt); // exponential backoff - console.log(`[webhook] Retrying in ${delay}ms...`); - await sleep(delay); + if (attempt < MAX_RETRIES - 1) { + const delay = BASE_DELAY_MS * Math.pow(2, attempt); + console.log(`[webhook] Retrying in ${delay}ms...`); + await sleep(delay); + } } - } - logger.error( - `[webhook] ✗ Failed to deliver ${payload.event} to ${config.url} after ${MAX_RETRIES} attempts.`, - lastError - ); + logger.error( + `[webhook] ✗ Failed to deliver ${payload.event} to ${config.url} after ${MAX_RETRIES} attempts.`, + lastError + ); + })()); } export async function dispatchToAll(