Only recreate instance after returning panic in wasm demo

Author: ictrobot

crates/aoc_wasm/web/aoc.mjs

@@ -36,8 +36,10 @@ export class Aoc {
     #instance;
     /** @type {WebAssembly.Memory} */
     #memory;
-    /** @type {Worker[]} */
+    /** @type {Worker[]|undefined} */
     #workers;
+    /** @type {number|undefined} */
+    #workerCount;
 
     /**
      * @param {WebAssembly.Module} module
@@ -101,38 +103,33 @@ export class Aoc {
 
     /**
      * @param {WebAssembly.Module} module
-     * @param {WebAssembly.Instance} [instance]
+     * @param {{instance?: WebAssembly.Instance, workerCount?: number}} [options]
      */
-    constructor(module, instance) {
+    constructor(module, options) {
+        options ??= {};
         const imports = WebAssembly.Module.imports(module);
         if (imports.length === 0) {
             this.#multithreaded = false;
             this.#module = module;
-            this.#instance = instance ?? new WebAssembly.Instance(module);
-            this.#memory = this.#exports.memory;
+            this.#instance = options.instance;
+            if (this.#instance) this.#memory = this.#exports.memory;
+            if (options.workerCount !== undefined) throw new Error("workerCount can only be provided for multithreaded modules");
         } else if (imports.length === 1 && imports[0].module === "env" && imports[0].name === "memory" && imports[0].kind === "memory") {
             this.#multithreaded = true;
             this.#module = module;
-            if (instance) throw new Error("Instance cannot be provided for multithreaded modules");
-            this.newInstance();
+            this.#workers = [];
+            this.#workerCount = options.workerCount ?? navigator.hardwareConcurrency;
+            if (options.instance) throw new Error("instance cannot be provided for multithreaded modules");
         } else {
             throw new Error("Unsupported module");
         }
     }
 
-    /** @param {number} [numWorkers] */
-    newInstance(numWorkers) {
-        if (this.#multithreaded) {
-            if (this.#workers?.length > 0) {
-                // Stop existing workers
-                for (const worker of this.#workers) {
-                    worker.terminate();
-                }
-                numWorkers ??= this.#workers.length;
-                this.#workers = [];
-            }
-            numWorkers ??= navigator.hardwareConcurrency;
+    newInstance() {
+        this.stopInstance();
 
+        if (this.#multithreaded) {
+            console.debug("creating multithreaded instance");
             this.#memory = new WebAssembly.Memory({initial: 96, maximum: 2048, shared: true});
             this.#instance = new WebAssembly.Instance(this.#module, {env: {memory: this.#memory}});
 
@@ -151,21 +148,45 @@ export class Aoc {
             //
             // Allocate all the stacks at once to avoid memory growing as workers start, which seems to cause problems.
             const stacks = [];
-            for (let i = 0; i < numWorkers; i++) {
+            for (let i = 0; i < this.#workerCount; i++) {
                 stacks.push(this.#exports.allocate_stack(stackSize + tlsSize, align));
             }
 
             this.#workers = [];
-            for (let i = 0; i < numWorkers; i++) {
+            for (let i = 0; i < this.#workerCount; i++) {
                 const worker = new Worker("./worker.mjs", {type: "module"});
                 worker.postMessage(["thread", this.#module, this.#memory, stacks[i] + stackSize]);
                 this.#workers.push(worker);
             }
         } else {
+            console.debug("creating single-threaded instance");
             this.#instance = new WebAssembly.Instance(this.#module);
+            this.#memory = this.#exports.memory;
+        }
+    }
+
+    ensureInstance() {
+        if (this.#instance === undefined) {
+            this.newInstance();
         }
     }
 
+    stopInstance() {
+        if (this.#multithreaded && this.#workers !== undefined) {
+            for (const worker of this.#workers) {
+                try {
+                    worker.terminate();
+                } catch (e) {
+                    console.warn(e);
+                }
+            }
+            this.#workers = [];
+        }
+
+        this.#instance = undefined;
+        this.#memory = undefined;
+    }
+
     /**
      * @param {number} year
      * @param {number} day
@@ -176,6 +197,8 @@ export class Aoc {
      * @return {{success: true, part1: string, part2: string} | {success: false, error: string, stack?: string, panic_location?: string}}
      */
     run(year, day, input, isExample = false, part1 = true, part2 = true) {
+        this.ensureInstance();
+
         let success;
         try {
             this.#write(input);
@@ -191,7 +214,7 @@ export class Aoc {
                 console.warn(e2);
             }
 
-            this.newInstance()
+            this.stopInstance();
 
             if (panic_payload.length > 0) {
                 return {

crates/aoc_wasm/web/web.mjs

@@ -1,7 +1,7 @@
 import {Aoc} from "./aoc.mjs";
 
 const COMMIT = "${GIT_COMMIT}";
-console.log("Commit", COMMIT);
+console.log("commit", COMMIT);
 
 const MODULE_PATHS = [
     "./aoc-simd128.wasm",
@@ -15,11 +15,11 @@ let module;
 for (const path of MODULE_PATHS) {
     try {
         module = await WebAssembly.compileStreaming(fetch(path));
-        console.log("Using " + path);
+        console.log("using " + path);
         break;
     } catch (err) {
         if (err instanceof WebAssembly.CompileError) {
-            console.warn("Compiling " + path + " failed: " + err);
+            console.warn("compiling " + path + " failed: " + err);
         } else {
             throw err;
         }
@@ -291,7 +291,7 @@ inputTextarea.addEventListener("drop", (e) => {
 
     const file = e.dataTransfer.files[0];
     if (file !== undefined) {
-        console.log("Loading input from file: ", file);
+        console.log("loading input from file: ", file);
 
         const reader = new FileReader();
         reader.onload = (event) => {
@@ -305,14 +305,14 @@ inputTextarea.addEventListener("drop", (e) => {
 
     const text = e.dataTransfer.getData("text/plain");
     if (text !== undefined) {
-        console.log("Loading input from drag text");
+        console.log("loading input from drag text");
         inputTextarea.value = text;
         saveInput();
 
         return;
     }
 
-    console.warn("No supported drag data");
+    console.warn("no supported drag data");
 });
 inputTextarea.addEventListener("change", saveInput);
 inputTextarea.addEventListener("input", saveInput);

crates/aoc_wasm/web/worker.mjs

@@ -7,13 +7,15 @@ onmessage = (e) => {
     switch (e.data.shift()) {
         case "init":
             instance = new Aoc(...e.data);
+            instance.ensureInstance();
             break;
         case "run":
             console.time("solution");
             const result = instance.run(...e.data);
             console.timeEnd("solution");
             postMessage(result);
             console.log(result);
+            instance.ensureInstance();
             break;
         case "thread":
             const [module, memory, ptr] = e.data;