Retry Wasm loads after a shared failure

src/core/wasm_image.js

@@ -57,12 +57,11 @@ class WasmImage {
   }
 
   async #instantiateWasm(failureCallback, imports, successCallback) {
-    const context = {
-      handler: WasmImage.#options.getHandler?.(),
-      options: WasmImage.#options,
-    };
     const loaded = await WasmLoader.load(
-      context,
+      {
+        handler: WasmImage.#options.getHandler?.(),
+        options: WasmImage.#options,
+      },
       this._filename,
       this._noWasmFilename,
       instantiate => successCallback({ exports: instantiate(imports) })

src/core/wasm_loader.js

@@ -16,16 +16,13 @@
 import { isNodeJS, warn } from "../shared/util.js";
 import { fetchBinaryData } from "./core_utils.js";
 
-/**
- * Loads Wasm modules, or their JS fallbacks.
- */
 class WasmLoader {
   static #modules = new Map();
 
   /**
-   * Fetch a binary file in the worker-thread or through the main-thread.
+   * Fetch a binary file in the worker or via the main thread.
    * @param {object} context - Document handler and options.
-   * @param {string} kind - Name of the option holding the base URL.
+   * @param {string} kind - Base URL option name.
    * @param {string} filename
    * @returns {Promise<Uint8Array>}
    */
@@ -41,16 +38,21 @@ class WasmLoader {
   }
 
   static #compile(context, filename) {
-    // Cache the promise, so that concurrent calls share a single compilation.
-    return this.#modules.getOrInsertComputed(filename, () =>
-      this.fetch(context, "wasmUrl", filename)
-        .then(bytes => WebAssembly.compile(bytes))
-        .catch(ex => {
-          // Allow a later call to retry.
-          this.#modules.delete(filename);
-          throw ex;
-        })
-    );
+    const cached = this.#modules.get(filename);
+    if (cached) {
+      // Retry with this context if the shared attempt fails.
+      return cached.catch(() => this.#compile(context, filename));
+    }
+    const promise = this.fetch(context, "wasmUrl", filename)
+      .then(bytes => WebAssembly.compile(bytes))
+      .catch(ex => {
+        // Don't cache failures.
+        this.#modules.delete(filename);
+        throw ex;
+      });
+    // Share compilation between concurrent callers.
+    this.#modules.set(filename, promise);
+    return promise;
   }
 
   static #importFallback(wasmUrl, filename) {
@@ -70,13 +72,12 @@ class WasmLoader {
   }
 
   /**
-   * Initialize a module with its Wasm file, or with its JS fallback.
+   * Initialize a Wasm module or its JS fallback.
    * @param {object} context - Document handler and options.
    * @param {string} filename - Wasm file.
    * @param {string} noWasmFilename - JS fallback file.
-   * @param {Function} init - Called with `instantiate(imports)`, which
-   *   returns the exports.
-   * @returns {Promise<boolean>} Whether `init` succeeded.
+   * @param {Function} init - Receives `instantiate(imports) => exports`.
+   * @returns {Promise<boolean>} Whether `init` was called without throwing.
    */
   static async load(context, filename, noWasmFilename, init) {
     const { useWasm, wasmUrl } = context.options;

test/unit/wasm_loader_spec.js

@@ -18,7 +18,7 @@ import { WASM_URL } from "./test_utils.js";
 import { WasmLoader } from "../../src/core/wasm_loader.js";
 
 describe("WasmLoader", function () {
-  // Smallest valid module: the magic number and the version.
+  // Wasm header: magic number and version.
   const EMPTY_WASM = new Uint8Array([
     0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00,
   ]);
@@ -30,33 +30,32 @@ describe("WasmLoader", function () {
     };
   }
 
-  it("fetches and compiles a Wasm file once for back-to-back calls", async function () {
-    const sendWithPromise = jasmine
+  let sendWithPromise, context, init;
+
+  beforeEach(function () {
+    sendWithPromise = jasmine
       .createSpy("sendWithPromise")
       .and.resolveTo(EMPTY_WASM);
-    const context = getContext(sendWithPromise);
-    const init1 = jasmine.createSpy("init1");
-    const init2 = jasmine.createSpy("init2");
+    context = getContext(sendWithPromise);
+    init = jasmine.createSpy("init");
+  });
 
+  it("shares a Wasm fetch between concurrent calls", async function () {
     const loaded = await Promise.all([
-      WasmLoader.load(context, "back_to_back.wasm", "unused.js", init1),
-      WasmLoader.load(context, "back_to_back.wasm", "unused.js", init2),
+      WasmLoader.load(context, "back_to_back.wasm", "unused.js", init),
+      WasmLoader.load(context, "back_to_back.wasm", "unused.js", init),
     ]);
     expect(loaded).toEqual([true, true]);
     expect(sendWithPromise).toHaveBeenCalledOnceWith("FetchBinaryData", {
       kind: "wasmUrl",
       filename: "back_to_back.wasm",
     });
-    expect(init1).toHaveBeenCalledOnceWith(jasmine.any(Function));
-    expect(init2).toHaveBeenCalledOnceWith(jasmine.any(Function));
+    expect(init).toHaveBeenCalledTimes(2);
+    expect(init).toHaveBeenCalledWith(jasmine.any(Function));
   });
 
-  it("fetches a Wasm file again after a failure", async function () {
-    const sendWithPromise = jasmine
-      .createSpy("sendWithPromise")
-      .and.rejectWith(new Error("Fetch failed."));
-    const context = getContext(sendWithPromise);
-    const init = jasmine.createSpy("init");
+  it("retries a failed fetch", async function () {
+    sendWithPromise.and.rejectWith(new Error("Fetch failed."));
 
     expect(
       await WasmLoader.load(context, "retry.wasm", "unused.js", init)
@@ -71,11 +70,39 @@ describe("WasmLoader", function () {
     expect(init).toHaveBeenCalledOnceWith(jasmine.any(Function));
   });
 
-  it("loads a JS fallback relative to the current directory in Node.js", async function () {
+  it("retries a failed shared fetch with its own context", async function () {
+    const failingContext = getContext(
+      jasmine.createSpy("failing").and.rejectWith(new Error("Fetch failed."))
+    );
+    const loaded = await Promise.all([
+      WasmLoader.load(failingContext, "concurrent.wasm", "unused.js", init),
+      WasmLoader.load(context, "concurrent.wasm", "unused.js", init),
+    ]);
+    expect(loaded).toEqual([false, true]);
+    expect(sendWithPromise).toHaveBeenCalledOnceWith("FetchBinaryData", {
+      kind: "wasmUrl",
+      filename: "concurrent.wasm",
+    });
+    expect(init).toHaveBeenCalledOnceWith(jasmine.any(Function));
+  });
+
+  it("stops retrying when all concurrent calls fail", async function () {
+    sendWithPromise.and.rejectWith(new Error("Fetch failed."));
+
+    const loaded = await Promise.all(
+      Array.from({ length: 3 }, () =>
+        WasmLoader.load(context, "failures.wasm", "unused.js", init)
+      )
+    );
+    expect(loaded).toEqual([false, false, false]);
+    expect(sendWithPromise).toHaveBeenCalledTimes(3);
+    expect(init).not.toHaveBeenCalled();
+  });
+
+  it("resolves Node.js fallback paths from the working directory", async function () {
     if (!isNodeJS) {
       pending("Only for Node.js.");
     }
-    const init = jasmine.createSpy("init");
     const loaded = await WasmLoader.load(
       {
         options: {