API: make ns.atExit add the callback to an array instead of setting it (#1059)

markdown/bitburner.ns.atexit.md

@@ -9,7 +9,7 @@ Add callback function when the script dies
 **Signature:**
 
 ```typescript
-atExit(f: () => void): void;
+atExit(f: () => void, id?: string): void;
 ```
 
 ## Parameters
@@ -17,6 +17,7 @@ atExit(f: () => void): void;
 |  Parameter | Type | Description |
 |  --- | --- | --- |
 |  f | () => void |  |
+|  id | string | _(Optional)_ |
 
 **Returns:**
 

markdown/bitburner.ns.md

@@ -56,7 +56,7 @@ export async function main(ns) {
 |  --- | --- |
 |  [alert(msg)](./bitburner.ns.alert.md) | Open up a message box. |
 |  [asleep(millis)](./bitburner.ns.asleep.md) | Suspends the script for n milliseconds. Doesn't block with concurrent calls. |
-|  [atExit(f)](./bitburner.ns.atexit.md) | Add callback function when the script dies |
+|  [atExit(f, id)](./bitburner.ns.atexit.md) | Add callback function when the script dies |
 |  [brutessh(host)](./bitburner.ns.brutessh.md) | Runs BruteSSH.exe on a server. |
 |  [clear(handle)](./bitburner.ns.clear.md) | Clear data from a file. |
 |  [clearLog()](./bitburner.ns.clearlog.md) | Clears the script’s logs. |

src/Netscript/WorkerScript.ts

@@ -78,8 +78,8 @@ export class WorkerScript {
   /** hostname on which this script is running */
   hostname: string;
 
-  /** Function called when the script ends. */
-  atExit: (() => void) | undefined = undefined;
+  /**Map of functions called when the script ends. */
+  atExit: Map<string, () => void> = new Map();
 
   constructor(runningScriptObj: RunningScript, pid: number, nsFuncsGenerator?: (ws: WorkerScript) => NSFull) {
     this.name = runningScriptObj.filename;

src/Netscript/killWorkerScript.ts

@@ -42,20 +42,25 @@ function stopAndCleanUpWorkerScript(ws: WorkerScript): void {
   if (ws.delay) clearTimeout(ws.delay);
   ws.delayReject?.(new ScriptDeath(ws));
   ws.env.runningFn = "";
+  const atExit = ws.atExit;
+  //Calling ns.exit inside ns.atExit can lead to recursion
+  //so the map must be cleared before looping
+  ws.atExit = new Map();
 
-  if (typeof ws.atExit === "function") {
+  for (const key of atExit.keys()) {
     try {
-      const atExit = ws.atExit;
-      ws.atExit = undefined;
-      atExit();
+      const callback = atExit.get(key);
+      if (typeof callback == "function") callback();
     } catch (e: unknown) {
       handleUnknownError(e, ws, "Error running atExit function.\n\n");
     }
-    if (ws.env.stopFlag) {
-      // If atExit() kills the script, we'll already be stopped, don't stop again.
-      return;
-    }
   }
+
+  if (ws.env.stopFlag) {
+    // If atExit() kills the script, we'll already be stopped, don't stop again.
+    return;
+  }
+
   ws.env.stopFlag = true;
   removeWorkerScript(ws);
 }

src/NetscriptFunctions.ts

@@ -1708,14 +1708,21 @@ export const ns: InternalAPI<NSFull> = {
     sinceInstall: Object.assign({}, Player.moneySourceA),
     sinceStart: Object.assign({}, Player.moneySourceB),
   }),
-  atExit: (ctx) => (f) => {
-    if (typeof f !== "function") {
-      throw helpers.errorMessage(ctx, "argument should be function");
-    }
-    ctx.workerScript.atExit = () => {
-      f();
-    }; // Wrap the user function to prevent WorkerScript leaking as 'this'
-  },
+  atExit:
+    (ctx) =>
+    (f, id = "default") => {
+      if (typeof f !== "function") {
+        throw helpers.errorMessage(ctx, "argument should be function");
+      }
+
+      if (typeof id !== "string") {
+        throw helpers.errorMessage(ctx, "id should be a string");
+      }
+
+      ctx.workerScript.atExit.set(id, () => {
+        f();
+      }); // Wrap the user function to prevent WorkerScript leaking as 'this'
+    },
   mv: (ctx) => (_host, _source, _destination) => {
     const hostname = helpers.string(ctx, "host", _host);
     const server = helpers.getServer(ctx, hostname);

src/ScriptEditor/NetscriptDefinitions.d.ts

@@ -7246,7 +7246,7 @@ export interface NS {
    *
    * Add callback to be executed when the script dies.
    */
-  atExit(f: () => void): void;
+  atExit(f: () => void, id?: string): void;
 
   /**
    * Move a file on the target server.

test/jest/Netscript/RunScript.test.ts

@@ -105,7 +105,10 @@ test.each([
     // await script death.
     const ws = workerScripts.get(pid);
     expect(ws).toBeDefined();
-    const result = await Promise.race([alerted, new Promise((resolve) => (ws.atExit = resolve))]);
+    const result = await Promise.race([
+      alerted,
+      new Promise<void>((resolve) => (ws!.atExit = new Map([["default", resolve]]))),
+    ]);
     // If an error alert was thrown, we catch it here.
     expect(result).not.toBeDefined();
     expect(runningScript.logs).toEqual(expectedLog);