fix(cli): handle shutdown events gracefully (#6051)

Previously, the `wing it` didn't wait for the Console to close and release its resources.

apps/wing/package.json

@@ -47,6 +47,7 @@
     "dotenv-expand": "^10.0.0",
     "glob": "^10.3.10",
     "inquirer": "^8",
+    "lodash.once": "^4.1.1",
     "nanoid": "^3.3.6",
     "npm-packlist": "^8.0.0",
     "open": "^8.4.2",
@@ -59,6 +60,7 @@
   "devDependencies": {
     "@types/debug": "^4.1.8",
     "@types/inquirer": "^9.0.7",
+    "@types/lodash.once": "^4.1.9",
     "@types/node": "^20.11.0",
     "@types/node-persist": "^3.1.4",
     "@types/npm-packlist": "^7.0.1",

apps/wing/src/commands/run.ts

@@ -4,7 +4,9 @@ import { createConsoleApp } from "@wingconsole/app";
 import { BuiltinPlatform } from "@winglang/compiler";
 import { debug } from "debug";
 import { glob } from "glob";
+import once from "lodash.once";
 import { parseNumericString } from "../util";
+import { beforeShutdown } from "../util.before-shutdown.js";
 
 /**
  * Options for the `run` command.
@@ -77,7 +79,7 @@ export async function run(entrypoint?: string, options?: RunOptions) {
     requestedPort,
     hostUtils: {
       async openExternal(url: string) {
-        await open(url);
+        open(url);
       },
     },
     platform: options?.platform,
@@ -87,11 +89,9 @@ export async function run(entrypoint?: string, options?: RunOptions) {
   const url = `http://localhost:${port}/`;
   console.log(`The Wing Console is running at ${url}`);
 
-  const onExit = async (exitCode: number) => {
+  const onExit = once(async (exitCode: number) => {
     await close(() => process.exit(exitCode));
-  };
+  });
 
-  process.once("exit", (c) => void onExit(c));
-  process.once("SIGTERM", () => void onExit(0));
-  process.once("SIGINT", () => void onExit(0));
+  beforeShutdown(() => onExit(0));
 }

apps/wing/src/util.before-shutdown.ts

@@ -0,0 +1,79 @@
+/**
+ * Based on https://gist.github.com/nfantone/1eaa803772025df69d07f4dbf5df7e58.
+ */
+type BeforeShutdownListener = (codeOrSignal: string | number) => Promise<void> | void;
+
+/**
+ * Time in milliseconds to wait before forcing shutdown.
+ */
+const SHUTDOWN_TIMEOUT = 15_000;
+
+/**
+ * A queue of listener callbacks to execute before shutting
+ * down the process.
+ */
+const shutdownListeners: BeforeShutdownListener[] = [];
+
+/**
+ * Listen for signals and execute given `fn` function once.
+ * @param  fn Function to execute on shutdown.
+ */
+const processOnce = (fn: BeforeShutdownListener) => {
+  process.once("SIGINT", (signal) => void fn(signal));
+  process.once("SIGTERM", (signal) => void fn(signal));
+  process.once("exit", (code) => void fn(code));
+  process.once("beforeExit", (code) => void fn(code));
+};
+
+/**
+ * Sets a forced shutdown mechanism that will exit the process after `timeout` milliseconds.
+ * @param timeout Time to wait before forcing shutdown (milliseconds)
+ */
+const forceExitAfter = (timeout: number) => () => {
+  setTimeout(() => {
+    // Force shutdown after timeout
+    console.warn(`Could not close resources gracefully after ${timeout}ms: forcing shutdown`);
+    return process.exit(1);
+  }, timeout).unref();
+};
+
+/**
+ * Main process shutdown handler. Will invoke every previously registered async shutdown listener
+ * in the queue and exit with a code of `0`. Any `Promise` rejections from any listener will
+ * be logged out as a warning, but won't prevent other callbacks from executing.
+ */
+async function shutdownHandler(codeOrSignal: string | number) {
+  for (const listener of shutdownListeners) {
+    try {
+      await listener(codeOrSignal);
+    } catch (error) {
+      console.warn(
+        `A shutdown handler failed before completing with: ${
+          error instanceof Error ? error.message : error
+        }`
+      );
+    }
+  }
+
+  return process.exit(0);
+}
+
+/**
+ * Registers a new shutdown listener to be invoked before exiting
+ * the main process. Listener handlers are guaranteed to be called in the order
+ * they were registered.
+ * @param listener The shutdown listener to register.
+ * @returns Echoes back the supplied `listener`.
+ */
+export function beforeShutdown(listener: BeforeShutdownListener) {
+  shutdownListeners.push(listener);
+  return listener;
+}
+
+// Register shutdown callback that kills the process after `SHUTDOWN_TIMEOUT` milliseconds
+// This prevents custom shutdown handlers from hanging the process indefinitely
+processOnce(forceExitAfter(SHUTDOWN_TIMEOUT));
+
+// Register process shutdown callback
+// Will listen to incoming signal events and execute all registered handlers in the stack
+processOnce(shutdownHandler);