@cfasim-ui/docs 0.4.5 → 0.4.6

package/wasm/wasmWorkerApi.ts

@@ -10,9 +10,17 @@ interface WorkerMessage {
 
 let lastId = 0;
 
-const worker = new Worker(new URL("./wasm.worker.ts", import.meta.url), {
-  type: "module",
-});
+function newWorker(): Worker {
+  return new Worker(new URL("./wasm.worker.ts", import.meta.url), {
+    type: "module",
+  });
+}
+
+let worker = newWorker();
+
+// Rejects keyed by request id so `cancelWasm` can fail in-flight promises
+// when we terminate the worker mid-run.
+const pendingRejects = new Map<number, (reason: Error) => void>();
 
 export function runWasm(
   model: string,
@@ -21,10 +29,16 @@ export function runWasm(
 ): Promise<unknown> {
   return new Promise((resolve, reject) => {
     const id = ++lastId;
+    // Capture the current worker so a later `cancelWasm` that swaps the
+    // worker reference doesn't leave the listener attached to a fresh
+    // (unrelated) worker.
+    const myWorker = worker;
+    pendingRejects.set(id, reject);
 
     function listener(event: MessageEvent<TransferableResponse>) {
       if (event.data?.id !== id) return;
-      worker.removeEventListener("message", listener);
+      myWorker.removeEventListener("message", listener);
+      pendingRejects.delete(id);
       if (event.data.error) {
         reject(new Error(event.data.error));
       } else {
@@ -32,7 +46,25 @@ export function runWasm(
       }
     }
 
-    worker.addEventListener("message", listener);
-    worker.postMessage({ id, model, fn, args } as WorkerMessage);
+    myWorker.addEventListener("message", listener);
+    myWorker.postMessage({ id, model, fn, args } as WorkerMessage);
   });
 }
+
+/**
+ * Terminate the worker and reject any in-flight `runWasm` promises with
+ * `Error("cancelled")`. A fresh worker is spawned for subsequent calls,
+ * which means the wasm module cache is lost and the next call pays the
+ * one-time module init cost. Use only when the running computation is
+ * known to be obsolete (e.g. the user has changed inputs).
+ */
+export function cancelWasm(): void {
+  if (pendingRejects.size === 0) return;
+  worker.terminate();
+  const rejects = Array.from(pendingRejects.values());
+  pendingRejects.clear();
+  worker = newWorker();
+  for (const reject of rejects) {
+    reject(new Error("cancelled"));
+  }
+}