stock-bot/libs/core/shutdown/src/shutdown.ts

import { getLogger } from '@stock-bot/logger';
import { SHUTDOWN_DEFAULTS } from './constants';
import type { ShutdownCallback, ShutdownOptions } from './types';

interface CallbackEntry {
  callback: ShutdownCallback;
  priority: number;
  name?: string;
}

export class Shutdown {
  private static instance: Shutdown | null = null;
  private readonly logger = getLogger('shutdown');
  private isShuttingDown = false;
  private shutdownTimeout: number;
  private callbacks: CallbackEntry[] = [];
  private signalHandlersRegistered = false;
  private shutdownPromise: Promise<void> | null = null;
  private triggerSignal: string | null = null;

  constructor(options: ShutdownOptions = {}) {
    this.shutdownTimeout = options.timeout || SHUTDOWN_DEFAULTS.TIMEOUT;
    if (options.autoRegister !== false) {
      // Register signal handlers immediately in constructor
      this.setupSignalHandlers();
    }
  }

  static getInstance(options?: ShutdownOptions): Shutdown {
    if (!Shutdown.instance) {
      Shutdown.instance = new Shutdown(options);
    }
    return Shutdown.instance;
  }

  /**
   * Get current shutdown status
   */
  getStatus(): { isShuttingDown: boolean; signal: string | null; callbackCount: number } {
    return {
      isShuttingDown: this.isShuttingDown,
      signal: this.triggerSignal,
      callbackCount: this.callbacks.length
    };
  }

  /**
   * Register a cleanup callback
   */
  onShutdown(
    callback: ShutdownCallback,
    priority: number = SHUTDOWN_DEFAULTS.MEDIUM_PRIORITY,
    name?: string
  ): void {
    if (this.isShuttingDown) {
      return;
    }
    this.callbacks.push({ callback, priority, name });
  }

  /**
   * Initiate graceful shutdown
   */
  async shutdown(signal?: string): Promise<void> {
    // If already shutting down, return the existing promise
    if (this.shutdownPromise) {
      this.logger.warn('Shutdown already in progress, ignoring additional request', { 
        originalSignal: this.triggerSignal,
        newSignal: signal 
      });
      return this.shutdownPromise;
    }

    if (this.isShuttingDown) {
      this.logger.warn('Shutdown flag already set but no promise, creating new shutdown');
    }

    this.isShuttingDown = true;
    this.triggerSignal = signal || 'manual';
    
    this.logger.info('Starting graceful shutdown', { 
      signal: this.triggerSignal,
      timeout: this.shutdownTimeout,
      callbackCount: this.callbacks.length 
    });

    // Create and store the shutdown promise
    this.shutdownPromise = this.performShutdown();
    
    try {
      await this.shutdownPromise;
      this.logger.info('Graceful shutdown completed successfully');
    } catch (error) {
      this.logger.error('Shutdown failed', { error, signal: this.triggerSignal });
      throw error;
    }
  }

  private async performShutdown(): Promise<void> {
    const timeout = new Promise<never>((_, reject) =>
      setTimeout(() => reject(new Error('Shutdown timeout')), this.shutdownTimeout)
    );

    try {
      await Promise.race([this.executeCallbacks(), timeout]);
    } catch (error) {
      throw error;
    }
  }

  private async executeCallbacks(): Promise<void> {
    const sorted = [...this.callbacks].sort((a, b) => a.priority - b.priority);
    
    this.logger.debug('Executing shutdown callbacks', { 
      count: sorted.length,
      callbacks: sorted.map(c => ({ name: c.name || 'unnamed', priority: c.priority }))
    });

    let completed = 0;
    let failed = 0;

    for (const { callback, name, priority } of sorted) {
      const callbackName = name || 'unnamed';
      try {
        this.logger.debug(`Starting shutdown callback: ${callbackName}`, { priority });
        const startTime = Date.now();
        
        await callback();
        
        const duration = Date.now() - startTime;
        completed++;
        this.logger.debug(`Completed shutdown callback: ${callbackName}`, { duration, priority });
      } catch (error) {
        failed++;
        this.logger.error(`Shutdown callback failed: ${callbackName}`, { error, priority });
      }
    }

    this.logger.info('Shutdown callbacks execution completed', { 
      total: sorted.length,
      completed,
      failed 
    });
  }

  private setupSignalHandlers(): void {
    if (this.signalHandlersRegistered) {
      this.logger.debug('Signal handlers already registered');
      return;
    }

    const signals: NodeJS.Signals[] = ['SIGTERM', 'SIGINT'];

    signals.forEach(signal => {
      // Use process.on instead of process.once to handle multiple signals
      process.on(signal, async () => {
        this.logger.info(`Received ${signal} signal`);
        
        // If already shutting down, ignore the signal
        if (this.isShuttingDown || this.shutdownPromise) {
          this.logger.info(`Already shutting down, ignoring ${signal}`);
          return;
        }

        try {
          await this.shutdown(signal);
          
          // Add small delay to ensure logs are flushed
          await new Promise(resolve => setTimeout(resolve, 100));
          
          this.logger.info('Exiting process with code 0');
          process.exit(0);
        } catch (error) {
          this.logger.error('Shutdown error, exiting with code 1', { error });
          
          // Add small delay to ensure error logs are flushed
          await new Promise(resolve => setTimeout(resolve, 100));
          
          process.exit(1);
        }
      });
      
      this.logger.debug(`Registered handler for ${signal}`);
    });

    this.signalHandlersRegistered = true;
    this.logger.info('Signal handlers registered for graceful shutdown', { signals });
  }
}