diff --git a/.env.example b/.env.example index 079ff56b..b05eb3d2 100644 --- a/.env.example +++ b/.env.example @@ -96,6 +96,12 @@ PUBLIC_RATE_LIMIT_TRUST_PROXY=false # 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 ba361c0d..f0252c1c 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'; @@ -105,6 +106,7 @@ import { RequestIdInterceptor } from './common/interceptors/request-id.intercept ), }), + 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 66ad0c28..85aa0702 100644 --- a/src/common/locks/redis-lock.util.spec.ts +++ b/src/common/locks/redis-lock.util.spec.ts @@ -179,8 +179,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 fb5a97d4..e86ec988 100644 --- a/src/config/env.validation.ts +++ b/src/config/env.validation.ts @@ -136,6 +136,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 7a29a0ba..1ca792ab 100644 --- a/src/database/prisma.service.spec.ts +++ b/src/database/prisma.service.spec.ts @@ -25,6 +25,7 @@ import { withQueryTimeout, } from './query-timeout.extension'; import { buildDatasourceUrl } from './datasource-url'; +import { ShutdownCoordinator, ShutdownResource } from '../common/shutdown/shutdown-coordinator.service'; import { ConnectionPoolExhaustedError, DatabaseTimeoutError } from './database.errors'; const BASE_URL = 'postgresql://user:pass@localhost:5432/astroid?schema=public'; @@ -63,10 +64,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, shutdown); const service = new PrismaService(configService as unknown as ConfigService); Object.setPrototypeOf(service, PrismaService.prototype); return service; @@ -254,50 +256,78 @@ describe('PrismaService', () => { expect(service.workerClient).toBeDefined(); }); - 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(); - }); + 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('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('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); }); - await expect(service.onModuleInit()).rejects.toThrow( - 'Database migrations are not up to date: 1 pending migration(s)', - ); - }); + 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)); - it('fails startup after exhausting database connection attempts', async () => { - vi.useFakeTimers(); - const service = buildPrismaService(); - const apiConnect = vi.spyOn(service, '$connect').mockRejectedValue(new Error('unavailable')); + 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); + }); + 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(); + }); - const initialization = expect(service.onModuleInit()).rejects.toThrow('unavailable'); - await vi.runAllTimersAsync(); - await initialization; + 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 = withMethods(buildPrismaService()); + const apiConnect = vi.spyOn(service, '$connect').mockRejectedValue(new Error('unavailable')); - expect(apiConnect).toHaveBeenCalledTimes(databaseConfig.connectionRetryAttempts); - expect(checkMigrationStatusMock).not.toHaveBeenCalled(); + const initialization = expect(service.onModuleInit()).rejects.toThrow('unavailable'); + await vi.runAllTimersAsync(); + await initialization; + + expect(apiConnect).toHaveBeenCalledTimes(databaseConfig.connectionRetryAttempts); + expect(checkMigrationStatusMock).not.toHaveBeenCalled(); + }); }); }); diff --git a/src/database/prisma.service.ts b/src/database/prisma.service.ts index 6879ac76..545666da 100644 --- a/src/database/prisma.service.ts +++ b/src/database/prisma.service.ts @@ -4,10 +4,12 @@ import { 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 { createQueryMetricsExtension } from './query-metrics.extension'; import { createQueryTimeoutExtension } from './query-timeout.extension'; @@ -52,10 +54,12 @@ export class PrismaService extends PrismaClient implements OnModuleInit, OnModul */ readonly workerClient: PrismaClient; + /** True when the {@link ShutdownCoordinator} owns the disconnect timing. */ + private readonly coordinatedShutdown: boolean; private readonly migrationCheck: MigrationCheckMode; private readonly migrationsDir?: string; - constructor(configService: ConfigService) { + constructor(configService: ConfigService, @Optional() shutdown?: ShutdownCoordinator) { const database = configService.getOrThrow('database'); const url = buildDatasourceUrl(database.url, { @@ -107,6 +111,10 @@ export class PrismaService extends PrismaClient implements OnModuleInit, OnModul }), ) 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() }); this.migrationCheck = database.migrationCheck; this.migrationsDir = database.migrationsDir; } @@ -187,6 +195,15 @@ export class PrismaService extends PrismaClient implements OnModuleInit, OnModul } async onModuleDestroy(): Promise { + // 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(); + } + } + + /** 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 b747cac9..0c15d5f5 100644 --- a/src/main.ts +++ b/src/main.ts @@ -6,10 +6,11 @@ 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'; import { TOTAL_COUNT_HEADER } from './common/constants/headers'; import { assertValidEnvironment, EnvironmentValidationError } from './config/env.validation'; +import { PrismaService } from './database/prisma.service'; async function bootstrap() { // Fail fast on missing or malformed configuration, before any module is @@ -119,6 +120,14 @@ async function bootstrap() { 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 667cfa9c..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'; @@ -13,6 +14,9 @@ import { TokenVerificationCacheService } from './services/token-verification-cac import { CacheService } from '../../common/cache/cache.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,6 +34,22 @@ import { PasskeyService } from './services/passkey.service'; imports: [PassportModule.register({ defaultStrategy: 'jwt' }), JwtModule.register({})], controllers: [AuthController, PasskeyController], providers: [ + { + provide: Redis, + inject: [ShutdownCoordinator], + useFactory: (shutdown: ShutdownCoordinator): Redis => { + const config = redisConfig(); + 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; + }, + }, CacheService, TokenVerificationCacheService, AuthService,