lifecycle.ts

// Single responsibility: process lifecycle and graceful drain. Every role runs the same three
// phases on SIGTERM — stop accepting, finish in-flight, close resources — under one deadline,
// and reports the same /healthz + /readyz state.

import { type Clock, systemClock } from './clock';
import { UltimateError } from './errors';
import { finiteCount } from './finite-option';
import { settleWithin } from './lifecycle-deadline';
import { lifecycleDrained } from './lifecycle-errors';
import { type LogFields, type Logger, logger as rootLogger } from './logger';

export type HealthState = 'starting' | 'ready' | 'draining' | 'stopped';

/** Ordered. `accept` runs first, `close` last. */
export type ShutdownPhase = 'accept' | 'inflight' | 'close';

export const SHUTDOWN_PHASES: readonly ShutdownPhase[] = ['accept', 'inflight', 'close'];

/** Signals Ultimate reacts to. Narrower than `NodeJS.Signals` on purpose. */
export type ProcessSignal = 'SIGTERM' | 'SIGINT' | 'SIGHUP' | 'SIGQUIT';

export interface ShutdownReason {
  readonly signal: string;
  /**
   * Real monotonic ms (`systemClock`) after which hooks are abandoned — deliberately NOT the
   * injected clock. The budget this bounds is `terminationGracePeriodSeconds`, counted by the
   * kubelet in real seconds, so a frozen clock must be unable to extend it: read off `clock` a
   * test that advanced an hour of fake time handed the drain a 16-minute grace period, while
   * `waitForIdle` went on sleeping on a real `setTimeout`. `clock` still owns `uptimeMs`.
   */
  readonly deadlineAt: number;
}

export type ShutdownHook = (reason: ShutdownReason) => void | Promise<void>;

export interface OnShutdownOptions {
  readonly phase?: ShutdownPhase | undefined;
}

export interface LifecycleOptions {
  /**
   * The whole drain's budget — the in-flight wait AND every hook, in every phase. 25s by default,
   * and **enforced whether or not an app sets it**: `ShutdownReason.deadlineAt` was always computed
   * and handed to every hook, so the deadline was declared by the design and only the enforcement
   * was missing. No hook reads `deadlineAt`, which is why it has to be imposed here.
   *
   * The lever is a LARGER value, not the absence of one: a `worker` holding a 10-minute job wants
   * `configureLifecycle({ deadlineMs: 600_000 })` and a `terminationGracePeriodSeconds` at least as
   * large. Left at 25s it is abandoned and the process exits clean — the row's visibility lease
   * lapses and another worker re-claims it, which is what at-least-once already promises. The
   * alternative is not "the job finishes": it is the same duplicate, delivered by SIGKILL at the
   * kubelet's grace period, with no log line naming what overran.
   *
   * Screened where it is assigned: a whole number of milliseconds, 0 or more. `0` is "drain now".
   */
  readonly deadlineMs?: number | undefined;
  readonly clock?: Clock | undefined;
  readonly logger?: Logger | undefined;
}

export type ReadinessStatus = 'ok' | 'failing';

/**
 * Synchronous, and that is the design, not a limitation. **Do not widen this to
 * `() => Promise<boolean>`** — the signature is the mechanism.
 *
 * A readiness endpoint that does I/O is a liveness bomb. A probe that awaits a network call takes
 * as long as the dependency does, so a slow database makes the endpoint miss its `timeoutSeconds`,
 * the kubelet reads that as unready, and capacity is pulled from an already-struggling system —
 * the outage the probe existed to prevent, caused by the probe. Worse under a liveness probe
 * sharing the handler: the pod is killed and restarts into the same slow database, cold.
 *
 * So the owner of the dependency keeps a boolean fresh — a pool exposes `isOpen`, a background
 * poller flips a flag on its own schedule with its own timeout — and this reads it. That puts the
 * waiting where a timeout can be tuned, and leaves this path unable to block. A check that throws
 * is `failing`.
 */
export type ReadinessCheck = () => boolean;

