From 34db6e8b26f82f24e66264ad41720765714a3e41 Mon Sep 17 00:00:00 2001 From: David Walker Date: Fri, 16 Aug 2024 19:10:20 -0700 Subject: [PATCH] NETSCRIPT: Rework disableLog for efficiency (#1589) The current implementation was naive; disableLog("ALL") was storing a key for every function, and iterating over a different object to do it (when iterating over objects is quite slow). The common cases of Bitburner (and especially batching, where efficiency matters most) are either never disabling anything, or disabling "ALL". This optimizes for these two cases, at the expense of slightly more complicated code to deal with the less-common edge cases. --- src/Netscript/WorkerScript.ts | 2 +- src/NetscriptFunctions.ts | 47 +++++++++----- test/jest/Netscript/DisableLog.test.ts | 87 ++++++++++++++++++++++++++ 3 files changed, 119 insertions(+), 17 deletions(-) create mode 100644 test/jest/Netscript/DisableLog.test.ts diff --git a/src/Netscript/WorkerScript.ts b/src/Netscript/WorkerScript.ts index afd1f55e1..772dd4f96 100644 --- a/src/Netscript/WorkerScript.ts +++ b/src/Netscript/WorkerScript.ts @@ -131,7 +131,7 @@ export class WorkerScript { } shouldLog(fn: string): boolean { - return this.disableLogs[fn] == null; + return !(this.disableLogs.ALL || this.disableLogs[fn]); } log(func: string, txt: () => string): void { diff --git a/src/NetscriptFunctions.ts b/src/NetscriptFunctions.ts index d44b68520..7ec2da28b 100644 --- a/src/NetscriptFunctions.ts +++ b/src/NetscriptFunctions.ts @@ -495,37 +495,46 @@ export const ns: InternalAPI = { }, disableLog: (ctx) => (_fn) => { const fn = helpers.string(ctx, "fn", _fn); + if (possibleLogs[fn] === undefined) { + throw helpers.errorMessage(ctx, `Invalid argument: ${fn}.`); + } if (fn === "ALL") { - for (const fn of Object.keys(possibleLogs)) { + ctx.workerScript.disableLogs = allDisabled; + // No need to log here, it's been disabled. + } else { + // We don't track individual log entries when all are disabled. + if (!ctx.workerScript.disableLogs["ALL"]) { ctx.workerScript.disableLogs[fn] = true; + helpers.log(ctx, () => `Disabled logging for ${fn}`); } - helpers.log(ctx, () => `Disabled logging for all functions`); - } else if (possibleLogs[fn] === undefined) { - throw helpers.errorMessage(ctx, `Invalid argument: ${fn}.`); - } else { - ctx.workerScript.disableLogs[fn] = true; - helpers.log(ctx, () => `Disabled logging for ${fn}`); } }, enableLog: (ctx) => (_fn) => { const fn = helpers.string(ctx, "fn", _fn); + if (possibleLogs[fn] === undefined) { + throw helpers.errorMessage(ctx, `Invalid argument: ${fn}.`); + } if (fn === "ALL") { - for (const fn of Object.keys(possibleLogs)) { - delete ctx.workerScript.disableLogs[fn]; - } + ctx.workerScript.disableLogs = {}; helpers.log(ctx, () => `Enabled logging for all functions`); - } else if (possibleLogs[fn] === undefined) { - throw helpers.errorMessage(ctx, `Invalid argument: ${fn}.`); + } else { + if (ctx.workerScript.disableLogs["ALL"]) { + // As an optimization, we normally store only that key, but we have to + // expand it out to all keys at this point. + // Conveniently, possibleLogs serves as a model for "all keys disabled." + ctx.workerScript.disableLogs = Object.assign({}, possibleLogs, { ALL: false, [fn]: false }); + } else { + ctx.workerScript.disableLogs[fn] = false; + } + helpers.log(ctx, () => `Enabled logging for ${fn}`); } - delete ctx.workerScript.disableLogs[fn]; - helpers.log(ctx, () => `Enabled logging for ${fn}`); }, isLogEnabled: (ctx) => (_fn) => { const fn = helpers.string(ctx, "fn", _fn); if (possibleLogs[fn] === undefined) { throw helpers.errorMessage(ctx, `Invalid argument: ${fn}.`); } - return !ctx.workerScript.disableLogs[fn]; + return ctx.workerScript.shouldLog(fn); }, getScriptLogs: (ctx) => @@ -1840,7 +1849,13 @@ export function NetscriptFunctions(ws: WorkerScript): NSFull { return NSProxy(ws, ns, [], { args: ws.args.slice(), pid: ws.pid, enums }); } -const possibleLogs = Object.fromEntries([...getFunctionNames(ns, "")].map((a) => [a, true])); +const possibleLogs = Object.fromEntries(getFunctionNames(ns, "").map((a) => [a, true])); +possibleLogs.ALL = true; + +// We reuse this object for *all* scripts that disable all keys, to prevent memory growth. +// Any script that needs a custom set of values will use a fresh object. +const allDisabled = { ALL: true } as const; + /** Provides an array of all function names on a nested object */ function getFunctionNames(obj: object, prefix: string): string[] { const functionNames: string[] = []; diff --git a/test/jest/Netscript/DisableLog.test.ts b/test/jest/Netscript/DisableLog.test.ts new file mode 100644 index 000000000..c7ae21acf --- /dev/null +++ b/test/jest/Netscript/DisableLog.test.ts @@ -0,0 +1,87 @@ +import type { Script } from "../../../src/Script/Script"; +import type { ScriptFilePath } from "../../../src/Paths/ScriptFilePath"; +import { FormatsNeedToChange } from "../../../src/ui/formatNumber"; +import { Server } from "../../../src/Server/Server"; +import { RunningScript } from "../../../src/Script/RunningScript"; +import { CompletedProgramName } from "../../../src/Programs/Enums"; +import { WorkerScript } from "../../../src/Netscript/WorkerScript"; +import { NetscriptFunctions } from "../../../src/NetscriptFunctions"; +import { AddToAllServers, DeleteServer } from "../../../src/Server/AllServers"; + +test("Edge cases of disableLog", function () { + // Ensure that formatting functions work properly + FormatsNeedToChange.emit(); + let server; + try { + server = new Server({ hostname: "home", adminRights: true, maxRam: 8 }); + server.programs.push(CompletedProgramName.bruteSsh, CompletedProgramName.ftpCrack); + AddToAllServers(server); + // We don't need this script to be runnable, it just needs to exist so that + // we can create a RunningScript and WorkerScript object. + expect(server.writeToScriptFile("test.js" as ScriptFilePath, "")).toEqual({ overwritten: false }); + const script = server.scripts.get("test.js" as ScriptFilePath) as Script; + + const runningScript = new RunningScript(script, 2); + const ws = new WorkerScript(runningScript, 1, NetscriptFunctions); + + const ns = ws.env.vars; + + // Generate logs in a specific pattern that checks edge cases in + // disableLog. We want to check various combinations of things that + // are/aren't disabled, as well as previously disabled. + ns.brutessh("home"); + ns.ftpcrack("home"); + ns.print("before disableLog ALL"); + + expect(() => ns.disableLog("all")).toThrow("Invalid argument: all."); + ns.disableLog("ALL"); + + ns.brutessh("home"); + ns.ftpcrack("home"); + ns.print("after disableLog ALL"); + + ns.enableLog("brutessh"); + ns.brutessh("home"); + ns.ftpcrack("home"); + ns.print("after enableLog brutessh"); + + ns.disableLog("brutessh"); + ns.enableLog("ftpcrack"); + ns.brutessh("home"); + ns.ftpcrack("home"); + ns.print("after enableLog ftpcrack"); + + ns.enableLog("ftpcrack"); + ns.brutessh("home"); + ns.ftpcrack("home"); + + ns.print("after redundant enable"); + ns.disableLog("ALL"); + ns.brutessh("home"); + ns.ftpcrack("home"); + ns.enableLog("ALL"); + ns.brutessh("home"); + ns.ftpcrack("home"); + ns.print("end"); + + expect(runningScript.logs).toEqual([ + "brutessh: Executed BruteSSH.exe on 'home' to open SSH port (22).", + "ftpcrack: Executed FTPCrack.exe on 'home' to open FTP port (21).", + "before disableLog ALL", + "disableLog: Invalid argument: all.", + "after disableLog ALL", + "brutessh: SSH Port (22) already opened on 'home'.", + "after enableLog brutessh", + "ftpcrack: FTP Port (21) already opened on 'home'.", + "after enableLog ftpcrack", + "ftpcrack: FTP Port (21) already opened on 'home'.", + "after redundant enable", + "enableLog: Enabled logging for all functions", + "brutessh: SSH Port (22) already opened on 'home'.", + "ftpcrack: FTP Port (21) already opened on 'home'.", + "end", + ]); + } finally { + DeleteServer(server.hostname); + } +});