Fix graceful shutdown to drain requests before disconnecting Prisma

backend/src/__tests__/gracefulShutdown.test.ts

@@ -0,0 +1,103 @@
+import http from 'http';
+import { GracefulShutdownHandler } from '../gracefulShutdown';
+
+function listen(server: http.Server): Promise<number> {
+  return new Promise((resolve) => {
+    server.listen(0, () => {
+      const address = server.address();
+      resolve(typeof address === 'object' && address ? address.port : 0);
+    });
+  });
+}
+
+describe('GracefulShutdownHandler', () => {
+  let exitSpy: jest.SpyInstance;
+
+  beforeEach(() => {
+    exitSpy = jest.spyOn(process, 'exit').mockImplementation(((_code?: number) => undefined) as never);
+  });
+
+  afterEach(() => {
+    exitSpy.mockRestore();
+    process.removeAllListeners('SIGTERM');
+    process.removeAllListeners('SIGINT');
+  });
+
+  it('does not run cleanup (e.g. prisma.$disconnect) until in-flight requests have finished committing', async () => {
+    // Reproduces the fixed bug: a request holding an open
+    // prisma.vault.update transaction must be allowed to finish before the
+    // DB connection is torn down, or the transaction is killed mid-commit
+    // and the row stays locked (P2034 on the next deploy's migration).
+    const events: string[] = [];
+    let transactionCommitted = false;
+
+    const server = http.createServer((_req, res) => {
+      setTimeout(() => {
+        transactionCommitted = true;
+        events.push('transaction-committed');
+        res.writeHead(200);
+        res.end('ok');
+      }, 100);
+    });
+
+    const port = await listen(server);
+    const handler = new GracefulShutdownHandler(2000);
+    handler.register(server);
+
+    handler.onShutdown(async () => {
+      events.push('prisma-disconnect');
+      if (!transactionCommitted) {
+        throw new Error('prisma disconnected before the in-flight transaction committed');
+      }
+    });
+
+    expect(handler.isShuttingDown()).toBe(false);
+
+    // agent: false disables keep-alive so the underlying socket closes as
+    // soon as the response ends, instead of idling and only being reaped by
+    // the (much longer, 2000ms) drain-timeout force-destroy below.
+    const requestFinished = new Promise<void>((resolve, reject) => {
+      http
+        .get(`http://127.0.0.1:${port}`, { agent: false }, (res) => {
+          res.on('data', () => {});
+          res.on('end', resolve);
+        })
+        .on('error', reject);
+    });
+
+    // Let the server start handling the request before signalling shutdown.
+    await new Promise((r) => setTimeout(r, 20));
+    process.emit('SIGTERM');
+
+    // The shutdown gate flips synchronously, before the drain/cleanup work.
+    expect(handler.isShuttingDown()).toBe(true);
+
+    await requestFinished;
+
+    // Give the post-drain cleanup step a chance to run.
+    await new Promise((r) => setTimeout(r, 100));
+
+    expect(events).toEqual(['transaction-committed', 'prisma-disconnect']);
+    expect(exitSpy).toHaveBeenCalledWith(0);
+  });
+
+  it('bounds cleanup to its own timeout and still exits if a task hangs', async () => {
+    const server = http.createServer((_req, res) => res.end('ok'));
+    await listen(server);
+
+    // drainTimeout=2000 (no active connections, so drain finishes
+    // immediately); cleanupTimeout=50 so a hung cleanup task doesn't leave
+    // this test waiting out the real 10s production default.
+    const handler = new GracefulShutdownHandler(2000, 50);
+    handler.register(server);
+
+    handler.onShutdown(() => new Promise(() => {})); // never resolves
+
+    process.emit('SIGTERM');
+
+    await new Promise((r) => setTimeout(r, 150));
+
+    expect(handler.isShuttingDown()).toBe(true);
+    expect(exitSpy).toHaveBeenCalledWith(0);
+  });
+});

backend/src/__tests__/shutdownGate.test.ts

@@ -0,0 +1,32 @@
+import request from 'supertest';
+import app, { shutdownHandler } from '../index';
+
+describe('shutdown gate middleware', () => {
+  afterEach(() => {
+    // Reach past the public API to reset internal state between tests —
+    // isShuttingDown() has no public setter by design (it should only ever
+    // flip via a real SIGTERM/SIGINT).
+    (shutdownHandler as unknown as { shuttingDown: boolean }).shuttingDown = false;
+  });
+
+  it('serves normal routes and returns 200 on /health before shutdown begins', async () => {
+    expect(shutdownHandler.isShuttingDown()).toBe(false);
+
+    const health = await request(app).get('/health');
+    expect(health.status).toBe(200);
+    const summary = await request(app).get('/api/v1/vault/summary');
+    expect(summary.status).not.toBe(503);
+  });
+
+  it('rejects new requests with 503 + Retry-After once shutdown begins, but keeps serving /health', async () => {
+    (shutdownHandler as unknown as { shuttingDown: boolean }).shuttingDown = true;
+
+    const summary = await request(app).get('/api/v1/vault/summary');
+    expect(summary.status).toBe(503);
+    expect(summary.headers['retry-after']).toBe('10');
+    expect(summary.body.code).toBe('SHUTTING_DOWN');
+
+    const health = await request(app).get('/health');
+    expect(health.status).toBe(200);
+  });
+});