export interface HealthReport {
  readonly state: HealthState;
  readonly ready: boolean;
  readonly uptimeMs: number;
  readonly inflight: number;
  readonly buildId: string;
  /** Named, because "alert on check failures BY CHECK NAME" is not writable against a boolean. */
  readonly checks: Readonly<Record<string, ReadinessStatus>>;
  /**
   * How many checks are registered. `checks: {}` reads identically for "every check passed" and
   * "nobody registered one", and only the second is a `/readyz` that means no more than "the
   * socket is bound" — which is what the chart's and compose's healthchecks route traffic on.
   * Reported rather than enforced: an empty registry is still ready, so a role that genuinely has
   * no dependency does not have to invent a check to boot.
   */
  readonly registered: number;
}

export interface HealthPayload {
  readonly ok: boolean;
  /** The status code the HTTP layer should return. Core stays HTTP-free; this is just data. */
  readonly status: number;
  readonly body: HealthReport;
}

interface Registration {
  readonly name: string;
  readonly phase: ShutdownPhase;
  readonly hook: ShutdownHook;
}

const DEFAULT_DEADLINE_MS = 25_000;

let deadlineMs = DEFAULT_DEADLINE_MS;
let clock: Clock = systemClock;
let log: Logger = rootLogger;
let state: HealthState = 'starting';
let startedAtMono = clock.monotonic();
let inflight = 0;
let registrations: Registration[] = [];
let drainPromise: Promise<void> | undefined;
let idleWaiters: (() => void)[] = [];
const readiness = new Map<string, ReadinessCheck>();

export function configureLifecycle(options: LifecycleOptions): void {
  // Screened above the write, never beside the arithmetic: `Math.max(0, deadlineAt - monotonic())`
  // PROPAGATES a NaN into `setTimeout(fn, NaN)`, i.e. 0 — measured, one in-flight operation dropped
  // and a 300ms close hook ABANDONED 111ms into a 25s budget, while `X_SHUTDOWN_TIMEOUT` rendered
  // `NaNms` and told the operator to RAISE a budget that was never a number. `min: 0` because 0 is
  // a real budget — drain now, no grace — and `@ultimat3/http`'s `drainTimeoutMs` accepts 0 and
  // hands it straight here, so a floor of 1 would refuse at boot what that package declares.
  if (options.deadlineMs !== undefined) {
    deadlineMs = finiteCount('configureLifecycle', 'deadlineMs', options.deadlineMs, 0);
  }
  if (options.clock !== undefined) {
    clock = options.clock;
    startedAtMono = clock.monotonic();
  }
  if (options.logger !== undefined) log = options.logger;
}

export function lifecycleState(): HealthState {
  return state;
}

/**
 * "This process bound its socket." NOT "this process can serve a request" — that is what the
 * readiness checks answer. Before them, `markReady()` was the whole of `/readyz`, so a pod went
 * green the instant it bound and the load balancer sent traffic into a Postgres pool that had not
 * opened a connection yet; `maxUnavailable: 0` does not help when readiness lies.
 *
 * **A drained lifecycle refuses this, and that is the whole of "one process, one lifecycle."**
 * `state` never leaves `stopped` and `drain()` memoizes, so a role that marked ready after a drain
 * used to be told nothing and go on to bind a socket answering 503 to everything, with no drain
 * left to close it. `X_LIFECYCLE_DRAINED` is that mistake named at the moment it is made, rather
 * than a lifecycle that can be restarted: two live lifecycles racing one shutdown is a worse
 * mechanism than a boot that fails.
 */
export function markReady(): void {
  if (isDraining()) throw lifecycleDrained(state === 'draining' ? 'draining' : 'stopped');
  if (state === 'starting') state = 'ready';
}

/**
 * Register a named readiness check. Returns its unregister — the same shape as `onShutdown`, and
 * owned by whoever can be started twice, for the same reason.
 */
export function registerReadinessCheck(name: string, check: ReadinessCheck): () => void {
  if (readiness.has(name)) {
    throw new UltimateError({
      code: 'X_READINESS_CHECK_DUPLICATE',
      cause: `a readiness check named "${name}" is already registered (have: ${[...readiness.keys()].join(', ')})`,
      fix: `name the second check for what it actually probes, e.g. registerReadinessCheck('${name}-replica', check) — or hold the unregister the first registration returned and call it first`,
      meta: { name },
    });
  }
  readiness.set(name, check);
  return () => {
    if (readiness.get(name) === check) readiness.delete(name);
  };
}

