From 4aaf845fca5990ec2f48df59accbbd7e174fc457 Mon Sep 17 00:00:00 2001 From: Shy <48560522+shyguy1412@users.noreply.github.com> Date: Wed, 6 Mar 2024 01:22:45 +0100 Subject: [PATCH] API: make ns.atExit add the callback to an array instead of setting it (#1059) --- markdown/bitburner.ns.atexit.md | 3 ++- markdown/bitburner.ns.md | 2 +- src/Netscript/WorkerScript.ts | 4 ++-- src/Netscript/killWorkerScript.ts | 21 ++++++++++++-------- src/NetscriptFunctions.ts | 23 ++++++++++++++-------- src/ScriptEditor/NetscriptDefinitions.d.ts | 2 +- test/jest/Netscript/RunScript.test.ts | 5 ++++- 7 files changed, 38 insertions(+), 22 deletions(-) diff --git a/markdown/bitburner.ns.atexit.md b/markdown/bitburner.ns.atexit.md index 21d5cd077..f6636dc5e 100644 --- a/markdown/bitburner.ns.atexit.md +++ b/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:** diff --git a/markdown/bitburner.ns.md b/markdown/bitburner.ns.md index 0865dffff..6c52c9494 100644 --- a/markdown/bitburner.ns.md +++ b/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. | diff --git a/src/Netscript/WorkerScript.ts b/src/Netscript/WorkerScript.ts index 3a9e207d7..28f6c0c3e 100644 --- a/src/Netscript/WorkerScript.ts +++ b/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 void> = new Map(); constructor(runningScriptObj: RunningScript, pid: number, nsFuncsGenerator?: (ws: WorkerScript) => NSFull) { this.name = runningScriptObj.filename; diff --git a/src/Netscript/killWorkerScript.ts b/src/Netscript/killWorkerScript.ts index 46cc84b09..57e9bb84c 100644 --- a/src/Netscript/killWorkerScript.ts +++ b/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); } diff --git a/src/NetscriptFunctions.ts b/src/NetscriptFunctions.ts index c46d24253..316c51764 100644 --- a/src/NetscriptFunctions.ts +++ b/src/NetscriptFunctions.ts @@ -1708,14 +1708,21 @@ export const ns: InternalAPI = { 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); diff --git a/src/ScriptEditor/NetscriptDefinitions.d.ts b/src/ScriptEditor/NetscriptDefinitions.d.ts index 3ebca2a3d..875fd064c 100644 --- a/src/ScriptEditor/NetscriptDefinitions.d.ts +++ b/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. diff --git a/test/jest/Netscript/RunScript.test.ts b/test/jest/Netscript/RunScript.test.ts index 2c1b72c72..7e520cf7f 100644 --- a/test/jest/Netscript/RunScript.test.ts +++ b/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((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);