termination-handler.ts

import { eachPromise, formatDate, formatTask, getLogger } from './utils';

export interface DurableServerTask {
  id: number;
  startedAt: Date;
  description: string;
  details: any;
  end: () => void;
}

const logger = getLogger();

export const trappableSignals: Partial<Record<NodeJS.Signals, number>> = {
  SIGINT: 2,
  SIGTERM: 15,
};

export const reportOnlySignals: NodeJS.Signals[] = [
  'SIGABRT',
  'SIGALRM',
  'SIGBREAK',
  'SIGBUS',
  'SIGCHLD',
  'SIGCONT',
  'SIGFPE',
  'SIGHUP',
  'SIGILL',
  'SIGINFO',
  'SIGIO',
  'SIGIOT',
  'SIGLOST',
  'SIGPIPE',
  'SIGPOLL',
  'SIGPWR',
  'SIGQUIT',
  'SIGSEGV',
  'SIGSTKFLT',
  'SIGSYS',
  'SIGTRAP',
  'SIGTSTP',
  'SIGTTIN',
  'SIGTTOU',
  'SIGUNUSED',
  'SIGURG',
  'SIGUSR1',
  'SIGUSR2',
  'SIGVTALRM',
  'SIGWINCH',
  'SIGXCPU',
  'SIGXFSZ',
  'SIGBREAK',
  'SIGLOST',
  'SIGINFO',
];

export class TerminationHandler {
  checkRateSeconds = 0.5;
  reportRateSeconds = 30;
  lastReportTime: null | number;

  nextTaskId = 0;
  tasks: DurableServerTask[];

  shutDownHandlers: Array<() => Promise<unknown>>;
  shuttingDownSince: null | Date;
  shuttingDownTaskCount: null | number;

  constructor() {
    this.lastReportTime = null;
    this.tasks = [];
    this.shutDownHandlers = [];
    this.shuttingDownSince = null;
    this.shuttingDownTaskCount = null;
  }

  beginDurableTask(description: string, details = {}) {
    const newTask: DurableServerTask = {
      id: this.nextTaskId++,
      description,
      details,
      startedAt: new Date(),
      end: () => this.endDurableTask(newTask),
    };

    this.tasks.push(newTask);
    logger.info(`New durable task: ${formatTask(newTask)}`);

    return newTask;
  }

  addShutdownHandler(handler: () => Promise<unknown>) {
    this.shutDownHandlers.push(handler);
  }

  start() {
    const signals = Object.keys(trappableSignals) as NodeJS.Signals[];
    signals.forEach((signal) =>
      process.on(signal, () => this.handleSignal(signal)),
    );
    reportOnlySignals.forEach((signal) => {
      process.on(signal, () =>
        logger.debug(
          `(PID ${process.pid}/${process.ppid}) Received signal ${signal}`,
        ),
      );
    });
  }

  // Internal methods below

  private endDurableTask(endingTask: DurableServerTask) {
    if (this.shuttingDownSince === null) {
      logger.info(`Ending durable task: ${formatTask(endingTask)}`);
    } else {
      logger.warn(`Ending durable task: ${formatTask(endingTask)}`);
    }

    if (!this.tasks.includes(endingTask)) {
      logger.warn(
        `(PID ${process.pid}/${process.ppid}) Attempted to end the same task twice. This indicates a bug in the code beginning and ending the task. `,
      );
    }

    this.tasks = this.tasks.filter(
      (existingTask) => existingTask !== endingTask,
    );
  }

  async handleSignal(signal: NodeJS.Signals) {
    if (this.shuttingDownSince) {
      logger.warn(
        `(PID ${process.pid}/${process.ppid}) Received ${signal} while already gracefully shutting down, so displaying shutdown status instead:`,
      );
      this.report();
    } else {
      logger.warn(
        `(PID ${process.pid}/${process.ppid}) Received ${signal}. Beginning graceful termination.`,
      );
      this.shuttingDownSince = new Date();
      this.shuttingDownTaskCount = this.tasks.length;

      await eachPromise(this.shutDownHandlers);

      await this.waitForTasks();

      const signalValue = trappableSignals[signal]!;
      process.exit(128 + signalValue);
    }
  }

  async waitForTasks() {
    if (this.tasks.length > 0) {
      while (this.tasks.length > 0) {
        const now = Date.now();

        if (
          this.lastReportTime === null ||
          now - this.lastReportTime >= this.reportRateSeconds * 1000
        ) {
          this.report();
          this.lastReportTime = now;
        }

        await new Promise((resolve) =>
          setTimeout(resolve, this.checkRateSeconds * 1000),
        );
      }

      this.report();
    }
  }

  report() {
    const now = formatDate(new Date());

    if (this.tasks.length) {
      logger.warn(
        `[${now}] Terminating: (PID ${process.pid}/${process.ppid}) waiting for tasks to complete, ${this.tasks.length}/${this.shuttingDownTaskCount} remaining. This message will repeat every ${this.reportRateSeconds} seconds.`,
      );

      this.tasks.forEach((task, index) => {
        const taskText = formatTask(task);
        logger.warn(`\t[${index + 1}/${this.tasks.length}] ${taskText}`);
      });
    } else {
      const deltaSeconds =
        Math.floor((Date.now() - this.shuttingDownSince!.getTime()) / 100) / 10;
      logger.warn(
        `[${now}] Termination: (PID ${process.pid}/${process.ppid}) complete, ${this.shuttingDownTaskCount} tasks ended gracefully in ${deltaSeconds} seconds.`,
      );
    }
  }
}