/** Test-only: registered checks. A count that climbs across a start/stop cycle is a leak. */
export function readinessCheckCount(): number {
  return readiness.size;
}

/**
 * Every line this file emits, and the only way it emits one. `log` is an injection seam
 * (`configureLifecycle({ logger })`), so an app's `Logger` decides whether a log call can throw —
 * and a throw here does not lose a line, it replaces the event. Inside `drain()` it rejected
 * `drainPromise`: `state` never reached 'stopped', the memo re-rejected for every later caller,
 * and on Bun the unhandled rejection ended the process the drain was trying to end cleanly.
 * Inside `readinessChecks()` it replaced the probe's answer with a throw.
 *
 * A lifecycle that cannot report is still a lifecycle: the line falls back to core's own
 * `rootLogger`, which is total by construction (`logger.ts`), and failing that is dropped.
 */
function report(level: 'info' | 'warn' | 'error', message: string, fields: LogFields): void {
  try {
    log[level](message, fields);
    return;
  } catch {
    // Fall through — the injected sink is gone, and the fallback below is the last one there is.
  }
  if (log === rootLogger) return;
  try {
    rootLogger[level](message, fields);
  } catch {
    // Both sinks are gone. Dropping the line is the only remaining option that still ends the
    // process, which is the outcome every caller of this file depends on.
  }
}

/**
 * Every check, run now, by name. A check that throws is `failing` — never an unhandled error.
 *
 * Built through `Object.fromEntries`, never by assigning `results[name]`: assignment to the one
 * name `__proto__` sets the PROTOTYPE instead of adding a key, so that check vanished from the
 * report, `ready` was computed over an empty object — vacuously true — and a failing check
 * answered 200. `fromEntries` defines own properties and has no such name.
 */
export function readinessChecks(): Readonly<Record<string, ReadinessStatus>> {
  const results: [string, ReadinessStatus][] = [];
  for (const [name, check] of readiness) {
    try {
      results.push([name, check() ? 'ok' : 'failing']);
    } catch (thrown) {
      results.push([name, 'failing']);
      report('warn', 'readiness check threw', { check: name, error: thrown });
    }
  }
  return Object.fromEntries(results);
}

export function inflightCount(): number {
  return inflight;
}

/** Test-only: drains still waiting on in-flight work. A count stuck above zero is a leak. */
export function idleWaiterCount(): number {
  return idleWaiters.length;
}

/**
 * Test-only: hooks still registered. A count that climbs across a start/stop cycle is a leak —
 * the registration retains its closure, and the next drain runs it against a torn-down resource.
 */
export function shutdownHookCount(): number {
  return registrations.length;
}

/** Register a drain hook. Returns an unregister function. */
export function onShutdown(
  name: string,
  hook: ShutdownHook,
  options?: OnShutdownOptions,
): () => void {
  const registration: Registration = { name, phase: options?.phase ?? 'close', hook };
  registrations.push(registration);
  return () => {
    registrations = registrations.filter((candidate) => candidate !== registration);
  };
}

/**
 * Mark a unit of work in flight. Call the returned function when it completes — drain waits
 * for the count to reach zero before closing resources.
 */
export function beginWork(): () => void {
  inflight += 1;
  let done = false;
  return () => {
    if (done) return;
    done = true;
    inflight -= 1;
    if (inflight === 0) {
      const waiters = idleWaiters;
      idleWaiters = [];
      for (const waiter of waiters) waiter();
    }
  };
}

/** True when new work must be refused — the HTTP layer answers 503 while this holds. */
export function isDraining(): boolean {
  return state === 'draining' || state === 'stopped';
}

function waitForIdle(timeoutMs: number): Promise<boolean> {
  if (inflight === 0) return Promise.resolve(true);
  return new Promise<boolean>((resolve) => {
    const waiter = (): void => {
      clearTimeout(timer);
      resolve(true);
    };
    // A drain that times out must not leave its waiter in the queue forever — the next
    // `beginWork()` to reach zero would still hold and invoke it, a dangling closure over a
    // promise nothing is awaiting anymore.
    const timer = setTimeout(() => {
      idleWaiters = idleWaiters.filter((candidate) => candidate !== waiter);
      resolve(false);
    }, timeoutMs);
    idleWaiters.push(waiter);
  });
}

