Proper graceful shutdown without dropping requests

[realfresh]

I've finally done it, graceful shutdown that ensures all in-flight requests are completed, connections closed properly, then finally exiting.

This is specific to Kubernetes, and the code below is taken from our codebase, some the imports may not make sense. But I'm sure you will be able to adapt it to your project correctly.

import {type Context, type Hono} from "hono";
import type {BlankSchema, Schema} from "hono/types";
import {EventEmitter} from "node:events";
import {type IncomingMessage, type Server} from "node:http";
import {type Socket} from "node:net";
import {sleep} from "...";
import {MultiSet} from "...";
import {type DurationInput, ms} from "...";
import {serve} from "@hono/node-server";
import {hlog} from "./hlog.js";
import {type HonoNodeEnv} from "./HonoAppBase.js";
import {HonoShutdownManager} from "./shutdown.js";

const ENABLE_TRACKER_LOGGING =
  process.env.HONO_SHUTDOWN_TRACKER_LOGGING === "true" ||
  process.env.HONO_SHUTDOWN_GRACEFUL_TEST === "true";

//~ Utility Functions
const onShutdownError = (err?: unknown) => {
  hlog.shut.error(err);
  process.exit(1);
};

const onShutdownTimeout = () => {
  hlog.shut.warn("Shutdown grace period elapsed, exiting now...");
  process.exit(1);
};

const onShutdownSuccess = () => {
  hlog.shut.info("Shutdown complete, exiting...");
  process.exit(0);
};

const closeHttpServer = (server: Server) => {
  return new Promise<void>((resolve) => {
    server.close((err) => {
      if (err) hlog.shut.error(err);
      hlog.shut.info(`-> server closed`);
      resolve();
    });
    server.closeIdleConnections();
  });
};

const executeShutdownHooks = async () => {
  const {hooks} = HonoShutdownManager;
  if (hooks.length > 0) {
    hlog.shut.info(`Executing shutdown hooks ${hooks.length}`);
    const promises = hooks.map((hook) => hook());
    await Promise.all(promises);
  }
};

//~ Utility Classes
abstract class Tracker extends EventEmitter {
  isShuttingDown = false;

  constructor() {
    super();
    HonoShutdownManager.onShutdown(() => {
      this.isShuttingDown = true;
    });
  }

  protected debug = (msg: string) => {
    if (ENABLE_TRACKER_LOGGING || this.isShuttingDown) {
      hlog.shut.debug(msg);
    }
  };
}

class ConnectionTracker extends Tracker {
  entries = new Set<Socket>();

  track = (req: IncomingMessage) => {
    const conn = req.socket;

    if (this.entries.has(conn)) {
      return;
    }

    this.entries.add(conn);
    this.debug(`++conn [${this.entries.size}]`);

    conn.on("close", () => {
      this.entries.delete(conn);
      this.debug(`--conn [${this.entries.size}]`);
    });
  };

  destroy = async () => {
    if (this.entries.size > 0) {
      hlog.shut.info(`-> destroying sockets: (${this.entries.size})`);
    }
    this.entries.forEach((conn) => conn.destroy());
  };
}

class RequestTracker extends Tracker {
  active = 0;
  paths = new MultiSet<string>();

  increment = (c: Context) => {
    this.active++;
    this.paths.add(c.req.path);
    this.debug(`++req (${this.active})`);
  };

  decrement = (c: Context) => {
    this.active--;
    this.paths.remove(c.req.path);
    this.debug(`--req (${this.active})`);

    if (this.active === 0) {
      this.emit("complete");
    }
  };

  /** Wait for all ongoing requests to finish */
  onComplete = async () => {
    if (this.active > 0) {
      await Promise.race([
        // The event allows for a fast exit if all requests are completed,
        // but the interval is a backup mechanism to avoid race conditions
        this.onCompleteEvent(),
        this.onCompleteInterval(),
      ]);
    }

    hlog.shut.info(`-> outstanding requests completed`);
  };

  private onCompleteEvent = () => {
    return new Promise((resolve) => this.once("complete", resolve));
  };

  private onCompleteInterval = async () => {
    while (this.active > 0) {
      this.logPaths();
      hlog.shut.info(`-> requests remaining (${this.active})`);
      await sleep(1000);
    }
  };

  private logPaths = () => {
    const paths: Record<string, number> = {};

    this.paths.forEachMultiplicity((count, path) => {
      paths[path] = count;
    });

    hlog.shut.info("-> paths in-flight", {paths});
  };
}

