feat(graceful-shutdown): add enhanced production-ready shutdown sequence

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);
+    });
   });
 });

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',
+    );
   }
 }

lib/index.ts

@@ -10,3 +10,4 @@ export {
   HealthCheckStatus,
   HealthCheckResult,
 } from './health-check';
+export { GracefulShutdownService } from './graceful-shutdown-timeout/graceful-shutdown-timeout.service';

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;
 }

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,
     };
   }
 }