/**
 * The budget every drain is bounded by — `DEFAULT_DEADLINE_MS` until an app raises it. There is no
 * unbounded state: `ShutdownReason.deadlineAt` was always computed and handed to every hook, so the
 * deadline was declared by the design all along and only the enforcement was missing.
 *
 * The ONE place the budget is decided, and exported so a test can pin it: 25s is far above any
 * drain a test can wait out, so the default needs a probe and not only a stopwatch.
 */
export function drainDeadlineMs(): number {
  return deadlineMs;
}

/**
 * What is left of that budget. Read per hook, not per phase: the deadline bounds the WHOLE drain,
 * so a hook that spent it leaves nothing for the ones behind it — which is what
 * `terminationGracePeriodSeconds` means, and what makes the SUM of the phases bounded rather than
 * each one of them separately. Returns `number`, never `number | undefined`: "no budget" is not a
 * state this file has, and the type is what keeps it from becoming one again.
 */
function remainingBudget(reason: ShutdownReason): number {
  return Math.max(0, reason.deadlineAt - systemClock.monotonic());
}

async function runPhase(phase: ShutdownPhase, reason: ShutdownReason): Promise<void> {
  for (const registration of registrations.filter((entry) => entry.phase === phase)) {
    const outcome = await settleWithin(() => registration.hook(reason), remainingBudget(reason));
    if (outcome.kind === 'failed') {
      report('error', 'shutdown hook failed', {
        hook: registration.name,
        phase,
        error: outcome.error,
      });
      continue;
    }
    if (outcome.kind === 'abandoned') {
      // Abandoned, not merely logged. A deadline that waited anyway would leave the kubelet to
      // SIGKILL this process — the every-deploy duplicate that draining exists to prevent — so
      // the drain moves on and the hook is left running with nobody reading it. The cost of that
      // choice is real and named in the cause: a write it had in flight may be half done.
      report('warn', 'X_SHUTDOWN_TIMEOUT', {
        code: 'X_SHUTDOWN_TIMEOUT',
        cause: `the "${registration.name}" shutdown hook (phase: ${phase}) was still running at the ${deadlineMs}ms drain deadline and has been ABANDONED — the process exits without it, so anything it had in flight may be incomplete`,
        fix: `raise the budget past the work this hook does — configureLifecycle({ deadlineMs: 600_000 }) for a 10-minute job — and set terminationGracePeriodSeconds to at least as many seconds, or make the "${registration.name}" hook return once it has stopped accepting work rather than once it has finished`,
        hook: registration.name,
        phase,
      });
    }
  }
}

/** The three phases, in order, under one budget. Never rejects — `drain()` depends on that. */
async function runDrain(signal: string, reason: ShutdownReason): Promise<void> {
  try {
    report('info', 'draining', { signal, deadlineMs, inflight });
    await runPhase('accept', reason);

    // Real monotonic, like `deadlineAt` itself: `waitForIdle` sleeps on a real `setTimeout`, and
    // a budget read off an injected clock is a number that timer will never honour.
    const remaining = Math.max(0, reason.deadlineAt - systemClock.monotonic());
    const idle = await waitForIdle(remaining);
    if (!idle) {
      report('warn', 'X_SHUTDOWN_TIMEOUT', {
        code: 'X_SHUTDOWN_TIMEOUT',
        cause: `${inflight} in-flight operations still running after ${deadlineMs}ms`,
        fix: 'raise the budget past the slowest handler — configureLifecycle({ deadlineMs: 600_000 }) for a 10-minute one — and set terminationGracePeriodSeconds to at least as many seconds, or shorten the handler',
      });
    }

    await runPhase('inflight', reason);
    await runPhase('close', reason);
  } catch (thrown) {
    // Nothing above should reach here — every hook is caught by `settleWithin` and every line
    // goes through `report`. If something does, the drain still ENDS: a rejected `drainPromise`
    // is a memo that re-rejects for every later caller and an unhandled rejection that kills the
    // process mid-drain, which is strictly worse than a drain that finished badly and said so.
    report('error', 'drain failed', { signal, error: thrown });
  } finally {
    state = 'stopped';
  }
  report('info', 'stopped', { signal });
}

