From c442996ef05b49a756a9a0155a1db3101e75ac68 Mon Sep 17 00:00:00 2001
From: Kazama <agentgangstar123@gmail.com>
Date: Thu, 30 Jul 2026 11:13:53 +0530
Subject: [PATCH] feat(reports): drain requests during shutdown

---
 src/routes/reports.ts      | 33 ++++++++++++++++++++++++++++++++-
 src/server.ts              |  6 ++++++
 tests/reportsDrain.test.ts | 26 ++++++++++++++++++++++++++
 3 files changed, 64 insertions(+), 1 deletion(-)
 create mode 100644 tests/reportsDrain.test.ts

diff --git a/src/routes/reports.ts b/src/routes/reports.ts
index becd2fa5..8a6054a2 100644
--- a/src/routes/reports.ts
+++ b/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({
diff --git a/src/server.ts b/src/server.ts
index af5f8bee..4e15b6ec 100644
--- a/src/server.ts
+++ b/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();
diff --git a/tests/reportsDrain.test.ts b/tests/reportsDrain.test.ts
new file mode 100644
index 00000000..59807cd0
--- /dev/null
+++ b/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();
+  });
+});