cube-ms/src/graceful_shutdown.js

Production-ready microservice framework with health monitoring, validation, error handling, and Docker Swarm support

import connectionManager from './db_connection_manager.js';

class GracefulShutdown {
  constructor() {
    this.isShuttingDown = false;
    this.activeConnections = new Set();
    this.server = null;
    this.shutdownCallbacks = [];
    this.shutdownTimeout = 10000; // 10 seconds
  }

  setServer(server) {
    this.server = server;
    this.setupServerTracking();
  }

  setupServerTracking() {
    if (!this.server) return;

    // Track active connections
    this.server.on('connection', (socket) => {
      this.activeConnections.add(socket);
      
      socket.on('close', () => {
        this.activeConnections.delete(socket);
      });
    });
  }

  registerShutdownCallback(callback, name = 'unknown') {
    this.shutdownCallbacks.push({ callback, name });
  }

  async gracefulShutdown(signal) {
    if (this.isShuttingDown) {
      console.log(`Already shutting down, ignoring ${signal}`);
      return;
    }

    this.isShuttingDown = true;
    console.log(`\n🔄 Received ${signal}, initiating graceful shutdown...`);

    const shutdownPromise = this.performShutdown();
    const timeoutPromise = new Promise((_, reject) => 
      setTimeout(() => reject(new Error('Shutdown timeout')), this.shutdownTimeout)
    );

    try {
      await Promise.race([shutdownPromise, timeoutPromise]);
      console.log('✅ Graceful shutdown completed');
      process.exit(0);
    } catch (error) {
      console.error('❌ Shutdown timeout reached, forcing exit', error.message);
      process.exit(1);
    }
  }

  async performShutdown() {
    const steps = [
      { name: 'Stop accepting new requests', fn: () => this.stopAcceptingRequests() },
      { name: 'Execute shutdown callbacks', fn: () => this.executeShutdownCallbacks() },
      { name: 'Close active connections', fn: () => this.closeActiveConnections() },
      { name: 'Close server', fn: () => this.closeServer() },
      { name: 'Close database connections', fn: () => connectionManager.closeAllConnections() }
    ];

    for (const step of steps) {
      try {
        console.log(`🔄 ${step.name}...`);
        await step.fn();
        console.log(`✅ ${step.name} completed`);
      } catch (error) {
        console.error(`❌ ${step.name} failed:`, error.message);
      }
    }
  }

  async stopAcceptingRequests() {
    if (this.server) {
      this.server.close();
    }
  }

  async executeShutdownCallbacks() {
    const promises = this.shutdownCallbacks.map(async ({ callback, name }) => {
      try {
        await callback();
        console.log(`  ✅ ${name} shutdown callback completed`);
      } catch (error) {
        console.error(`  ❌ ${name} shutdown callback failed:`, error.message);
      }
    });

    await Promise.all(promises);
  }

  async closeActiveConnections() {
    if (this.activeConnections.size === 0) return;

    console.log(`  Closing ${this.activeConnections.size} active connections...`);
    
    // Give connections 5 seconds to close naturally
    const gracePeriod = new Promise(resolve => setTimeout(resolve, 5000));
    
    await Promise.race([
      gracePeriod,
      new Promise(resolve => {
        const checkConnections = () => {
          if (this.activeConnections.size === 0) {
            resolve();
          } else {
            setTimeout(checkConnections, 100);
          }
        };
        checkConnections();
      })
    ]);

    // Force close remaining connections
    for (const socket of this.activeConnections) {
      socket.destroy();
    }
  }

  async closeServer() {
    if (!this.server) return;

    return new Promise((resolve, reject) => {
      this.server.close((error) => {
        if (error) {
          reject(error);
        } else {
          resolve();
        }
      });
    });
  }

  setupSignalHandlers() {
    const signals = ['SIGTERM', 'SIGINT', 'SIGUSR2'];
    
    for (const signal of signals) {
      process.on(signal, () => this.gracefulShutdown(signal));
    }

    // Handle PM2 graceful reload
    process.on('message', (msg) => {
      if (msg === 'shutdown') {
        this.gracefulShutdown('PM2_SHUTDOWN');
      }
    });
  }

  isShuttingDownNow() {
    return this.isShuttingDown;
  }
}

// Singleton instance
const gracefulShutdown = new GracefulShutdown();

export default gracefulShutdown;