/**
 * Idempotent: concurrent signals join the same drain, and so does a RE-ENTRANT one.
 *
 * The memo is published before `runDrain` is called, and that ordering is the whole of this
 * function. A hook may call back in here — `handle.stop()` in `@ultimat3/http` is `drain('manual')`
 * and an `accept` hook is exactly where a server stops listening — and `settleWithin` invokes a
 * hook SYNCHRONOUSLY, so the old `drainPromise = (async () => …)()` had not assigned yet when the
 * first hook ran: the re-entrant call saw `undefined`, started a second whole drain, and recursed
 * ~4,700 deep until the stack ran out, every level swallowed by `settleWithin` as
 * `shutdown hook failed`. Same rule as `packages/jobs/src/worker.ts` — guard and registration in
 * one synchronous step.
 */
export function drain(signal = 'manual'): Promise<void> {
  if (drainPromise !== undefined) return drainPromise;
  state = 'draining';
  const reason: ShutdownReason = { signal, deadlineAt: systemClock.monotonic() + deadlineMs };
  let published!: () => void;
  drainPromise = new Promise<void>((resolve) => {
    published = resolve;
  });
  // Both settle paths, for the reason `installSignalHandlers` gives below: `runDrain` cannot
  // reject today — that is its `try/finally`, not luck — and a rejected memo would re-reject for
  // every later caller and end the process the drain was trying to end cleanly.
  void runDrain(signal, reason).then(published, published);
  return drainPromise;
}

export interface SignalHandlerOptions {
  readonly signals?: readonly ProcessSignal[] | undefined;
  /** Call `process.exit()` once drained. Off in tests. */
  readonly exit?: boolean | undefined;
}

/** Install SIGTERM/SIGINT handling. Returns an uninstall function. */
export function installSignalHandlers(options?: SignalHandlerOptions): () => void {
  const signals: readonly ProcessSignal[] = options?.signals ?? ['SIGTERM', 'SIGINT'];
  const handlers = new Map<ProcessSignal, () => void>();

  for (const signal of signals) {
    const handler = (): void => {
      // Attached on BOTH settle paths, for the reason `settleWithin` gives: an unhandled rejection
      // ends the process before the drain does, and the exit is what the kubelet is waiting for.
      // `drain()` cannot reject today — that is the `try/finally` above, not luck — and this is
      // the one line that keeps it true when someone changes the body.
      const done = (): void => {
        if (options?.exit === true) process.exit(0);
      };
      void drain(signal).then(done, done);
    };
    handlers.set(signal, handler);
    process.on(signal, handler);
  }

  return () => {
    for (const [signal, handler] of handlers) process.off(signal, handler);
  };
}

export function healthReport(): HealthReport {
  const checks = readinessChecks();
  return {
    state,
    // `ready` is the same predicate `/readyz` answers on, so a body and its status can never
    // disagree — a 200 whose body says `ready: false` is the bug this shares one source to avoid.
    ready: state === 'ready' && Object.values(checks).every((status) => status === 'ok'),
    uptimeMs: Math.round(clock.monotonic() - startedAtMono),
    inflight,
    buildId: process.env['BUILD_ID'] ?? 'dev',
    checks,
    registered: readiness.size,
  };
}

/**
 * Liveness: the process exists and is not wedged. Stays 200 while draining, and deliberately
 * ignores the checks — a database outage that failed liveness everywhere would restart the whole
 * fleet into the same outage, with cold caches and no connections.
 */
export function healthzPayload(): HealthPayload {
  const body = healthReport();
  const ok = state !== 'stopped';
  return { ok, status: ok ? 200 : 503, body };
}

/** Readiness: may this instance receive traffic? 503 while starting, draining or any check fails. */
export function readyzPayload(): HealthPayload {
  const body = healthReport();
  return { ok: body.ready, status: body.ready ? 200 : 503, body };
}

/** Test-only: forget all hooks and return to `starting`. */
export function resetLifecycle(): void {
  deadlineMs = DEFAULT_DEADLINE_MS;
  clock = systemClock;
  log = rootLogger;
  state = 'starting';
  startedAtMono = clock.monotonic();
  inflight = 0;
  registrations = [];
  drainPromise = undefined;
  idleWaiters = [];
  readiness.clear();
}