[dev/run] add ability to register callback for clean up (#35119) (#35195)

* [dev/run] add ability to register callback for clean up

* remove old onExit option

* update docs to mention `addCleanupTask`

* update yarn.lock

package.json

@@ -356,6 +356,7 @@
     "eslint-plugin-prettier": "^3.0.1",
     "eslint-plugin-react": "^7.12.4",
     "eslint-plugin-react-hooks": "^1.6.0",
+    "exit-hook": "^2.1.0",
     "faker": "1.1.0",
     "fetch-mock": "7.3.0",
     "geckodriver": "^1.16.1",

src/dev/run/README.md

@@ -65,7 +65,7 @@ $ node scripts/my_task
 
 ## API
 
-- ***`run(fn: async ({ flags: Flags, log: ToolingLog }) => Promise<void>, options: Options)`***
+- ***`run(fn: async ({ flags: Flags, log: ToolingLog, addCleanupTask }) => Promise<void>, options: Options)`***
   
     Execte an async function, passing it the parsed flags and a tooling log that is configured to the requested logging level. If the returned promise is rejected with an error created by `createFailError(...)` or `createFlagError(...)` the process will exit as described by the error, otherwise the process will exit with code 1.
     
@@ -78,6 +78,14 @@ $ node scripts/my_task
 
       The parsed CLI flags, created by [`getopts`](https://www.npmjs.com/package/getopts). Includes the default flags for controlling the log level of the ToolingLog, and `flags.unexpected`, which is an array of flag names which were passed but not expected.
 
+    - *`addCleanupTask: (task: CleanupTask) => void`*:
+
+      A function that registers a task to be called __once__ as soon as one of the following occurs:
+
+      1. `fn` resolve/rejects
+      2. something calls `process.exit()`
+      3. the user hits <kbd>ctrl</kbd>+<kbd>c</kbd> in their terminal causing the SIGINT signal to be sent to the process
+
     **`Options`:**
     - *`usage: string`*
 

src/dev/run/run.ts

@@ -17,11 +17,17 @@
  * under the License.
  */
 
+// @ts-ignore @types are outdated and module is super simple
+import exitHook from 'exit-hook';
+
 import { pickLevelFromFlags, ToolingLog } from '@kbn/dev-utils';
 import { createFlagError, isFailError } from './fail';
 import { Flags, getFlags, getHelp } from './flags';
 
-type RunFn = (args: { log: ToolingLog; flags: Flags }) => Promise<void> | void;
+type CleanupTask = () => void;
+type RunFn = (
+  args: { log: ToolingLog; flags: Flags; addCleanupTask: (task: CleanupTask) => void }
+) => Promise<void> | void;
 
 export interface Options {
   usage?: string;
@@ -50,13 +56,13 @@ export async function run(fn: RunFn, options: Options = {}) {
     writeTo: process.stdout,
   });
 
-  try {
-    if (!allowUnexpected && flags.unexpected.length) {
-      throw createFlagError(`Unknown flag(s) "${flags.unexpected.join('", "')}"`);
-    }
+  process.on('unhandledRejection', error => {
+    log.error('UNHANDLED PROMISE REJECTION');
+    log.error(error);
+    process.exit(1);
+  });
 
-    await fn({ log, flags });
-  } catch (error) {
+  const handleErrorWithoutExit = (error: any) => {
     if (isFailError(error)) {
       log.error(error.message);
 
@@ -64,11 +70,46 @@ export async function run(fn: RunFn, options: Options = {}) {
         log.write(getHelp(options));
       }
 
-      process.exit(error.exitCode);
+      process.exitCode = error.exitCode;
     } else {
       log.error('UNHANDLED ERROR');
       log.error(error);
-      process.exit(1);
+      process.exitCode = 1;
     }
+  };
+
+  const doCleanup = () => {
+    const tasks = cleanupTasks.slice(0);
+    cleanupTasks.length = 0;
+
+    for (const task of tasks) {
+      try {
+        task();
+      } catch (error) {
+        handleErrorWithoutExit(error);
+      }
+    }
+  };
+
+  const unhookExit: CleanupTask = exitHook(doCleanup);
+  const cleanupTasks: CleanupTask[] = [unhookExit];
+
+  try {
+    if (!allowUnexpected && flags.unexpected.length) {
+      throw createFlagError(`Unknown flag(s) "${flags.unexpected.join('", "')}"`);
+    }
+
+    try {
+      await fn({
+        log,
+        flags,
+        addCleanupTask: (task: CleanupTask) => cleanupTasks.push(task),
+      });
+    } finally {
+      doCleanup();
+    }
+  } catch (error) {
+    handleErrorWithoutExit(error);
+    process.exit();
   }
 }

yarn.lock

@@ -10417,6 +10417,11 @@ exit-hook@^1.0.0:
   resolved "https://registry.yarnpkg.com/exit-hook/-/exit-hook-1.1.1.tgz#f05ca233b48c05d54fff07765df8507e95c02ff8"
   integrity sha1-8FyiM7SMBdVP/wd2XfhQfpXAL/g=
 
+exit-hook@^2.1.0:
+  version "2.1.0"
+  resolved "https://registry.yarnpkg.com/exit-hook/-/exit-hook-2.1.0.tgz#2be08d8d01220050878577bfa017e104a6c3bcf3"
+  integrity sha512-JCSm9Znc/KW6hoKYHOIqLxM2Z88+AQcabo07rJHZSyXcQIq6HsXkSWRVZRp13RFkGVIDcz1DRIbKR5cnU1uzCA==
+
 exit@^0.1.2, exit@~0.1.1:
   version "0.1.2"
   resolved "https://registry.yarnpkg.com/exit/-/exit-0.1.2.tgz#0632638f8d877cc82107d30a0fff1a17cba1cd0c"