From be4d70cc664c6ac632067231256dbfdbf1b6148c Mon Sep 17 00:00:00 2001 From: AdaBliss Date: Mon, 28 Sep 2026 09:27:46 -0700 Subject: [PATCH 1/3] feat(shutdown): drain HTTP and BullMQ work, then close queues, Redis and Prisma in order Add a ShutdownCoordinator that owns SIGTERM/SIGINT handling and runs a fixed, bounded sequence: 1. HTTP: stop accepting connections, close idle keep-alive sockets, mark in-flight responses `Connection: close` and wait for them. 2. Workers: pause every BullMQ worker (no new jobs) and wait for active jobs, then close it. Steps 1-2 share SHUTDOWN_GRACE_PERIOD_MS (default 20000). On expiry, remaining connections are destroyed and busy workers are force-closed, so their jobs are retried via stalled-job recovery. 3. app.close(): run every Nest lifecycle hook. From beforeApplicationShutdown the coordinator closes registered resources phase by phase: queues -> redis -> database, each bounded. The process exits 0 after a clean drain and 1 if the grace period expired or any step failed. Repeated signals are ignored while shutdown runs once. Logs carry resource names and timings only, never job data or URLs. Nest's app.enableShutdownHooks() is not used: it disconnects Prisma in onModuleDestroy before the HTTP server stops, has no grace period and exits by re-raising the signal. Providers keep ownership of their connections and register how to close them: PrismaService (database), the shared LocksModule Redis client and the auth Redis client (redis). Queues and workers created by @nestjs/bullmq are discovered at bootstrap. RedisLock no longer disconnects the shared client it does not own, and PrismaService's beforeExit hook is replaced by the coordinator. Closes #60 --- .env.example | 6 + docs/graceful-shutdown.md | 61 +++ src/app.module.ts | 2 + src/common/locks/locks.module.ts | 13 +- src/common/locks/redis-lock.util.spec.ts | 6 +- src/common/locks/redis-lock.util.ts | 10 +- .../shutdown/close-redis-client.spec.ts | 48 +++ src/common/shutdown/close-redis-client.ts | 21 + src/common/shutdown/index.ts | 3 + .../shutdown-coordinator.service.spec.ts | 396 ++++++++++++++++++ .../shutdown/shutdown-coordinator.service.ts | 366 ++++++++++++++++ src/common/shutdown/shutdown.module.ts | 16 + src/config/env.validation.ts | 8 + src/config/index.ts | 3 + src/config/shutdown.config.ts | 14 + src/database/prisma.service.spec.ts | 34 +- src/database/prisma.service.ts | 29 +- src/main.ts | 14 +- src/modules/auth/auth.module.ts | 9 +- 19 files changed, 1028 insertions(+), 31 deletions(-) create mode 100644 docs/graceful-shutdown.md create mode 100644 src/common/shutdown/close-redis-client.spec.ts create mode 100644 src/common/shutdown/close-redis-client.ts create mode 100644 src/common/shutdown/index.ts create mode 100644 src/common/shutdown/shutdown-coordinator.service.spec.ts create mode 100644 src/common/shutdown/shutdown-coordinator.service.ts create mode 100644 src/common/shutdown/shutdown.module.ts create mode 100644 src/config/shutdown.config.ts diff --git a/.env.example b/.env.example index 40dc8b43..fc6fba49 100644 --- a/.env.example +++ b/.env.example @@ -77,6 +77,12 @@ RATE_LIMIT_MAX_REQUESTS=120 # Prometheus server outside the private network. METRICS_ALLOWED_IPS=127.0.0.1/32,::1/128,10.0.0.0/8,172.16.0.0/12,192.168.0.0/16 +# Graceful shutdown +# On SIGTERM/SIGINT, in-flight HTTP requests and BullMQ jobs get this long (ms) +# to finish before being forcibly terminated. Keep it ~10s below the +# orchestrator's kill deadline (Kubernetes default: 30s). +SHUTDOWN_GRACE_PERIOD_MS=20000 + # AI Provider (Nvidia NIM API compatible with OpenAI SDK) AI_PROVIDER=nvidia AI_PROVIDER_KEY=your_nvidia_api_key_starting_with_nvapi diff --git a/docs/graceful-shutdown.md b/docs/graceful-shutdown.md new file mode 100644 index 00000000..baa378a5 --- /dev/null +++ b/docs/graceful-shutdown.md @@ -0,0 +1,61 @@ +# Graceful Shutdown + +On `SIGTERM` or `SIGINT`, the API drains in-flight work before exiting, so a +deployment never abandons an HTTP request or a BullMQ job halfway through a +financial operation. + +The sequence is run by `ShutdownCoordinator` (`src/common/shutdown`), installed +from `src/main.ts`: + +| Step | What happens | Bound | +| --- | --- | --- | +| 1. HTTP | The server stops accepting connections. Idle keep-alive connections are closed and in-flight responses carry `Connection: close`. | Grace period | +| 2. Workers | Every BullMQ worker is paused (no new jobs are fetched) and its active jobs finish, then the worker is closed. | Grace period (shared with step 1) | +| 3. Lifecycle hooks | `app.close()` runs Nest's `onModuleDestroy` / `beforeApplicationShutdown` / `onApplicationShutdown` hooks, so every provider releases what it owns. | 15s | +| 4. Queues | BullMQ queues are closed. | 3s per resource | +| 5. Redis | Shared Redis clients are closed with `QUIT`. | 3s per resource | +| 6. Database | Both Prisma connection pools are disconnected. | 3s per resource | + +Steps 4-6 run inside step 3 (from `beforeApplicationShutdown`), so they also +run when a module is closed outside a signal, for example in tests. + +## Exit codes + +- `0`: everything drained and closed. +- `1`: the grace period expired with requests or jobs still in progress, or a + resource failed to close. The logs name the step and resource that failed. + They include resource names and timings only, never job payloads or + connection strings. + +When the grace period expires, remaining HTTP connections are destroyed and +busy workers are force-closed. The interrupted jobs' locks lapse, and BullMQ's +stalled-job recovery re-queues them for another worker. + +Repeated signals while shutdown is in progress are logged and ignored. +Shutdown runs once and is always bounded, so it cannot hang. + +## Configuration + +| Variable | Default | Description | +| --- | --- | --- | +| `SHUTDOWN_GRACE_PERIOD_MS` | `20000` | Time in-flight HTTP requests and BullMQ jobs get to finish. | + +Keep the grace period about 10 seconds below your orchestrator's kill +deadline, to leave room for steps 3-6. Kubernetes' default +`terminationGracePeriodSeconds` is 30. + +## Adding a resource + +Providers keep ownership of the connections they create. They register how to +close each connection, and the coordinator decides when: + +```ts +shutdown.register({ + name: 'redis:my-feature', // shown in logs; never a URL or credential + phase: 'redis', // 'queues' | 'redis' | 'database' + close: () => closeRedisClient(client), +}); +``` + +Workers and queues created with `@nestjs/bullmq` (`@Processor`, +`BullModule.registerQueue`) are discovered automatically. diff --git a/src/app.module.ts b/src/app.module.ts index 443b946b..39bb78d3 100644 --- a/src/app.module.ts +++ b/src/app.module.ts @@ -11,6 +11,7 @@ import { RedisThrottlerStorage } from './common/throttler/redis-throttler.storag import { DatabaseModule } from './database/database.module'; import { EventsModule } from './events/events.module'; import { LocksModule } from './common/locks/locks.module'; +import { ShutdownModule } from './common/shutdown/shutdown.module'; import { REDIS_CLIENT } from './common/locks/locks.constants'; import { EncryptionModule } from './common/encryption/encryption.module'; import { RequestIdMiddleware } from './middleware/request-id.middleware'; @@ -103,6 +104,7 @@ import { AuditLogInterceptor } from './common/interceptors/audit-log.interceptor ), }), + ShutdownModule, DatabaseModule, EventsModule, LocksModule, diff --git a/src/common/locks/locks.module.ts b/src/common/locks/locks.module.ts index 237c6f72..83f6c082 100644 --- a/src/common/locks/locks.module.ts +++ b/src/common/locks/locks.module.ts @@ -3,6 +3,8 @@ import { APP_INTERCEPTOR } from '@nestjs/core'; import { ConfigService } from '@nestjs/config'; import Redis from 'ioredis'; import { RedisConfig } from '../../config/redis.config'; +import { ShutdownCoordinator } from '../shutdown/shutdown-coordinator.service'; +import { closeRedisClient } from '../shutdown/close-redis-client'; import { REDIS_CLIENT } from './locks.constants'; import { RedisLock } from './redis-lock.util'; import { AgentLockInterceptor } from './agent-lock.interceptor'; @@ -14,7 +16,7 @@ import { TransactionLockInterceptor } from './transaction-lock.interceptor'; * Global distributed-locking infrastructure. * * Provides a single shared ioredis client and the {@link RedisLock} service to - * every module, and registers the {@link AgentLockInterceptor} and + * every module (the client is closed by the {@link ShutdownCoordinator}), and registers the {@link AgentLockInterceptor} and * {@link BudgetLockInterceptor} that enforce `@UseAgentLock()` and * `@UseBudgetLock()` on any decorated controller method. */ @@ -23,10 +25,13 @@ import { TransactionLockInterceptor } from './transaction-lock.interceptor'; providers: [ { provide: REDIS_CLIENT, - inject: [ConfigService], - useFactory: (config: ConfigService) => { + inject: [ConfigService, ShutdownCoordinator], + useFactory: (config: ConfigService, shutdown: ShutdownCoordinator) => { const { host, port, password, db } = config.getOrThrow('redis'); - return new Redis({ host, port, password, db }); + const client = new Redis({ host, port, password, db }); + // Closed after queues and before the database; see ShutdownCoordinator. + shutdown.register({ name: 'redis:shared', phase: 'redis', close: () => closeRedisClient(client) }); + return client; }, }, RedisLock, diff --git a/src/common/locks/redis-lock.util.spec.ts b/src/common/locks/redis-lock.util.spec.ts index f13dbbcd..e6ad413d 100644 --- a/src/common/locks/redis-lock.util.spec.ts +++ b/src/common/locks/redis-lock.util.spec.ts @@ -142,8 +142,8 @@ describe('RedisLock', () => { }); }); - it('disconnects the shared client on module destroy', () => { - lock.onModuleDestroy(); - expect(redis.disconnect).toHaveBeenCalled(); + it('does not close the shared client it does not own', () => { + expect((lock as unknown as Record).onModuleDestroy).toBeUndefined(); + expect(redis.disconnect).not.toHaveBeenCalled(); }); }); diff --git a/src/common/locks/redis-lock.util.ts b/src/common/locks/redis-lock.util.ts index 4d60b24d..a00c50e2 100644 --- a/src/common/locks/redis-lock.util.ts +++ b/src/common/locks/redis-lock.util.ts @@ -1,4 +1,4 @@ -import { Inject, Injectable, OnModuleDestroy } from '@nestjs/common'; +import { Inject, Injectable } from '@nestjs/common'; import { Redis } from 'ioredis'; import { randomUUID } from 'crypto'; import { LockNotAcquiredException } from '../exceptions/domain.exception'; @@ -42,13 +42,11 @@ export type LockRelease = () => Promise; * ``` */ @Injectable() -export class RedisLock implements OnModuleDestroy { +export class RedisLock { + // The shared client is owned by LocksModule and closed in the shutdown + // coordinator's `redis` phase, after queues have drained. constructor(@Inject(REDIS_CLIENT) private readonly redis: Redis) {} - onModuleDestroy(): void { - this.redis.disconnect(); - } - /** * Acquires a distributed lock with automatic expiration. * @param key - Lock key (should be unique per resource). diff --git a/src/common/shutdown/close-redis-client.spec.ts b/src/common/shutdown/close-redis-client.spec.ts new file mode 100644 index 00000000..4977c2c4 --- /dev/null +++ b/src/common/shutdown/close-redis-client.spec.ts @@ -0,0 +1,48 @@ +import { describe, expect, it, vi } from 'vitest'; +import Redis from 'ioredis'; +import { closeRedisClient } from './close-redis-client'; + +function client(status: string, quit: () => Promise = async () => 'OK') { + return { status, quit: vi.fn(quit), disconnect: vi.fn() }; +} + +describe('closeRedisClient', () => { + it('sends QUIT to a connected client so pending replies are delivered', async () => { + const redis = client('ready'); + + await closeRedisClient(redis as unknown as Redis); + + expect(redis.quit).toHaveBeenCalledTimes(1); + expect(redis.disconnect).not.toHaveBeenCalled(); + }); + + it.each(['wait', 'connecting', 'reconnecting'])( + 'disconnects a client in the %s state without QUIT', + async (status) => { + const redis = client(status); + + await closeRedisClient(redis as unknown as Redis); + + expect(redis.quit).not.toHaveBeenCalled(); + expect(redis.disconnect).toHaveBeenCalledTimes(1); + }, + ); + + it('is a no-op for a client that is already closed', async () => { + const redis = client('end'); + + await closeRedisClient(redis as unknown as Redis); + + expect(redis.quit).not.toHaveBeenCalled(); + expect(redis.disconnect).not.toHaveBeenCalled(); + }); + + it('falls back to disconnect when QUIT fails', async () => { + const redis = client('ready', async () => { + throw new Error('Connection is closed.'); + }); + + await expect(closeRedisClient(redis as unknown as Redis)).resolves.toBeUndefined(); + expect(redis.disconnect).toHaveBeenCalledTimes(1); + }); +}); diff --git a/src/common/shutdown/close-redis-client.ts b/src/common/shutdown/close-redis-client.ts new file mode 100644 index 00000000..837ab1f8 --- /dev/null +++ b/src/common/shutdown/close-redis-client.ts @@ -0,0 +1,21 @@ +import Redis from 'ioredis'; + +/** + * Closes an ioredis client gracefully: `QUIT` lets pending replies arrive + * before the socket closes. Clients that never connected (lazyConnect) or are + * already closed are simply disconnected, so this is safe to call repeatedly. + */ +export async function closeRedisClient(client: Redis): Promise { + if (client.status === 'end') { + return; + } + if (client.status !== 'ready') { + client.disconnect(); + return; + } + try { + await client.quit(); + } catch { + client.disconnect(); + } +} diff --git a/src/common/shutdown/index.ts b/src/common/shutdown/index.ts new file mode 100644 index 00000000..73b951af --- /dev/null +++ b/src/common/shutdown/index.ts @@ -0,0 +1,3 @@ +export * from './shutdown.module'; +export * from './shutdown-coordinator.service'; +export * from './close-redis-client'; diff --git a/src/common/shutdown/shutdown-coordinator.service.spec.ts b/src/common/shutdown/shutdown-coordinator.service.spec.ts new file mode 100644 index 00000000..5a98968c --- /dev/null +++ b/src/common/shutdown/shutdown-coordinator.service.spec.ts @@ -0,0 +1,396 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { INestApplication, Logger } from '@nestjs/common'; +import { ConfigService } from '@nestjs/config'; +import { DiscoveryService } from '@nestjs/core'; +import { WorkerHost } from '@nestjs/bullmq'; +import { Queue } from 'bullmq'; +import { ShutdownCoordinator, ShutdownPhase } from './shutdown-coordinator.service'; + +/** Resolves on demand, to model work that is still in flight. */ +function deferred() { + let resolve!: () => void; + const promise = new Promise((r) => (resolve = r)); + return { promise, resolve }; +} + +/** Lets pending promise callbacks run. */ +const flush = () => new Promise((resolve) => setImmediate(resolve)); + +class TestProcessor extends WorkerHost { + async process(): Promise { + return undefined; + } +} + +describe('ShutdownCoordinator', () => { + const GRACE_MS = 1_000; + let events: string[]; + let coordinator: ShutdownCoordinator; + let providers: Array<{ instance: unknown }>; + + let httpDone: ReturnType; + let server: { + listening: boolean; + close: ReturnType; + closeIdleConnections: ReturnType; + closeAllConnections: ReturnType; + prependListener: ReturnType; + }; + let app: { getHttpServer: () => typeof server; close: ReturnType }; + let worker: { + name: string; + pause: ReturnType; + close: ReturnType; + }; + let queueClose: ReturnType; + let errorLog: ReturnType; + + const resource = (name: string, phase: ShutdownPhase, close?: () => Promise) => + coordinator.register({ + name, + phase, + close: + close ?? + (async () => { + events.push(`${phase}:${name}`); + }), + }); + + beforeEach(() => { + events = []; + errorLog = vi.spyOn(Logger.prototype, 'error').mockImplementation(() => undefined); + vi.spyOn(Logger.prototype, 'log').mockImplementation(() => undefined); + vi.spyOn(Logger.prototype, 'warn').mockImplementation(() => undefined); + + httpDone = deferred(); + httpDone.resolve(); + server = { + listening: true, + close: vi.fn((callback: () => void) => { + events.push('http:stop-accepting'); + void httpDone.promise.then(() => { + events.push('http:drained'); + callback(); + }); + }), + closeIdleConnections: vi.fn(), + closeAllConnections: vi.fn(() => events.push('http:force-closed')), + prependListener: vi.fn(), + }; + + worker = { + name: 'webhooks', + pause: vi.fn(async () => { + events.push('worker:paused'); + }), + close: vi.fn(async (force?: boolean) => { + events.push(force ? 'worker:force-closed' : 'worker:closed'); + }), + }; + const host = new TestProcessor(); + Object.defineProperty(host, 'worker', { value: worker }); + + const queue = Object.create(Queue.prototype) as Queue; + queueClose = vi.fn(async () => { + events.push('queues:queue:webhooks'); + }); + Object.defineProperty(queue, 'name', { value: 'webhooks' }); + Object.defineProperty(queue, 'close', { value: queueClose }); + + providers = [ + { instance: host }, + { instance: queue }, + { instance: undefined }, + { instance: {} }, + ]; + + const config = { + getOrThrow: vi.fn().mockReturnValue({ gracePeriodMs: GRACE_MS }), + } as unknown as ConfigService; + const discovery = { getProviders: () => providers } as unknown as DiscoveryService; + coordinator = new ShutdownCoordinator(config, discovery); + + // Mirrors Nest's app.close(): lifecycle hooks, then beforeApplicationShutdown. + app = { + getHttpServer: () => server, + close: vi.fn(async () => { + events.push('app:lifecycle-hooks'); + await coordinator.beforeApplicationShutdown(); + }), + }; + + // Registered out of order on purpose: phases, not registration order, + // decide the close order. + resource('prisma', 'database'); + resource('redis:shared', 'redis'); + coordinator.onApplicationBootstrap(); + + coordinator.enableShutdownHooks(app as unknown as INestApplication, { + signals: [], + exit: vi.fn(), + }); + }); + + afterEach(() => { + vi.useRealTimers(); + vi.restoreAllMocks(); + }); + + it('drains and closes everything in order, then reports a clean exit', async () => { + await expect(coordinator.shutdown('SIGTERM')).resolves.toBe(0); + + expect(events).toEqual([ + 'http:stop-accepting', + 'http:drained', + 'worker:paused', + 'worker:closed', + 'app:lifecycle-hooks', + 'queues:queue:webhooks', + 'redis:redis:shared', + 'database:prisma', + ]); + expect(server.closeIdleConnections).toHaveBeenCalled(); + expect(server.closeAllConnections).not.toHaveBeenCalled(); + expect(worker.close).toHaveBeenCalledWith(); + expect(errorLog).not.toHaveBeenCalled(); + }); + + it('waits for in-flight requests and jobs before releasing queues, Redis and Prisma', async () => { + httpDone = deferred(); + const job = deferred(); + worker.pause.mockImplementation(async () => { + await job.promise; + events.push('worker:paused'); + }); + + const result = coordinator.shutdown('SIGTERM'); + await flush(); + expect(events).toEqual(['http:stop-accepting']); + expect(coordinator.isShuttingDown).toBe(true); + + httpDone.resolve(); + await flush(); + expect(events).toEqual(['http:stop-accepting', 'http:drained']); + expect(worker.pause).toHaveBeenCalled(); + + job.resolve(); + await expect(result).resolves.toBe(0); + expect(events.slice(2)).toEqual([ + 'worker:paused', + 'worker:closed', + 'app:lifecycle-hooks', + 'queues:queue:webhooks', + 'redis:redis:shared', + 'database:prisma', + ]); + }); + + it('force-closes workers whose jobs outlive the grace period and exits non-zero', async () => { + vi.useFakeTimers(); + worker.pause.mockReturnValue(new Promise(() => undefined)); + + const result = coordinator.shutdown('SIGTERM'); + await vi.advanceTimersByTimeAsync(GRACE_MS); + + await expect(result).resolves.toBe(1); + expect(worker.close).toHaveBeenCalledWith(true); + expect(events).toEqual([ + 'http:stop-accepting', + 'http:drained', + 'worker:force-closed', + 'app:lifecycle-hooks', + 'queues:queue:webhooks', + 'redis:redis:shared', + 'database:prisma', + ]); + expect(errorLog).toHaveBeenCalledWith( + expect.stringContaining( + 'Worker "webhooks": active jobs did not finish within the 1000ms grace period', + ), + ); + }); + + it('closes lingering HTTP connections when requests outlive the grace period', async () => { + vi.useFakeTimers(); + httpDone = deferred(); + + const result = coordinator.shutdown('SIGTERM'); + await vi.advanceTimersByTimeAsync(GRACE_MS); + + await expect(result).resolves.toBe(1); + expect(server.closeAllConnections).toHaveBeenCalled(); + expect(events.slice(0, 2)).toEqual(['http:stop-accepting', 'http:force-closed']); + expect(events).toContain('database:prisma'); + }); + + it('shares one grace period between HTTP and worker draining', async () => { + vi.useFakeTimers(); + httpDone = deferred(); + setTimeout(() => httpDone.resolve(), 800); + worker.pause.mockImplementation(() => new Promise((resolve) => setTimeout(resolve, 500))); + + const result = coordinator.shutdown('SIGTERM'); + await vi.advanceTimersByTimeAsync(GRACE_MS); + + // 800ms of HTTP drain leaves 200ms, less than the 500ms the job needs. + await expect(result).resolves.toBe(1); + expect(worker.close).toHaveBeenCalledWith(true); + }); + + it('keeps closing later phases when a resource fails, then exits non-zero', async () => { + const failing = new Error('connection reset'); + Object.assign(failing, { payload: { secret: 'job-data-must-not-be-logged' } }); + resource('redis:auth', 'redis', async () => { + throw failing; + }); + + await expect(coordinator.shutdown('SIGTERM')).resolves.toBe(1); + + expect(events).toContain('database:prisma'); + const logged = errorLog.mock.calls.map((call) => String(call[0])).join('\n'); + expect(logged).toContain('Closing redis resource "redis:auth" failed: Error: connection reset'); + expect(logged).not.toContain('job-data-must-not-be-logged'); + }); + + it('bounds a resource close that never settles', async () => { + vi.useFakeTimers(); + resource('redis:auth', 'redis', () => new Promise(() => undefined)); + + const result = coordinator.shutdown('SIGTERM'); + await vi.advanceTimersByTimeAsync(ShutdownCoordinator.RESOURCE_CLOSE_TIMEOUT_MS); + + await expect(result).resolves.toBe(1); + expect(events).toContain('database:prisma'); + }); + + it('exits non-zero when a worker fails to close', async () => { + worker.close.mockRejectedValue(new Error('redis gone')); + + await expect(coordinator.shutdown('SIGTERM')).resolves.toBe(1); + expect(events).toContain('database:prisma'); + }); + + it('still closes resources when the application lifecycle hooks hang', async () => { + vi.useFakeTimers(); + app.close.mockReturnValue(new Promise(() => undefined)); + + const result = coordinator.shutdown('SIGTERM'); + await vi.advanceTimersByTimeAsync(ShutdownCoordinator.APP_CLOSE_TIMEOUT_MS); + + await expect(result).resolves.toBe(1); + expect(events.slice(-3)).toEqual([ + 'queues:queue:webhooks', + 'redis:redis:shared', + 'database:prisma', + ]); + }); + + it('runs cleanup once when shutdown is requested repeatedly', async () => { + const first = coordinator.shutdown('SIGTERM'); + const second = coordinator.shutdown('SIGINT'); + + expect(second).toBe(first); + await expect(first).resolves.toBe(0); + await expect(coordinator.shutdown('SIGTERM')).resolves.toBe(0); + + expect(server.close).toHaveBeenCalledTimes(1); + expect(worker.pause).toHaveBeenCalledTimes(1); + expect(app.close).toHaveBeenCalledTimes(1); + expect(queueClose).toHaveBeenCalledTimes(1); + expect(events.filter((event) => event === 'database:prisma')).toHaveLength(1); + }); + + it('closes resources in phase order when the app is closed without a signal', async () => { + await coordinator.beforeApplicationShutdown(); + await coordinator.beforeApplicationShutdown(); + + expect(events).toEqual(['queues:queue:webhooks', 'redis:redis:shared', 'database:prisma']); + }); + + it('ignores registrations that arrive after resources started closing', async () => { + await coordinator.beforeApplicationShutdown(); + const late = vi.fn(async () => undefined); + + coordinator.register({ name: 'late', phase: 'redis', close: late }); + await coordinator.beforeApplicationShutdown(); + + expect(late).not.toHaveBeenCalled(); + }); + + it('skips HTTP draining when the server is not listening', async () => { + server.listening = false; + + await expect(coordinator.shutdown('SIGTERM')).resolves.toBe(0); + expect(server.close).not.toHaveBeenCalled(); + }); + + describe('signal handling', () => { + let handlers: Map void>; + let exit: ReturnType; + let fresh: ShutdownCoordinator; + + beforeEach(() => { + handlers = new Map(); + vi.spyOn(process, 'on').mockImplementation(((signal: string, handler: () => void) => { + handlers.set(signal, handler); + return process; + }) as typeof process.on); + exit = vi.fn(); + + fresh = new ShutdownCoordinator( + { getOrThrow: () => ({ gracePeriodMs: GRACE_MS }) } as unknown as ConfigService, + { getProviders: () => [] } as unknown as DiscoveryService, + ); + server.listening = false; + fresh.enableShutdownHooks( + { + getHttpServer: () => server, + close: vi.fn(async () => undefined), + } as unknown as INestApplication, + { exit }, + ); + }); + + it('listens for SIGTERM and SIGINT', () => { + expect([...handlers.keys()]).toEqual(['SIGTERM', 'SIGINT']); + }); + + it('exits with the shutdown result, once, however many signals arrive', async () => { + handlers.get('SIGTERM')?.(); + handlers.get('SIGINT')?.(); + handlers.get('SIGTERM')?.(); + await vi.waitFor(() => expect(exit).toHaveBeenCalled()); + await flush(); + + expect(exit).toHaveBeenCalledTimes(1); + expect(exit).toHaveBeenCalledWith(0); + }); + + it('does not install handlers twice', () => { + const before = handlers.size; + vi.mocked(process.on).mockClear(); + + fresh.enableShutdownHooks({ getHttpServer: () => server } as unknown as INestApplication, { + exit, + }); + + expect(process.on).not.toHaveBeenCalled(); + expect(handlers.size).toBe(before); + }); + + it('asks keep-alive clients to disconnect once shutdown begins', async () => { + const onRequest = server.prependListener.mock.calls.at(-1)?.[1] as ( + req: unknown, + res: { headersSent: boolean; setHeader: ReturnType }, + ) => void; + const res = { headersSent: false, setHeader: vi.fn() }; + + onRequest({}, res); + expect(res.setHeader).not.toHaveBeenCalled(); + + void fresh.shutdown('SIGTERM'); + onRequest({}, res); + expect(res.setHeader).toHaveBeenCalledWith('Connection', 'close'); + }); + }); +}); diff --git a/src/common/shutdown/shutdown-coordinator.service.ts b/src/common/shutdown/shutdown-coordinator.service.ts new file mode 100644 index 00000000..8a5a9816 --- /dev/null +++ b/src/common/shutdown/shutdown-coordinator.service.ts @@ -0,0 +1,366 @@ +import { + BeforeApplicationShutdown, + INestApplication, + Injectable, + Logger, + OnApplicationBootstrap, +} from '@nestjs/common'; +import { ConfigService } from '@nestjs/config'; +import { DiscoveryService } from '@nestjs/core'; +import { WorkerHost } from '@nestjs/bullmq'; +import { Queue, Worker } from 'bullmq'; +import { Server, ServerResponse } from 'http'; +import { ShutdownConfig } from '../../config/shutdown.config'; + +/** + * Resource phases, closed in this order: the reverse of the order in which the + * application depends on them. Queues publish through Redis-backed + * connections and are drained first; the shared Redis client and finally the + * database go last, once nothing can issue work against them. + */ +export const SHUTDOWN_RESOURCE_PHASES = ['queues', 'redis', 'database'] as const; + +export type ShutdownPhase = (typeof SHUTDOWN_RESOURCE_PHASES)[number]; + +/** A connection or handle released by the coordinator during shutdown. */ +export interface ShutdownResource { + /** Stable, non-sensitive label used in logs. Never a URL or credential. */ + name: string; + phase: ShutdownPhase; + /** Releases the resource. Owned and implemented by the provider that created it. */ + close(): Promise; +} + +export type ShutdownSignal = 'SIGTERM' | 'SIGINT'; + +export interface ShutdownHookOptions { + signals?: readonly ShutdownSignal[]; + /** Called with the exit code once shutdown completes. Defaults to `process.exit`. */ + exit?: (code: number) => void; +} + +type Outcome = 'settled' | 'timeout'; + +/** + * Coordinates graceful shutdown on SIGTERM / SIGINT so deployments do not + * abandon in-flight API requests or BullMQ jobs. + * + * Nest's built-in `enableShutdownHooks()` runs `onModuleDestroy` hooks (which + * disconnect Prisma) before it closes the HTTP server, has no grace period, + * and exits by re-raising the signal. The coordinator instead runs a fixed + * sequence, each step bounded so shutdown can never hang: + * + * 1. HTTP stop accepting connections; let in-flight requests finish + * 2. Workers pause every BullMQ worker and let active jobs finish + * (steps 1-2 share the configurable grace period; anything + * still running when it expires is forcibly terminated) + * 3. `app.close()`, which runs Nest's lifecycle hooks so every provider + * releases what it owns. From `beforeApplicationShutdown`, the + * coordinator closes registered resources phase by phase: + * queues -> redis -> database + * + * The process then exits 0 after a clean drain, or 1 if the grace period + * expired or any step failed. + * + * Providers keep ownership of their connections: they register a `close` + * callback via {@link register}, and the coordinator only decides when it + * runs. Queues and workers created by `@nestjs/bullmq` are discovered + * automatically at bootstrap. + */ +@Injectable() +export class ShutdownCoordinator implements OnApplicationBootstrap, BeforeApplicationShutdown { + /** Upper bound for closing a single registered resource. */ + static readonly RESOURCE_CLOSE_TIMEOUT_MS = 3_000; + /** Upper bound for `app.close()`, which includes every resource phase. */ + static readonly APP_CLOSE_TIMEOUT_MS = 15_000; + static readonly DEFAULT_SIGNALS: readonly ShutdownSignal[] = ['SIGTERM', 'SIGINT']; + + private readonly logger = new Logger(ShutdownCoordinator.name); + private readonly gracePeriodMs: number; + private readonly resources: ShutdownResource[] = []; + private readonly workers = new Set(); + + private app?: INestApplication; + private shuttingDown = false; + private shutdownPromise?: Promise; + private resourcesPromise?: Promise; + + constructor( + config: ConfigService, + private readonly discovery: DiscoveryService, + ) { + this.gracePeriodMs = config.getOrThrow('shutdown').gracePeriodMs; + } + + /** True once shutdown has begun. */ + get isShuttingDown(): boolean { + return this.shuttingDown; + } + + /** Registers a resource to be closed in its phase during shutdown. */ + register(resource: ShutdownResource): void { + if (this.resourcesPromise) { + this.logger.warn( + `Ignoring registration of "${resource.name}": resources are already closing`, + ); + return; + } + this.resources.push(resource); + } + + /** Discovers the BullMQ workers and queues created by `@nestjs/bullmq`. */ + onApplicationBootstrap(): void { + for (const { instance } of this.discovery.getProviders()) { + if (instance instanceof WorkerHost) { + const worker = workerOf(instance); + if (worker) { + this.workers.add(worker); + } + } else if (instance instanceof Queue) { + const queue = instance; + this.register({ name: `queue:${queue.name}`, phase: 'queues', close: () => queue.close() }); + } + } + } + + /** + * Installs SIGTERM / SIGINT handlers that run {@link shutdown} and exit with + * its result. Called once from `main.ts` after the HTTP server is listening. + * Repeated signals are logged and ignored while shutdown is in progress. + */ + enableShutdownHooks(app: INestApplication, options: ShutdownHookOptions = {}): void { + if (this.app) { + return; + } + this.app = app; + const exit = options.exit ?? ((code: number) => process.exit(code)); + + // Ask keep-alive clients to disconnect after their current response, so + // connections drain instead of carrying new requests. Prepended so it runs + // before the application handler can send headers. + const server = app.getHttpServer() as Server; + server.prependListener('request', (_req, res: ServerResponse) => { + if (this.shuttingDown && !res.headersSent) { + res.setHeader('Connection', 'close'); + } + }); + + for (const signal of options.signals ?? ShutdownCoordinator.DEFAULT_SIGNALS) { + process.on(signal, () => { + if (this.shutdownPromise) { + this.logger.warn(`Received ${signal} while shutdown is already in progress; ignoring`); + return; + } + void this.shutdown(signal).then(exit); + }); + } + } + + /** + * Runs the shutdown sequence once and resolves with the process exit code: + * 0 after a clean drain, 1 if the grace period expired or a step failed. + * Concurrent and repeated calls share the same run. + */ + shutdown(reason = 'shutdown request'): Promise { + if (!this.shutdownPromise) { + this.shutdownPromise = this.run(reason); + } + return this.shutdownPromise; + } + + /** + * Closes registered resources phase by phase. Runs inside `app.close()` + * after every `onModuleDestroy` hook, so it also releases connections when + * a module is closed outside a signal (for example in tests). Idempotent. + */ + async beforeApplicationShutdown(): Promise { + await this.closeResources(); + } + + private async run(reason: string): Promise { + this.shuttingDown = true; + const startedAt = Date.now(); + const deadline = startedAt + this.gracePeriodMs; + this.logger.log(`Received ${reason}; shutting down (grace period ${this.gracePeriodMs}ms)`); + + const failures: string[] = []; + try { + failures.push(...(await this.drainHttp(deadline))); + failures.push(...(await this.drainWorkers(deadline))); + failures.push(...(await this.closeApplication())); + } catch (error) { + failures.push(`unexpected error: ${describe(error)}`); + } + + const elapsedMs = Date.now() - startedAt; + if (failures.length > 0) { + this.logger.error( + `Shutdown finished with ${failures.length} problem(s) after ${elapsedMs}ms: ${failures.join('; ')}`, + ); + return 1; + } + this.logger.log(`Shutdown complete in ${elapsedMs}ms`); + return 0; + } + + /** Stops accepting connections and waits for in-flight requests. */ + private async drainHttp(deadline: number): Promise { + const server = this.app?.getHttpServer() as Server | undefined; + if (!server?.listening) { + return []; + } + + const startedAt = Date.now(); + const closed = new Promise((resolve) => server.close(() => resolve())); + server.closeIdleConnections(); + + if ((await settleBy(closed, deadline)) === 'timeout') { + server.closeAllConnections(); + this.logger.error( + `HTTP: in-flight requests did not finish within the ${this.gracePeriodMs}ms grace period; ` + + 'closed remaining connections', + ); + return ['http: grace period expired with requests in flight']; + } + this.logger.log( + `HTTP: stopped accepting connections and drained in-flight requests (${Date.now() - startedAt}ms)`, + ); + return []; + } + + /** + * Pauses every worker (no new jobs are fetched) and waits for active jobs, + * then closes it. Workers still busy at the deadline are force-closed; their + * jobs' locks lapse and BullMQ's stalled-job recovery re-queues them. + */ + private async drainWorkers(deadline: number): Promise { + const results = await Promise.all( + Array.from(this.workers, async (worker): Promise => { + const startedAt = Date.now(); + try { + if ((await settleBy(worker.pause(), deadline)) === 'timeout') { + this.logger.error( + `Worker "${worker.name}": active jobs did not finish within the ${this.gracePeriodMs}ms ` + + 'grace period; forcing close. Interrupted jobs will be retried after their lock expires', + ); + await settleWithin(worker.close(true), ShutdownCoordinator.RESOURCE_CLOSE_TIMEOUT_MS); + return `worker "${worker.name}": grace period expired with jobs in progress`; + } + if ( + (await settleWithin(worker.close(), ShutdownCoordinator.RESOURCE_CLOSE_TIMEOUT_MS)) === + 'timeout' + ) { + return `worker "${worker.name}": close timed out`; + } + this.logger.log( + `Worker "${worker.name}": drained and closed (${Date.now() - startedAt}ms)`, + ); + return null; + } catch (error) { + this.logger.error(`Worker "${worker.name}": failed to close: ${describe(error)}`); + return `worker "${worker.name}": ${describe(error)}`; + } + }), + ); + return results.filter((failure): failure is string => failure !== null); + } + + /** Runs Nest's lifecycle hooks (which close registered resources). */ + private async closeApplication(): Promise { + const failures: string[] = []; + if (this.app) { + const app = this.app; + try { + if ( + (await settleWithin(app.close(), ShutdownCoordinator.APP_CLOSE_TIMEOUT_MS)) === 'timeout' + ) { + this.logger.error( + `Application lifecycle hooks did not finish within ${ShutdownCoordinator.APP_CLOSE_TIMEOUT_MS}ms`, + ); + failures.push('application: lifecycle hooks timed out'); + } + } catch (error) { + this.logger.error(`Application lifecycle hooks failed: ${describe(error)}`); + failures.push(`application: ${describe(error)}`); + } + } + // Resource failures are collected from the (idempotent) phase run, whether + // it happened inside app.close() or has to happen now. + failures.push(...(await this.closeResources())); + return failures; + } + + private closeResources(): Promise { + if (!this.resourcesPromise) { + this.resourcesPromise = this.closeResourcePhases(); + } + return this.resourcesPromise; + } + + private async closeResourcePhases(): Promise { + const failures: string[] = []; + for (const phase of SHUTDOWN_RESOURCE_PHASES) { + const members = this.resources.filter((resource) => resource.phase === phase); + const results = await Promise.all(members.map((resource) => this.closeResource(resource))); + failures.push(...results.filter((failure): failure is string => failure !== null)); + } + return failures; + } + + private async closeResource(resource: ShutdownResource): Promise { + const startedAt = Date.now(); + try { + const outcome = await settleWithin( + Promise.resolve().then(() => resource.close()), + ShutdownCoordinator.RESOURCE_CLOSE_TIMEOUT_MS, + ); + if (outcome === 'timeout') { + this.logger.error( + `Closing ${resource.phase} resource "${resource.name}" timed out after ${ShutdownCoordinator.RESOURCE_CLOSE_TIMEOUT_MS}ms`, + ); + return `${resource.phase} "${resource.name}": close timed out`; + } + this.logger.log( + `Closed ${resource.phase} resource "${resource.name}" (${Date.now() - startedAt}ms)`, + ); + return null; + } catch (error) { + this.logger.error( + `Closing ${resource.phase} resource "${resource.name}" failed: ${describe(error)}`, + ); + return `${resource.phase} "${resource.name}": ${describe(error)}`; + } + } +} + +/** The BullMQ worker behind a `@Processor` host, if it has been created. */ +function workerOf(host: WorkerHost): Worker | undefined { + try { + return host.worker; + } catch { + // The getter throws until @nestjs/bullmq has registered the worker. + return undefined; + } +} + +/** Resolves when `work` settles or at `deadline` (epoch ms), whichever is first. Rejections propagate. */ +function settleBy(work: Promise, deadline: number): Promise { + return settleWithin(work, Math.max(0, deadline - Date.now())); +} + +function settleWithin(work: Promise, timeoutMs: number): Promise { + let timer: NodeJS.Timeout | undefined; + const timeout = new Promise((resolve) => { + // Deliberately not unref'd: the pending deadline must keep the process + // alive until the step is resolved one way or the other. + timer = setTimeout(() => resolve('timeout'), timeoutMs); + }); + return Promise.race([work.then((): Outcome => 'settled'), timeout]).finally(() => + clearTimeout(timer), + ); +} + +/** Error summary for logs: class and message only, never payloads. */ +function describe(error: unknown): string { + return error instanceof Error ? `${error.name}: ${error.message}` : String(error); +} diff --git a/src/common/shutdown/shutdown.module.ts b/src/common/shutdown/shutdown.module.ts new file mode 100644 index 00000000..7bab691d --- /dev/null +++ b/src/common/shutdown/shutdown.module.ts @@ -0,0 +1,16 @@ +import { Global, Module } from '@nestjs/common'; +import { DiscoveryModule } from '@nestjs/core'; +import { ShutdownCoordinator } from './shutdown-coordinator.service'; + +/** + * Global graceful-shutdown infrastructure. Providers that own connections + * inject {@link ShutdownCoordinator} and register how to close them; `main.ts` + * installs the signal handlers. + */ +@Global() +@Module({ + imports: [DiscoveryModule], + providers: [ShutdownCoordinator], + exports: [ShutdownCoordinator], +}) +export class ShutdownModule {} diff --git a/src/config/env.validation.ts b/src/config/env.validation.ts index 33be8784..7cf87a9e 100644 --- a/src/config/env.validation.ts +++ b/src/config/env.validation.ts @@ -101,6 +101,14 @@ export const metricsEnvSchema = z.object({ .default('127.0.0.1/32,::1/128,10.0.0.0/8,172.16.0.0/12,192.168.0.0/16'), }); +export const shutdownEnvSchema = z.object({ + // How long in-flight HTTP requests and BullMQ jobs may run after SIGTERM / + // SIGINT before they are forcibly terminated (ms). Keep this plus ~10s of + // resource-close headroom below the orchestrator's kill deadline (Kubernetes + // terminationGracePeriodSeconds defaults to 30s). + SHUTDOWN_GRACE_PERIOD_MS: z.coerce.number().int().positive().default(20_000), +}); + export const aiEnvSchema = z.object({ AI_PROVIDER: z.string().default('nvidia'), AI_PROVIDER_KEY: z.string().min(1, 'AI_PROVIDER_KEY is required'), diff --git a/src/config/index.ts b/src/config/index.ts index b89ea49e..5c7c95b6 100644 --- a/src/config/index.ts +++ b/src/config/index.ts @@ -11,6 +11,7 @@ import { metricsConfig } from './metrics.config'; import { encryptionConfig } from './encryption.config'; import { rateLimitConfig } from './rate-limit.config'; import { throttlerConfig } from './throttler.config'; +import { shutdownConfig } from './shutdown.config'; export * from './app.config'; export * from './database.config'; @@ -24,6 +25,7 @@ export * from './metrics.config'; export * from './encryption.config'; export * from './rate-limit.config'; export * from './throttler.config'; +export * from './shutdown.config'; /** * Global configuration module. Every slice is registered via `registerAs` and @@ -46,5 +48,6 @@ export const AppConfigModule = ConfigModule.forRoot({ encryptionConfig, rateLimitConfig, throttlerConfig, + shutdownConfig, ], }); diff --git a/src/config/shutdown.config.ts b/src/config/shutdown.config.ts new file mode 100644 index 00000000..da129e4c --- /dev/null +++ b/src/config/shutdown.config.ts @@ -0,0 +1,14 @@ +import { registerAs } from '@nestjs/config'; +import { shutdownEnvSchema, validateEnv } from './env.validation'; + +export type ShutdownConfig = { + /** Time in-flight HTTP requests and BullMQ jobs get to finish on shutdown. */ + gracePeriodMs: number; +}; + +export const shutdownConfig = registerAs('shutdown', (): ShutdownConfig => { + const env = validateEnv(shutdownEnvSchema, process.env); + return { + gracePeriodMs: env.SHUTDOWN_GRACE_PERIOD_MS, + }; +}); diff --git a/src/database/prisma.service.spec.ts b/src/database/prisma.service.spec.ts index fe386ae3..6fc60ce4 100644 --- a/src/database/prisma.service.spec.ts +++ b/src/database/prisma.service.spec.ts @@ -18,6 +18,7 @@ import { withQueryTimeout, } from './query-timeout.extension'; import { buildDatasourceUrl } from './datasource-url'; +import { ShutdownCoordinator, ShutdownResource } from '../common/shutdown/shutdown-coordinator.service'; import { ConnectionPoolExhaustedError, DatabaseTimeoutError, @@ -51,11 +52,11 @@ function createMockClient(): { }; } -function buildPrismaService(): PrismaService { +function buildPrismaService(shutdown?: ShutdownCoordinator): PrismaService { const configService = { getOrThrow: vi.fn().mockReturnValue(databaseConfig), }; - return new PrismaService(configService as unknown as ConfigService); + return new PrismaService(configService as unknown as ConfigService, shutdown); } describe('withQueryTimeout', () => { @@ -226,4 +227,33 @@ describe('PrismaService', () => { const service = buildPrismaService(); expect(service.workerClient).toBeDefined(); }); + + describe('shutdown', () => { + // The PrismaClient mock returns a plain object from `super()`, so restore + // PrismaService's prototype to exercise its lifecycle methods. + const withMethods = (service: PrismaService) => Object.setPrototypeOf(service, PrismaService.prototype) as PrismaService; + + it('disconnects both pools on module destroy when running without a coordinator', async () => { + const service = withMethods(buildPrismaService()); + + await service.onModuleDestroy(); + + expect(service.$disconnect).toHaveBeenCalledTimes(1); + expect(service.workerClient.$disconnect).toHaveBeenCalledTimes(1); + }); + + it('registers in the database phase and leaves disconnect timing to the coordinator', async () => { + const registered: ShutdownResource[] = []; + const shutdown = { register: (r: ShutdownResource) => registered.push(r) }; + const service = withMethods(buildPrismaService(shutdown as unknown as ShutdownCoordinator)); + + await service.onModuleDestroy(); + expect(service.$disconnect).not.toHaveBeenCalled(); + + expect(registered).toEqual([expect.objectContaining({ name: 'prisma', phase: 'database' })]); + await registered[0].close(); + expect(service.$disconnect).toHaveBeenCalledTimes(1); + expect(service.workerClient.$disconnect).toHaveBeenCalledTimes(1); + }); + }); }); diff --git a/src/database/prisma.service.ts b/src/database/prisma.service.ts index cbbea166..2c4a4b87 100644 --- a/src/database/prisma.service.ts +++ b/src/database/prisma.service.ts @@ -1,7 +1,8 @@ -import { INestApplication, Injectable, Logger, OnModuleDestroy, OnModuleInit } from '@nestjs/common'; +import { Injectable, Logger, OnModuleDestroy, OnModuleInit, Optional } from '@nestjs/common'; import { ConfigService } from '@nestjs/config'; import { PrismaClient } from '@prisma/client'; import { DatabaseConfig } from '../config/database.config'; +import { ShutdownCoordinator } from '../common/shutdown/shutdown-coordinator.service'; import { buildDatasourceUrl } from './datasource-url'; import { createQueryTimeoutExtension } from './query-timeout.extension'; import { @@ -41,7 +42,10 @@ export class PrismaService extends PrismaClient implements OnModuleInit, OnModul */ readonly workerClient: PrismaClient; - constructor(configService: ConfigService) { + /** True when the {@link ShutdownCoordinator} owns the disconnect timing. */ + private readonly coordinatedShutdown: boolean; + + constructor(configService: ConfigService, @Optional() shutdown?: ShutdownCoordinator) { const database = configService.getOrThrow('database'); const url = buildDatasourceUrl(database.url, { @@ -93,6 +97,11 @@ export class PrismaService extends PrismaClient implements OnModuleInit, OnModul poolTimeoutMs: database.poolTimeoutMs, }), ) as unknown as PrismaClient; + + // The database is the last resource released on shutdown, after workers, + // queues and Redis, so nothing can still be issuing queries. + this.coordinatedShutdown = shutdown !== undefined; + shutdown?.register({ name: 'prisma', phase: 'database', close: () => this.disconnectAll() }); } async onModuleInit(): Promise { @@ -139,14 +148,16 @@ export class PrismaService extends PrismaClient implements OnModuleInit, OnModul } async onModuleDestroy(): Promise { - await this.$disconnect(); - await this.workerClient.$disconnect(); + // Under the coordinator, disconnecting here would run before workers and + // queues have finished; it disconnects in the `database` phase instead. + if (!this.coordinatedShutdown) { + await this.disconnectAll(); + } } - /** Registers a Nest shutdown hook so the process closes the pool cleanly. */ - async enableShutdownHooks(app: INestApplication): Promise { - process.on('beforeExit', () => { - void app.close(); - }); + /** Closes both connection pools. Safe to call more than once. */ + async disconnectAll(): Promise { + await this.$disconnect(); + await this.workerClient.$disconnect(); } } diff --git a/src/main.ts b/src/main.ts index 91df3f89..8f6b5097 100644 --- a/src/main.ts +++ b/src/main.ts @@ -6,7 +6,7 @@ import { Logger as PinoLogger } from 'nestjs-pino'; import helmet from 'helmet'; import { Request, Response, NextFunction } from 'express'; import { AppModule } from './app.module'; -import { PrismaService } from './database/prisma.service'; +import { ShutdownCoordinator } from './common/shutdown/shutdown-coordinator.service'; import { AppConfig } from './config/app.config'; async function bootstrap() { @@ -85,11 +85,15 @@ async function bootstrap() { SwaggerModule.setup('docs', app, document); } - // Prisma shutdown hook - const prisma = app.get(PrismaService); - await prisma.enableShutdownHooks(app); - await app.listen(appConfig.port); + + // Graceful shutdown on SIGTERM / SIGINT: stop accepting HTTP work, drain + // in-flight requests and BullMQ jobs within SHUTDOWN_GRACE_PERIOD_MS, then + // close queues, Redis and Prisma in order. Used instead of Nest's + // app.enableShutdownHooks(), which disconnects the database before the HTTP + // server stops and has no grace period; the coordinator still runs every + // Nest lifecycle hook through app.close(). + app.get(ShutdownCoordinator).enableShutdownHooks(app); console.log(`🚀 Astroid API listening on port ${appConfig.port}`); console.log(`📚 Swagger docs: http://localhost:${appConfig.port}/docs`); } diff --git a/src/modules/auth/auth.module.ts b/src/modules/auth/auth.module.ts index 5f64f557..963aa38b 100644 --- a/src/modules/auth/auth.module.ts +++ b/src/modules/auth/auth.module.ts @@ -13,6 +13,8 @@ import { TokenBlacklistService } from './services/token-blacklist.service'; import { PasskeyController } from './controllers/passkey.controller'; import { PasskeyService } from './services/passkey.service'; import { redisConfig } from '../../config/redis.config'; +import { ShutdownCoordinator } from '../../common/shutdown/shutdown-coordinator.service'; +import { closeRedisClient } from '../../common/shutdown/close-redis-client'; /** * Authentication module. Registers passport-jwt and api-key strategies and a bare @@ -30,15 +32,18 @@ import { redisConfig } from '../../config/redis.config'; providers: [ { provide: Redis, - useFactory: (): Redis => { + inject: [ShutdownCoordinator], + useFactory: (shutdown: ShutdownCoordinator): Redis => { const config = redisConfig(); - return new Redis({ + const client = new Redis({ host: config.host, port: config.port, password: config.password || undefined, db: config.db, lazyConnect: true, }); + shutdown.register({ name: 'redis:auth', phase: 'redis', close: () => closeRedisClient(client) }); + return client; }, }, AuthService, From 223ecaa00d9edc174d623768260dc0065b091344 Mon Sep 17 00:00:00 2001 From: AdaBliss Date: Tue, 6 Oct 2026 02:16:23 -0700 Subject: [PATCH 2/3] Fix: Repair merge corruption in prisma.service.ts and dependents A prior merge left prisma.service.ts and its spec with duplicate imports, a garbled duplicate .$extends() chain spliced into the constructor, and a missing closing brace in disconnectAll/describe, which broke TypeScript compilation. Also restores the missing PrismaService import in main.ts and the missing Redis import in auth.module.ts, and fixes three prisma.service.spec.ts tests that were missing the withMethods() prototype-restoring wrapper needed to call lifecycle methods on the mocked PrismaClient. Co-Authored-By: Claude Sonnet 5 --- src/database/prisma.service.spec.ts | 83 ++++++++++++++--------------- src/database/prisma.service.ts | 12 ++--- src/main.ts | 1 + src/modules/auth/auth.module.ts | 1 + 4 files changed, 45 insertions(+), 52 deletions(-) diff --git a/src/database/prisma.service.spec.ts b/src/database/prisma.service.spec.ts index b861415a..1ca792ab 100644 --- a/src/database/prisma.service.spec.ts +++ b/src/database/prisma.service.spec.ts @@ -26,10 +26,6 @@ import { } from './query-timeout.extension'; import { buildDatasourceUrl } from './datasource-url'; import { ShutdownCoordinator, ShutdownResource } from '../common/shutdown/shutdown-coordinator.service'; -import { - ConnectionPoolExhaustedError, - DatabaseTimeoutError, -} from './database.errors'; import { ConnectionPoolExhaustedError, DatabaseTimeoutError } from './database.errors'; const BASE_URL = 'postgresql://user:pass@localhost:5432/astroid?schema=public'; @@ -287,50 +283,51 @@ describe('PrismaService', () => { expect(service.$disconnect).toHaveBeenCalledTimes(1); expect(service.workerClient.$disconnect).toHaveBeenCalledTimes(1); }); - it('retries a transient database connection failure during startup', async () => { - vi.useFakeTimers(); - const service = buildPrismaService(); - const apiConnect = vi - .spyOn(service, '$connect') - .mockRejectedValueOnce(new Error('database starting')) - .mockResolvedValue(undefined); - const workerConnect = vi.spyOn(service.workerClient, '$connect').mockResolvedValue(undefined); - - const initialization = service.onModuleInit(); - await vi.runAllTimersAsync(); - await initialization; - - expect(apiConnect).toHaveBeenCalledTimes(2); - expect(workerConnect).toHaveBeenCalledOnce(); - expect(checkMigrationStatusMock).toHaveBeenCalledOnce(); - }); - - it('fails startup when migration health reports pending migrations', async () => { - const service = buildPrismaService(); - checkMigrationStatusMock.mockResolvedValue({ - upToDate: false, - migrations: [], - pending: [{ name: 'pending_migration', applied: false, finished: false, error: null }], - failed: [], - message: '1 pending migration(s)', + it('retries a transient database connection failure during startup', async () => { + vi.useFakeTimers(); + const service = withMethods(buildPrismaService()); + const apiConnect = vi + .spyOn(service, '$connect') + .mockRejectedValueOnce(new Error('database starting')) + .mockResolvedValue(undefined); + const workerConnect = vi.spyOn(service.workerClient, '$connect').mockResolvedValue(undefined); + + const initialization = service.onModuleInit(); + await vi.runAllTimersAsync(); + await initialization; + + expect(apiConnect).toHaveBeenCalledTimes(2); + expect(workerConnect).toHaveBeenCalledOnce(); + expect(checkMigrationStatusMock).toHaveBeenCalledOnce(); }); - await expect(service.onModuleInit()).rejects.toThrow( - 'Database migrations are not up to date: 1 pending migration(s)', - ); - }); + it('fails startup when migration health reports pending migrations', async () => { + const service = withMethods(buildPrismaService()); + checkMigrationStatusMock.mockResolvedValue({ + upToDate: false, + migrations: [], + pending: [{ name: 'pending_migration', applied: false, finished: false, error: null }], + failed: [], + message: '1 pending migration(s)', + }); + + await expect(service.onModuleInit()).rejects.toThrow( + 'Database migrations are not up to date: 1 pending migration(s)', + ); + }); - it('fails startup after exhausting database connection attempts', async () => { - vi.useFakeTimers(); - const service = buildPrismaService(); - const apiConnect = vi.spyOn(service, '$connect').mockRejectedValue(new Error('unavailable')); + it('fails startup after exhausting database connection attempts', async () => { + vi.useFakeTimers(); + const service = withMethods(buildPrismaService()); + const apiConnect = vi.spyOn(service, '$connect').mockRejectedValue(new Error('unavailable')); - const initialization = expect(service.onModuleInit()).rejects.toThrow('unavailable'); - await vi.runAllTimersAsync(); - await initialization; + const initialization = expect(service.onModuleInit()).rejects.toThrow('unavailable'); + await vi.runAllTimersAsync(); + await initialization; - expect(apiConnect).toHaveBeenCalledTimes(databaseConfig.connectionRetryAttempts); - expect(checkMigrationStatusMock).not.toHaveBeenCalled(); + expect(apiConnect).toHaveBeenCalledTimes(databaseConfig.connectionRetryAttempts); + expect(checkMigrationStatusMock).not.toHaveBeenCalled(); + }); }); }); diff --git a/src/database/prisma.service.ts b/src/database/prisma.service.ts index e5e1c38e..158a33fe 100644 --- a/src/database/prisma.service.ts +++ b/src/database/prisma.service.ts @@ -1,10 +1,10 @@ -import { Injectable, Logger, OnModuleDestroy, OnModuleInit, Optional } from '@nestjs/common'; import { INestApplication, Injectable, Logger, OnModuleDestroy, OnModuleInit, + Optional, } from '@nestjs/common'; import { ConfigService } from '@nestjs/config'; import { PrismaClient } from '@prisma/client'; @@ -111,14 +111,6 @@ export class PrismaService extends PrismaClient implements OnModuleInit, OnModul // queues and Redis, so nothing can still be issuing queries. this.coordinatedShutdown = shutdown !== undefined; shutdown?.register({ name: 'prisma', phase: 'database', close: () => this.disconnectAll() }); - }) - .$extends(createQueryMetricsExtension({ slowQueryThresholdMs: database.slowQueryThresholdMs })) - .$extends( - createQueryTimeoutExtension({ - queryTimeoutMs: database.workerQueryTimeoutMs, - poolTimeoutMs: database.poolTimeoutMs, - }), - ) as unknown as PrismaClient; } async onModuleInit(): Promise { @@ -192,6 +184,8 @@ export class PrismaService extends PrismaClient implements OnModuleInit, OnModul async disconnectAll(): Promise { await this.$disconnect(); await this.workerClient.$disconnect(); + } + /** * Reads live connection counts for this database from Postgres' * `pg_stat_activity`. Prisma's Rust query engine doesn't expose pool diff --git a/src/main.ts b/src/main.ts index 477a6f99..fe3a29c6 100644 --- a/src/main.ts +++ b/src/main.ts @@ -10,6 +10,7 @@ import { ShutdownCoordinator } from './common/shutdown/shutdown-coordinator.serv import { AppConfig } from './config/app.config'; import { assertValidEnvironment, EnvironmentValidationError } from './config/env.validation'; import { DatabaseConfig } from './config/database.config'; +import { PrismaService } from './database/prisma.service'; async function bootstrap() { // Fail fast on missing or malformed configuration, before any module is diff --git a/src/modules/auth/auth.module.ts b/src/modules/auth/auth.module.ts index f885d073..20819cd7 100644 --- a/src/modules/auth/auth.module.ts +++ b/src/modules/auth/auth.module.ts @@ -1,4 +1,5 @@ import { Module } from '@nestjs/common'; +import { Redis } from 'ioredis'; import { JwtModule } from '@nestjs/jwt'; import { PassportModule } from '@nestjs/passport'; import { AuthController } from './auth.controller'; From 0ccd1bb14a04cf46a155523e8a13b9792f3cd604 Mon Sep 17 00:00:00 2001 From: Astroid Dev Date: Wed, 7 Oct 2026 09:22:41 +0100 Subject: [PATCH 3/3] Remove unused import --- src/main.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/src/main.ts b/src/main.ts index ff09f3c6..0c15d5f5 100644 --- a/src/main.ts +++ b/src/main.ts @@ -10,7 +10,6 @@ import { ShutdownCoordinator } from './common/shutdown/shutdown-coordinator.serv import { AppConfig } from './config/app.config'; import { TOTAL_COUNT_HEADER } from './common/constants/headers'; import { assertValidEnvironment, EnvironmentValidationError } from './config/env.validation'; -import { DatabaseConfig } from './config/database.config'; import { PrismaService } from './database/prisma.service'; async function bootstrap() {