feat(reports): drain requests during shutdown

src/routes/reports.ts

@@ -11,12 +11,42 @@
  * /api/reports/scheduled   — CRUD for user-owned scheduled report configs
  */
 
-import { Router } from "express";
+import { Router, type Request, type Response, type NextFunction } from "express";
 import { requireAuth } from "../middleware/requireAuth";
 import { createPerUserTokenBucketLimiter } from "../middleware/rateLimit";
 import { scheduledReportsRouter } from "./reports/scheduled";
 import { idempotency } from "../middleware/idempotency";
 
+let inFlightReportsRequests = 0;
+
+/** Wait for report handlers to finish before the database is closed. */
+export async function drainReportsRequests(timeoutMs = 10000): Promise<void> {
+  const start = Date.now();
+  while (inFlightReportsRequests > 0 && Date.now() - start <= timeoutMs) {
+    await new Promise((resolve) => setTimeout(resolve, 50));
+  }
+}
+
+/** Track all requests entering /api/reports, including auth failures. */
+export function reportsInFlightMiddleware(
+  _req: Request,
+  res: Response,
+  next: NextFunction,
+): void {
+  inFlightReportsRequests += 1;
+  let finished = false;
+  const cleanup = () => {
+    if (!finished) {
+      finished = true;
+      inFlightReportsRequests = Math.max(0, inFlightReportsRequests - 1);
+    }
+  };
+
+  res.once("finish", cleanup);
+  res.once("close", cleanup);
+  next();
+}
+
 export interface ReportsRouterOptions {
   rateLimit?: {
     capacity?: number;
@@ -27,6 +57,7 @@ export interface ReportsRouterOptions {
 export function createReportsRouter(options: ReportsRouterOptions = {}): Router {
   const router = Router();
 
+  router.use(reportsInFlightMiddleware);
   router.use(requireAuth);
   router.use(
     createPerUserTokenBucketLimiter({

src/server.ts

@@ -13,6 +13,7 @@ import { startPredictionsConfirmer } from "./workers/predictionsConfirmer";
 import { drainSearchRequests } from "./routes/search";
 import { drainExportsRequests } from "./routes/exports";
 import { drainFingerprintRequests } from "./routes/fingerprint";
+import { drainReportsRequests } from "./routes/reports";
 
 const app = createApp();
 let webhookWorker: WebhookWorker | null = null;
@@ -44,12 +45,17 @@ connectWithRetry()
         process.exit(1);
       }, 5000).unref();
 
+      // Stop accepting new connections while existing route handlers drain.
+      server.close();
+
       // Ensure in-flight /api/search requests finish
       await drainSearchRequests(4000);
       // Ensure in-flight /api/exports requests finish
       await drainExportsRequests(4000);
       // Ensure in-flight /api/fingerprint requests finish
       await drainFingerprintRequests(4000);
+      // Ensure in-flight /api/reports requests finish before closing Postgres
+      await drainReportsRequests(4000);
 
       stopScheduler();
       await closeDb();

tests/reportsDrain.test.ts

@@ -0,0 +1,26 @@
+import { EventEmitter } from "node:events";
+import type { Request, Response } from "express";
+import {
+  drainReportsRequests,
+  reportsInFlightMiddleware,
+} from "../src/routes/reports";
+
+describe("/api/reports graceful shutdown drain", () => {
+  it("waits for an in-flight response to finish", async () => {
+    const response = new EventEmitter() as unknown as Response;
+    const next = jest.fn();
+
+    reportsInFlightMiddleware({} as Request, response, next);
+    const draining = drainReportsRequests(500);
+
+    await new Promise((resolve) => setImmediate(resolve));
+    response.emit("finish");
+
+    await expect(draining).resolves.toBeUndefined();
+    expect(next).toHaveBeenCalledTimes(1);
+  });
+
+  it("returns immediately when no requests are in flight", async () => {
+    await expect(drainReportsRequests(50)).resolves.toBeUndefined();
+  });
+});