smartexit/ts/smartexit.classes.smartexit.ts

import * as plugins from './smartexit.plugins.js';
import { ProcessLifecycle } from './smartexit.classes.lifecycle.js';

export interface ISmartExitOptions {
  /** Completely disable logging for this instance. */
  silent?: boolean;
}

export type TProcessSignal =
  | 'SIGHUP'
  | 'SIGINT'
  | 'SIGQUIT'
  | 'SIGILL'
  | 'SIGTRAP'
  | 'SIGABRT'
  | 'SIGIOT'
  | 'SIGBUS'
  | 'SIGFPE'
  | 'SIGKILL'
  | 'SIGUSR1'
  | 'SIGSEGV'
  | 'SIGUSR2'
  | 'SIGPIPE'
  | 'SIGALRM'
  | 'SIGTERM'
  | 'SIGCHLD'
  | 'SIGCONT'
  | 'SIGSTOP'
  | 'SIGTSTP'
  | 'SIGTTIN'
  | 'SIGTTOU'
  | 'SIGURG'
  | 'SIGXCPU'
  | 'SIGXFSZ'
  | 'SIGVTALRM'
  | 'SIGPROF'
  | 'SIGWINCH'
  | 'SIGPOLL'
  | 'SIGIO'
  | 'SIGPWR'
  | 'SIGINFO'
  | 'SIGLOST'
  | 'SIGSYS'
  | 'SIGUNUSED';

/**
 * SmartExit — process and cleanup tracker.
 *
 * Lightweight: the constructor does NOT register signal handlers.
 * Each instance auto-registers with ProcessLifecycle so that global
 * shutdown (triggered by ProcessLifecycle.install()) kills all tracked processes.
 *
 * Libraries should create instances freely. Only the application entry point
 * should call `ProcessLifecycle.install()`.
 */
export class SmartExit {
  /** Kill an entire process tree by PID. */
  public static async killTreeByPid(pidArg: number, signalArg: TProcessSignal = 'SIGKILL') {
    if (process.platform === 'win32') {
      // Windows: use native taskkill for tree kill
      try {
        plugins.childProcess.execSync(`taskkill /T /F /PID ${pidArg}`, { stdio: 'ignore' });
      } catch {
        // Process already dead
      }
    } else {
      // POSIX: kill the entire process group via negative PID.
      // Works even for grandchildren reparented to PID 1.
      try {
        process.kill(-pidArg, signalArg);
      } catch (err: any) {
        if (err.code !== 'ESRCH') {
          // ESRCH = no such process/group, already dead — that's fine
          throw err;
        }
      }
    }
  }

  // Instance state
  public processesToEnd = new plugins.lik.ObjectMap<plugins.childProcess.ChildProcess>();
  public cleanupFunctions = new plugins.lik.ObjectMap<() => Promise<any>>();
  /** PIDs tracked independently for tree-killing during shutdown. */
  public trackedPids = new Set<number>();
  private options: ISmartExitOptions;

  private log(message: string, isError = false): void {
    if (this.options.silent) {
      return;
    }
    const prefix = '[smartexit]';
    if (isError) {
      console.error(`${prefix} ${message}`);
    } else {
      console.log(`${prefix} ${message}`);
    }
  }

  constructor(optionsArg: ISmartExitOptions = {}) {
    this.options = optionsArg;
    // Auto-register with the global ProcessLifecycle registry
    ProcessLifecycle.registerInstance(this);
  }

  /** Register a child process for cleanup on shutdown. */
  public addProcess(childProcessArg: plugins.childProcess.ChildProcess) {
    this.processesToEnd.add(childProcessArg);
    if (childProcessArg.pid) {
      this.trackedPids.add(childProcessArg.pid);
    }
  }

  /** Register an async cleanup function to run on shutdown. */
  public addCleanupFunction(cleanupFunctionArg: () => Promise<any>) {
    this.cleanupFunctions.add(cleanupFunctionArg);
  }

  /** Unregister a child process. */
  public removeProcess(childProcessArg: plugins.childProcess.ChildProcess) {
    this.processesToEnd.remove(childProcessArg);
    if (childProcessArg.pid) {
      this.trackedPids.delete(childProcessArg.pid);
    }
  }

  /**
   * Run cleanup functions, then kill all tracked process trees.
   * Called by ProcessLifecycle during global shutdown.
   * Can also be called manually.
   */
  public async killAll(): Promise<{ processesKilled: number; cleanupFunctionsRan: number }> {
    const processes = this.processesToEnd.getArray();
    const cleanupFuncs = this.cleanupFunctions.getArray();
    let processesKilled = 0;
    let cleanupFunctionsRan = 0;

    // Phase 1: Run cleanup functions (processes still alive)
    for (const cleanupFunction of cleanupFuncs) {
      try {
        await cleanupFunction();
        cleanupFunctionsRan++;
      } catch (err) {
        this.log(`Cleanup function failed: ${err}`, true);
      }
    }

    // Phase 2: Kill ALL tracked process trees by PID.
    // We use trackedPids (not processesToEnd) because the process object may have
    // been removed by smartshell's exit handler before shutdown runs.
    // We do NOT check .killed — the direct child may be dead but grandchildren alive.
    for (const pid of this.trackedPids) {
      try {
        await SmartExit.killTreeByPid(pid, 'SIGTERM');
        processesKilled++;
      } catch {
        // SIGTERM failed — force kill
        try {
          await SmartExit.killTreeByPid(pid, 'SIGKILL');
          processesKilled++;
        } catch {
          // Process tree already dead — fine
        }
      }
    }
    this.trackedPids.clear();

    return { processesKilled, cleanupFunctionsRan };
  }

  /** Remove this instance from the global ProcessLifecycle registry. */
  public deregister(): void {
    ProcessLifecycle.deregisterInstance(this);
  }
}