diff --git a/lib/graceful-shutdown-timeout/graceful-shutdown-timeout.service.spec.ts b/lib/graceful-shutdown-timeout/graceful-shutdown-timeout.service.spec.ts index f8767c19c..64d3f34ce 100644 --- a/lib/graceful-shutdown-timeout/graceful-shutdown-timeout.service.spec.ts +++ b/lib/graceful-shutdown-timeout/graceful-shutdown-timeout.service.spec.ts @@ -3,6 +3,8 @@ import { LoggerService } from '@nestjs/common'; import { GracefulShutdownService, TERMINUS_GRACEFUL_SHUTDOWN_TIMEOUT, + TERMINUS_ENABLE_ENHANCED_SHUTDOWN, + TERMINUS_BEFORE_SHUTDOWN_DELAY, } from './graceful-shutdown-timeout.service'; import { TERMINUS_LOGGER } from '../health-check/logger/logger.provider'; import { sleep } from '../utils'; @@ -22,26 +24,115 @@ describe('GracefulShutdownService', () => { let service: GracefulShutdownService; let logger: LoggerService; - beforeEach(async () => { - const module = await Test.createTestingModule({ - providers: [ - GracefulShutdownService, - { provide: TERMINUS_LOGGER, useValue: loggerMock }, - { provide: TERMINUS_GRACEFUL_SHUTDOWN_TIMEOUT, useValue: 1000 }, - ], - }).compile(); - - logger = module.get(TERMINUS_LOGGER); - service = module.get(GracefulShutdownService); + beforeEach(() => { + jest.clearAllMocks(); }); - it('should not trigger sleep if signal is not SIGTERM', async () => { - await service.beforeApplicationShutdown('SIGINT'); - expect(sleep).not.toHaveBeenCalled(); + describe('Standard graceful shutdown', () => { + beforeEach(async () => { + const module = await Test.createTestingModule({ + providers: [ + GracefulShutdownService, + { provide: TERMINUS_LOGGER, useValue: loggerMock }, + { provide: TERMINUS_GRACEFUL_SHUTDOWN_TIMEOUT, useValue: 1000 }, + { provide: TERMINUS_ENABLE_ENHANCED_SHUTDOWN, useValue: false }, + { provide: TERMINUS_BEFORE_SHUTDOWN_DELAY, useValue: 15000 }, + ], + }).compile(); + + logger = module.get(TERMINUS_LOGGER); + service = module.get(GracefulShutdownService); + }); + + it('should not trigger sleep if signal is not SIGTERM', async () => { + await service.beforeApplicationShutdown('SIGINT'); + expect(sleep).not.toHaveBeenCalled(); + }); + + it('should trigger sleep if signal is SIGTERM', async () => { + await service.beforeApplicationShutdown('SIGTERM'); + expect(sleep).toHaveBeenCalledWith(1000); + }); + + it('should not be shutting down initially', () => { + expect(service.isApplicationShuttingDown()).toBe(false); + }); }); - it('should trigger sleep if signal is SIGTERM', async () => { - await service.beforeApplicationShutdown('SIGTERM'); - expect(sleep).toHaveBeenCalledWith(1000); + describe('Enhanced graceful shutdown', () => { + beforeEach(async () => { + const module = await Test.createTestingModule({ + providers: [ + GracefulShutdownService, + { provide: TERMINUS_LOGGER, useValue: loggerMock }, + { provide: TERMINUS_GRACEFUL_SHUTDOWN_TIMEOUT, useValue: 5000 }, + { provide: TERMINUS_ENABLE_ENHANCED_SHUTDOWN, useValue: true }, + { provide: TERMINUS_BEFORE_SHUTDOWN_DELAY, useValue: 10000 }, + ], + }).compile(); + + logger = module.get(TERMINUS_LOGGER); + service = module.get(GracefulShutdownService); + }); + + it('should perform enhanced shutdown sequence', async () => { + expect(service.isApplicationShuttingDown()).toBe(false); + + await service.beforeApplicationShutdown('SIGTERM'); + + // Should wait for both delays + expect(sleep).toHaveBeenCalledTimes(2); + expect(sleep).toHaveBeenNthCalledWith(1, 10000); // beforeShutdownDelayMs + expect(sleep).toHaveBeenNthCalledWith(2, 5000); // gracefulShutdownTimeoutMs + }); + + it('should mark application as shutting down during enhanced shutdown', async () => { + const shutdownPromise = service.beforeApplicationShutdown('SIGTERM'); + + // After starting shutdown, should be marked as shutting down + expect(service.isApplicationShuttingDown()).toBe(true); + + await shutdownPromise; + }); + + it('should log appropriate messages during enhanced shutdown', async () => { + await service.beforeApplicationShutdown('SIGTERM'); + + expect(loggerMock.log).toHaveBeenCalledWith( + 'Received termination signal SIGTERM', + ); + expect(loggerMock.log).toHaveBeenCalledWith( + 'Enhanced graceful shutdown initiated - marking readiness probe as unhealthy', + ); + expect(loggerMock.log).toHaveBeenCalledWith( + 'Waiting 10000ms for load balancer to stop routing traffic', + ); + expect(loggerMock.log).toHaveBeenCalledWith( + 'Processing remaining requests for up to 5000ms', + ); + expect(loggerMock.log).toHaveBeenCalledWith( + 'Enhanced graceful shutdown complete, terminating application', + ); + }); + + it('should skip beforeShutdownDelayMs if set to 0', async () => { + const module = await Test.createTestingModule({ + providers: [ + GracefulShutdownService, + { provide: TERMINUS_LOGGER, useValue: loggerMock }, + { provide: TERMINUS_GRACEFUL_SHUTDOWN_TIMEOUT, useValue: 5000 }, + { provide: TERMINUS_ENABLE_ENHANCED_SHUTDOWN, useValue: true }, + { provide: TERMINUS_BEFORE_SHUTDOWN_DELAY, useValue: 0 }, + ], + }).compile(); + + service = module.get(GracefulShutdownService); + + await service.beforeApplicationShutdown('SIGTERM'); + + // Should only wait for gracefulShutdownTimeoutMs + expect(sleep).toHaveBeenCalledTimes(1); + expect(sleep).toHaveBeenCalledWith(5000); + }); }); }); diff --git a/lib/graceful-shutdown-timeout/graceful-shutdown-timeout.service.ts b/lib/graceful-shutdown-timeout/graceful-shutdown-timeout.service.ts index 6d6f9bbeb..6872e9f97 100644 --- a/lib/graceful-shutdown-timeout/graceful-shutdown-timeout.service.ts +++ b/lib/graceful-shutdown-timeout/graceful-shutdown-timeout.service.ts @@ -11,32 +11,87 @@ import { sleep } from '../utils'; export const TERMINUS_GRACEFUL_SHUTDOWN_TIMEOUT = 'TERMINUS_GRACEFUL_SHUTDOWN_TIMEOUT'; +export const TERMINUS_ENABLE_ENHANCED_SHUTDOWN = + 'TERMINUS_ENABLE_ENHANCED_SHUTDOWN'; + +export const TERMINUS_BEFORE_SHUTDOWN_DELAY = 'TERMINUS_BEFORE_SHUTDOWN_DELAY'; + /** * Handles Graceful shutdown timeout useful to await * for some time before the application shuts down. */ @Injectable() export class GracefulShutdownService implements BeforeApplicationShutdown { + private isShuttingDown = false; + constructor( @Inject(TERMINUS_LOGGER) private readonly logger: LoggerService, @Inject(TERMINUS_GRACEFUL_SHUTDOWN_TIMEOUT) private readonly gracefulShutdownTimeoutMs: number, + @Inject(TERMINUS_ENABLE_ENHANCED_SHUTDOWN) + private readonly enableEnhancedShutdown: boolean = false, + @Inject(TERMINUS_BEFORE_SHUTDOWN_DELAY) + private readonly beforeShutdownDelayMs: number = 15000, ) { if (this.logger instanceof ConsoleLogger) { this.logger.setContext(GracefulShutdownService.name); } } + /** + * Check if the application is currently shutting down + * Used to mark readiness probe as unhealthy during shutdown + */ + public isApplicationShuttingDown(): boolean { + return this.isShuttingDown; + } + async beforeApplicationShutdown(signal: string) { this.logger.log(`Received termination signal ${signal || ''}`); if (signal === 'SIGTERM') { + if (this.enableEnhancedShutdown) { + await this.performEnhancedShutdown(); + } else { + await this.performStandardGracefulShutdown(); + } + } + } + + private async performStandardGracefulShutdown() { + this.logger.log( + `Awaiting ${this.gracefulShutdownTimeoutMs}ms before shutdown`, + ); + await sleep(this.gracefulShutdownTimeoutMs); + this.logger.log(`Timeout reached, shutting down now`); + } + + private async performEnhancedShutdown() { + // Step 1: Mark application as shutting down (readiness probe will fail) + this.isShuttingDown = true; + this.logger.log( + 'Enhanced graceful shutdown initiated - marking readiness probe as unhealthy', + ); + + // Step 2: Wait for load balancer to stop routing traffic + if (this.beforeShutdownDelayMs > 0) { + this.logger.log( + `Waiting ${this.beforeShutdownDelayMs}ms for load balancer to stop routing traffic`, + ); + await sleep(this.beforeShutdownDelayMs); + } + + // Step 3: Wait for remaining requests to complete + if (this.gracefulShutdownTimeoutMs > 0) { this.logger.log( - `Awaiting ${this.gracefulShutdownTimeoutMs}ms before shutdown`, + `Processing remaining requests for up to ${this.gracefulShutdownTimeoutMs}ms`, ); await sleep(this.gracefulShutdownTimeoutMs); - this.logger.log(`Timeout reached, shutting down now`); } + + this.logger.log( + 'Enhanced graceful shutdown complete, terminating application', + ); } } diff --git a/lib/index.ts b/lib/index.ts index bead16219..0673dbf24 100644 --- a/lib/index.ts +++ b/lib/index.ts @@ -10,3 +10,4 @@ export { HealthCheckStatus, HealthCheckResult, } from './health-check'; +export { GracefulShutdownService } from './graceful-shutdown-timeout/graceful-shutdown-timeout.service'; diff --git a/lib/terminus-options.interface.ts b/lib/terminus-options.interface.ts index 06339328e..d77850fd9 100644 --- a/lib/terminus-options.interface.ts +++ b/lib/terminus-options.interface.ts @@ -25,4 +25,22 @@ export interface TerminusModuleOptions { * @default 0 */ gracefulShutdownTimeoutMs?: number; + /** + * Enable enhanced graceful shutdown sequence for production environments + * When enabled, the shutdown process will: + * 1. Mark readiness probe as unhealthy + * 2. Wait for load balancer to stop routing traffic + * 3. Process remaining requests + * 4. Close connections and shutdown + * @default false + */ + enableEnhancedShutdown?: boolean; + /** + * Time to wait (in ms) after marking readiness probe unhealthy + * before starting the shutdown process. + * This allows load balancers to detect and stop routing traffic. + * Only used when enableEnhancedShutdown is true. + * @default 15000 + */ + beforeShutdownDelayMs?: number; } diff --git a/lib/terminus.module.ts b/lib/terminus.module.ts index 1cdec9a24..06b70a741 100644 --- a/lib/terminus.module.ts +++ b/lib/terminus.module.ts @@ -2,6 +2,8 @@ import { type DynamicModule, Module, type Provider } from '@nestjs/common'; import { GracefulShutdownService, TERMINUS_GRACEFUL_SHUTDOWN_TIMEOUT, + TERMINUS_ENABLE_ENHANCED_SHUTDOWN, + TERMINUS_BEFORE_SHUTDOWN_DELAY, } from './graceful-shutdown-timeout/graceful-shutdown-timeout.service'; import { HealthCheckService } from './health-check'; import { getErrorLoggerProvider } from './health-check/error-logger/error-logger.provider'; @@ -44,6 +46,8 @@ export class TerminusModule { errorLogStyle = 'json', logger = true, gracefulShutdownTimeoutMs = 0, + enableEnhancedShutdown = false, + beforeShutdownDelayMs = 15000, } = options; const providers: Provider[] = [ @@ -52,19 +56,33 @@ export class TerminusModule { getLoggerProvider(logger), ]; - if (gracefulShutdownTimeoutMs > 0) { - providers.push({ + // Always provide the configuration values + providers.push( + { provide: TERMINUS_GRACEFUL_SHUTDOWN_TIMEOUT, useValue: gracefulShutdownTimeoutMs, - }); + }, + { + provide: TERMINUS_ENABLE_ENHANCED_SHUTDOWN, + useValue: enableEnhancedShutdown, + }, + { + provide: TERMINUS_BEFORE_SHUTDOWN_DELAY, + useValue: beforeShutdownDelayMs, + }, + ); + // Only add GracefulShutdownService if graceful shutdown is configured + const exportsArray: any[] = [...exports_]; + if (gracefulShutdownTimeoutMs > 0 || enableEnhancedShutdown) { providers.push(GracefulShutdownService); + exportsArray.push(GracefulShutdownService); } return { module: TerminusModule, providers, - exports: exports_, + exports: exportsArray, }; } }