diff --git a/src/routes/exports.ts b/src/routes/exports.ts index 8f04c8d..76ebddb 100644 --- a/src/routes/exports.ts +++ b/src/routes/exports.ts @@ -12,10 +12,56 @@ * /api/exports — Prediction history export (default alias) */ -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 { exportsPredictionsRouter } from "./exports/predictions"; +import { logger } from "../config/logger"; + +// In-flight request tracking for graceful shutdown drain +let inFlightExportsRequests = 0; + +/** + * Wait for all in-flight /api/exports requests to finish. + * @param timeoutMs Maximum time to wait before forcing resolution + */ +export async function drainExportsRequests(timeoutMs = 10000): Promise { + const start = Date.now(); + if (inFlightExportsRequests === 0) { + logger.info("No in-flight /api/exports requests to drain"); + return; + } + + logger.info({ inFlight: inFlightExportsRequests }, "Draining in-flight /api/exports requests..."); + + while (inFlightExportsRequests > 0) { + if (Date.now() - start > timeoutMs) { + logger.warn({ inFlight: inFlightExportsRequests }, "Timeout waiting for /api/exports requests to drain"); + break; + } + // Poll every 50ms + await new Promise((resolve) => setTimeout(resolve, 50)); + } + + if (inFlightExportsRequests === 0) { + logger.info("Successfully drained all /api/exports requests"); + } +} + +function exportsInFlightMiddleware(_req: Request, res: Response, next: NextFunction): void { + inFlightExportsRequests++; + let finished = false; + const cleanup = () => { + if (!finished) { + finished = true; + inFlightExportsRequests = Math.max(0, inFlightExportsRequests - 1); + } + }; + + res.once("finish", cleanup); + res.once("close", cleanup); + next(); +} export interface ExportsRouterOptions { rateLimit?: { @@ -27,6 +73,7 @@ export interface ExportsRouterOptions { export function createExportsRouter(options: ExportsRouterOptions = {}): Router { const router = Router(); + router.use(exportsInFlightMiddleware); router.use(requireAuth); router.use( createPerUserTokenBucketLimiter({ diff --git a/src/server.ts b/src/server.ts index 9d10f7c..9f37c33 100644 --- a/src/server.ts +++ b/src/server.ts @@ -10,6 +10,7 @@ import { backupVerificationWorker } from "./workers/backupVerificationWorker"; import { reconciliationWorker } from "./workers/reconciliationWorker"; import { startSlowQueryAlerter } from "./workers/slowQueryAlerter"; import { drainSearchRequests } from "./routes/search"; +import { drainExportsRequests } from "./routes/exports"; const app = createApp(); let webhookWorker: WebhookWorker | null = null; @@ -41,6 +42,8 @@ connectWithRetry() // Ensure in-flight /api/search requests finish await drainSearchRequests(4000); + // Ensure in-flight /api/exports requests finish + await drainExportsRequests(4000); stopScheduler(); await closeDb();