//~ Hono Runner For Graceful Shutdown
export type HonoRunnerParams = {
  port: number;
  onStart: () => void;
  onShutdown: () => void;
  /**
   * Determines if the server should gracefully handle shutdown.
   * If `false`, the server will call `process.exit(0)` immediately on SIGTERM.
   * That would effectively disable the entire logic in this module.
   * Useful in development for faster restarts / reloads.
   */
  shutdownGraceful: boolean;
  /**
   * The maximum amount of time to perform the shutdown process before timing out
   * and calling `process.exit(1)`. Should match `terminationGracePeriodSeconds`
   */
  shutdownTimeout: DurationInput;
  /**
   * Duration to delay invoking `sever.close()` which stops it from accepting new
   * connections. This value should equal the time it will take for the K8 readiness
   * probe to fail. This means K8 will stop forwarding traffic to this server.
   * Delaying for this period ensures that new connections forwarded by K8
   * are still handled instead of returning an error. The goal is to only close
   * the server when there are no more incoming requests.
   */
  shutdownServerCloseDelay: DurationInput;
};

export function HonoRunner<
  E extends HonoNodeEnv<{}>,
  S extends Schema = BlankSchema,
  BP extends string = "/",
>(app: Hono<E, S, BP>, params: HonoRunnerParams) {
  const state = {shutdown: false};
  const requests = new RequestTracker();
  const connections = new ConnectionTracker();

  const shutdownTimeout = ms(params.shutdownTimeout);

  // Add a 3-second buffer to given K8 time to mark the pod as not ready
  const shutdownCloseDelay = ms(params.shutdownServerCloseDelay) + ms("3s");

  // This gets added after health-check handlers, which means it won't track
  // health-check requests, which is ideal since we don't care if those
  // are completed fully during the shutdown
  app.use("*", async (c, next) => {
    const {incoming} = c.env;
    connections.track(incoming);
    requests.increment(c);

    // On Shutdown:
    if (state.shutdown) {
      c.header("connection", "close");
    }

    // Process the HTTP request:
    // This apparently doesn't throw any errors
    // so omit the try/catch to avoid performance overhead
    await next();

    // On Shutdown
    if (state.shutdown) {
      // 1. Close keep-alive connections
      // Setting this header informs keep-alive connections to disconnect after response
      // c.res.headers.set("connection", "close");
      c.header("connection", "close");

      // 2. Destroy the socket
      // Use `destroySoon()` which is more graceful than `destroy()` by
      // allowing the response to be sent first before closing the socket
      c.env.incoming.socket.destroySoon();
    }

    requests.decrement(c);
  });

  const handleShutdown = async (server: Server) => {
    hlog.shut.info(`Configure shutdown timeout`, {shutdownTimeout});
    setTimeout(onShutdownTimeout, shutdownTimeout);

    hlog.shut.info(`Wait for readiness probe to fail...`, {
      shutdownCloseDelay,
    });
    await sleep(shutdownCloseDelay);

    // Set the shutdown flag:
    // -> start sending 'connection: close' headers
    // -> destroy sockets after requests are complete
    // Do this only after the readiness probe has failed to ensure no new connections
    state.shutdown = true;

    hlog.shut.info(`Complete in-flight requests...`);
    await requests.onComplete();

    // This method uses `destroy()`, not `destroySoon()`, which is more aggressive
    // but in this case, it is fine since all the requests are completed
    hlog.shut.info(`Destroy remaining sockets...`);
    await connections.destroy();

    hlog.shut.info(`Close server and execute shutdown hooks...`);
    await Promise.all([closeHttpServer(server), executeShutdownHooks()]);

    onShutdownSuccess();
  };

  const handleSigterm = (server: Server) => {
    try {
      hlog.shut.info(`SIGTERM <- start shutdown`);
      params.onShutdown();

      if (params.shutdownGraceful) {
        void handleShutdown(server);
      } else {
        onShutdownSuccess();
      }
    } catch (err) {
      onShutdownError(err);
    }
  };

  const start = () => {
    const server = serve({
      fetch: app.fetch,
      port: params.port,
    }) as Server;

    hlog.server.info(`Server started on port :${params.port}`);

    HonoShutdownManager.catchShutdownSignal(() => {
      handleSigterm(server);
    });

    params.onStart();

    return server;
  };

  const addShutdownHook = HonoShutdownManager.addHook;

  return {start, addShutdownHook};
}