backend/src/gracefulShutdown.ts

@@ -2,14 +2,30 @@ import { Server } from 'http';
 import type { Socket } from 'net';
 import { logger } from './middleware/structuredLogging';
 
+// Default bound on the DB/job cleanup step so a hung disconnect can't wedge
+// the process indefinitely once connections have already drained.
+const DEFAULT_CLEANUP_TIMEOUT_MS = 10000;
+
 export class GracefulShutdownHandler {
   private drainTimeout: number;
+  private cleanupTimeout: number;
   private server: Server | null = null;
   private activeConnections = new Set<Socket>();
   private cleanupTasks: (() => Promise<void>)[] = [];
+  private shuttingDown = false;
 
-  constructor(drainTimeoutMs: number = 30000) {
+  constructor(drainTimeoutMs: number = 30000, cleanupTimeoutMs: number = DEFAULT_CLEANUP_TIMEOUT_MS) {
     this.drainTimeout = drainTimeoutMs;
+    this.cleanupTimeout = cleanupTimeoutMs;
+  }
+
+  /**
+   * True from the moment a shutdown signal is received. Request-handling
+   * code (e.g. the shutdown-gate middleware) uses this to start rejecting
+   * new work with a 503 before the listener actually stops accepting sockets.
+   */
+  isShuttingDown(): boolean {
+    return this.shuttingDown;
   }
 
   register(server: Server): void {
@@ -39,11 +55,20 @@ export class GracefulShutdownHandler {
   private async shutdown(signal: string): Promise<void> {
     logger.log('info', `${signal} received, starting graceful shutdown`);
 
+    // Flip the gate immediately: request-handling middleware checks this to
+    // start returning 503s before the listener has actually stopped accepting.
+    this.shuttingDown = true;
+
     if (!this.server) {
+      await this.runCleanupTasks();
       process.exit(0);
+      return;
     }
 
-    // Force close after drain timeout
+    // Force-destroy any remaining sockets if the drain takes too long, so a
+    // slow/stuck client can't block shutdown forever. This only forces the
+    // HTTP drain; cleanup still runs afterwards so DB connections close
+    // cleanly instead of being torn down mid-transaction.
     const drainTimer = setTimeout(() => {
       logger.log(
         'warn',
@@ -53,31 +78,61 @@ export class GracefulShutdownHandler {
       this.activeConnections.forEach((socket) => {
         socket.destroy();
       });
-
-      process.exit(1);
     }, this.drainTimeout);
 
     drainTimer.unref();
 
     try {
-      // Run registered cleanup tasks (e.g., database shutdown)
-      await Promise.all(this.cleanupTasks.map(task => task().catch(err => {
-        logger.log('error', 'Cleanup task failed during shutdown', {
-          error: err instanceof Error ? err.message : String(err)
+      // Stop accepting new connections and wait for in-flight requests
+      // (and any transaction they hold open) to finish before touching
+      // the database connection or background jobs.
+      await new Promise<void>((resolve) => {
+        this.server!.close(() => {
+          logger.log('info', 'Server closed, no longer accepting connections');
+          resolve();
         });
-      })));
-
-      // Stop accepting new connections and wait for existing ones to close
-      this.server.close(() => {
-        logger.log('info', 'Server closed, no longer accepting connections');
-        clearTimeout(drainTimer);
-        process.exit(0);
       });
+
+      clearTimeout(drainTimer);
+
+      // Only now run registered cleanup tasks (DB disconnect, job/poller
+      // shutdown) — in-flight work has already completed or been drained,
+      // so this can no longer interrupt a mid-commit transaction.
+      await this.runCleanupTasks();
+
+      process.exit(0);
     } catch (error) {
       logger.log('error', 'Error during graceful shutdown', {
         error: error instanceof Error ? error.message : String(error)
       });
       process.exit(1);
     }
   }
+
+  private async runCleanupTasks(): Promise<void> {
+    const cleanup = Promise.all(
+      this.cleanupTasks.map((task) =>
+        task().catch((err) => {
+          logger.log('error', 'Cleanup task failed during shutdown', {
+            error: err instanceof Error ? err.message : String(err),
+          });
+        }),
+      ),
+    );
+
+    let timeoutHandle: NodeJS.Timeout;
+    const timeout = new Promise<void>((resolve) => {
+      timeoutHandle = setTimeout(() => {
+        logger.log(
+          'warn',
+          `Cleanup tasks did not complete within ${this.cleanupTimeout}ms, proceeding with shutdown`,
+        );
+        resolve();
+      }, this.cleanupTimeout);
+      timeoutHandle.unref();
+    });
+
+    await Promise.race([cleanup, timeout]);
+    clearTimeout(timeoutHandle!);
+  }
 }

backend/src/index.ts

@@ -257,6 +257,11 @@ const cacheVaultMetricsTtl = parseInt(process.env.CACHE_TTL_MS || process.env.CA
 // Configure logger
 logger.configure(logLevel);
 
+// Coordinates SIGTERM/SIGINT: stops new requests immediately, drains
+// in-flight ones, then runs DB/job cleanup — see onShutdown registrations
+// near the bottom of this file.
+const shutdownHandler = new GracefulShutdownHandler(drainTimeout);
+
 void walletAliasMappingService.loadFromDatabase().catch((error) => {
   logger.log('error', 'Failed to warm wallet alias cache', {
     error: error instanceof Error ? error.message : String(error),
@@ -718,6 +723,25 @@ app.use((req: Request, res: Response, next: NextFunction) => {
   next();
 });
 
+// ─── Shutdown Gate ───────────────────────────────────────────────────────────
+// Once SIGTERM/SIGINT starts a drain, reject new requests with 503 instead of
+// routing them — the health check stays exempt so liveness probes keep seeing
+// 200 until the process actually exits.
+app.use((req: Request, res: Response, next: NextFunction) => {
+  if (shutdownHandler.isShuttingDown() && req.path !== '/health') {
+    res.set('Retry-After', '10');
+    res.status(503).json({
+      error: 'Service Unavailable',
+      status: 503,
+      code: 'SHUTTING_DOWN',
+      message: 'Server is shutting down and not accepting new requests',
+      retryable: true,
+    });
+    return;
+  }
+  next();
+});
+
 // Apply the Redis-backed default limiter (reads tier) globally (skip health/ready probes).
 app.use((req: Request, res: Response, next: NextFunction) => {
   if (req.path === '/health' || req.path === '/ready') return next();
@@ -4798,6 +4822,12 @@ if (process.env.NODE_ENV !== 'test' && process.env.VAULT_CONTRACT_ID) {
     pollIntervalMs: parseInt(process.env.EVENT_POLL_INTERVAL_MS || '10000', 10),
     batchSize: parseInt(process.env.EVENT_REPLAY_BATCH_SIZE || '100', 10),
   });
+
+  // stopEventPollingService() is a no-op if the service was never started,
+  // so this is safe to register unconditionally alongside the block above.
+  shutdownHandler.onShutdown(async () => {
+    stopEventPollingService();
+  });
 }
 
 // ─── Outbox Pattern Processor ─────────────────────────────────────────
@@ -4810,6 +4840,16 @@ if (process.env.NODE_ENV !== 'test') {
       error: err instanceof Error ? err.message : String(err),
     });
   });
+  eventOutboxService.start();
+
+  // Register graceful shutdown for the outbox processor so pending events
+  // are not abandoned when the process receives a termination signal.
+  // Runs only after in-flight HTTP requests have drained (see
+  // GracefulShutdownHandler), so it can't race a request still writing to
+  // the outbox.
+  shutdownHandler.onShutdown(async () => {
+    eventOutboxService.stop();
+  });
   const shutdownController = new AbortController();
   const shutdownSignal = shutdownController.signal;
   eventOutboxService.start({ signal: shutdownSignal });
@@ -5415,4 +5455,25 @@ app.use((req: Request, res: Response) => {
   });
 });
 
+// ─── Server Startup & Graceful Shutdown ────────────────────────────────────
+// Tests import `app` directly via supertest and never bind a real port.
+if (process.env.NODE_ENV !== 'test') {
+  // Disconnect Prisma last, after the HTTP listener has finished draining
+  // in-flight requests (see GracefulShutdownHandler) — this is what stops a
+  // termination signal from severing an open prisma.vault.update transaction
+  // mid-commit and leaving a row locked for the next deploy's migration.
+  shutdownHandler.onShutdown(async () => {
+    await prisma.$disconnect();
+  });
+
+  const server = app.listen(port, () => {
+    logger.log('info', `Server listening on port ${port}`);
+  });
+  shutdownHandler.register(server);
+}
+
+// Exported for tests: lets the shutdown-gate middleware behavior be verified
+// without sending a real process signal (which would also tear down other
+// suites' background jobs sharing this module in the Jest worker).
+export { shutdownHandler };
 export default app;