task: drain jobs and proxy streams on graceful shutdown

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.
 

src/index.test.ts

@@ -1,7 +1,7 @@
 /// <reference types="jest" />
 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<void>((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<string, () => 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);
+  });
+});

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<void>;
   logger?: Pick<typeof console, 'log' | 'warn' | 'error'>;
   timeoutMs?: number;
+  subsystems?: DrainableSubsystem[];
+}
+
+export interface DrainableSubsystem {
+  name: string;
+  beginShutdown: () => void | Promise<void>;
+  awaitIdle: () => Promise<void>;
+}
+
+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<void>((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<number> | null = null;
 
@@ -55,18 +125,64 @@ export function createGracefulShutdownHandler({
         }
       }, timeoutMs);
 
-      server.close(async (error?: Error) => {
-        clearTimeout(timeout);
+      const drainSubsystems = async (): Promise<boolean> => {
+        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<boolean>((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) => {

src/webhooks/webhook.dispatcher.test.ts

@@ -1,11 +1,12 @@
-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', () => {
     let originalFetch: typeof global.fetch;
 
     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();
+    });
 });