Offload the instrumentor to a Web Worker if possible

src/wasm-memprof.ts

@@ -1,27 +1,85 @@
-// @ts-check
-
 import * as bg from "../bindgen/pkg/wasm_memprof.js";
 import init from "../bindgen/pkg/wasm_memprof.js";
 import { perftools } from "./profile.pb.js";
 // @ts-ignore
 import bindgenWasm from "../bindgen/pkg/wasm_memprof_bg.wasm" assert { type: "binary" };
 
-let _initBindgen: Promise<void> | undefined;
+type Bindgen = {
+    instrument_allocator: (moduleBytes: Uint8Array) => Promise<Uint8Array>
+}
+let _initBindgen: Promise<Bindgen> | undefined;
 
 /**
  * Initialize the WebAssembly bindgen module.
  *
- * @param WebAssembly The WebAssembly object
+ * @param WebAssembly The original WebAssembly object
+ * @param options Options for the profiler
  * @returns A promise that resolves once the bindgen module is initialized
  */
-async function initBindgen(WebAssembly: typeof globalThis.WebAssembly): Promise<void> {
+async function initBindgen(WebAssembly: typeof globalThis.WebAssembly, options: WMProfOptions): Promise<Bindgen> {
+    // Use Web Worker if available and not disabled by the user for the
+    // following reasons:
+    //
+    // 1. The instrumentor is a heavy operation and can block the main
+    //    thread for a long time.
+    // 2. When the tracee module is very large, the instrumentor consumes
+    //    a lot of memory. However, even if the instrumentor calls `free`, the
+    //    Wasm memory space will not be shrunk due to the limitation of Wasm
+    //    itself. Therefore, the only way to deallocate the Wasm memory space is
+    //    to deallocate the WebAssembly instance itself.
+    //    However, glue code generated by wasm-bindgen does not expose a way to
+    //    deallocate the WebAssembly instance. So, we use a Web Worker to
+    //    isolate the JS memory space and deallocate the Worker instance to free
+    //    the Wasm instance.
+
+    if (typeof Worker === "undefined" || !(options.useWorker ?? true)) {
+        const bindgen = await __initBindgen(WebAssembly);
+        return bindgen;
+    }
+
+    const workerScript = `
+            onmessage = async (e) => {
+                const { selfScript, bytes } = e.data;
+                const { __initBindgen } = await import(selfScript);
+                const bindgen = await __initBindgen(globalThis.WebAssembly);
+                const instrumented = await bindgen.instrument_allocator(bytes);
+                postMessage(instrumented, [instrumented.buffer]);
+            }
+        `;
+    const blob = new Blob([workerScript], { type: "application/javascript" });
+    return {
+        async instrument_allocator(moduleBytes: Uint8Array): Promise<Uint8Array> {
+            const worker = new Worker(URL.createObjectURL(blob));
+            return new Promise((resolve, reject) => {
+                worker.onmessage = (e) => {
+                    resolve(e.data);
+                    worker.terminate();
+                }
+                worker.onerror = (e) => {
+                    reject(e);
+                }
+                worker.postMessage({
+                    selfScript: import.meta.url,
+                    bytes: moduleBytes
+                });
+            });
+        }
+    }
+}
+
+export async function __initBindgen(WebAssembly: typeof globalThis.WebAssembly) {
     if (!_initBindgen) {
         _initBindgen = (async () => {
             // @ts-ignore: We expect bindgenWasm to be a Uint8Array but
             // wasm-pack emits a d.ts for the wasm module that doesn't
             // work for us.
             const bgModule = await WebAssembly.compile(bindgenWasm);
             await init(bgModule);
+            return {
+                async instrument_allocator(moduleBytes: Uint8Array): Promise<Uint8Array> {
+                    return bg.instrument_allocator(moduleBytes);
+                }
+            }
         })();
     }
     return _initBindgen;
@@ -35,10 +93,10 @@ async function initBindgen(WebAssembly: typeof globalThis.WebAssembly): Promise<
  * @param WebAssembly The WebAssembly object
  * @returns A promise that resolves to the instrumented module bytes
  */
-async function instrumentModule(moduleBytes: BufferSource, WebAssembly: typeof globalThis.WebAssembly): Promise<Uint8Array> {
+async function instrumentModule(moduleBytes: BufferSource, WebAssembly: typeof globalThis.WebAssembly, options: WMProfOptions): Promise<Uint8Array> {
     const arrayBuffer = moduleBytes instanceof ArrayBuffer ? moduleBytes : moduleBytes.buffer;
-    await initBindgen(WebAssembly);
-    return bg.instrument_allocator(new Uint8Array(arrayBuffer));
+    const bg = await initBindgen(WebAssembly, options);
+    return await bg.instrument_allocator(new Uint8Array(arrayBuffer));
 }
 
 /**
@@ -264,7 +322,14 @@ type WMProfOptions = {
      * @param name raw function name (typically mangled)
      * @returns demangled function name
      */
-    demangler?: (name: string) => string
+    demangler?: (name: string) => string,
+
+    /**
+     * Use Web Worker to isolate the instrumentor from the main thread.
+     *
+     * @default true
+     */
+    useWorker?: boolean
 };
 
 type AllocationInfo = {
@@ -343,12 +408,12 @@ export class WMProf {
 
         const newMethods = {
             compile: async (bufferSource) => {
-                const buffer = await instrumentModule(bufferSource, WebAssembly);
+                const buffer = await instrumentModule(bufferSource, WebAssembly, options);
                 return WebAssembly.compile(buffer);
             },
             compileStreaming: async (response) => {
                 const buffer = await (await response).arrayBuffer();
-                const instrumented = await instrumentModule(buffer, WebAssembly);
+                const instrumented = await instrumentModule(buffer, WebAssembly, options);
                 return WebAssembly.compile(instrumented);
             },
             instantiate: async (bufferSource, importObject) => {
@@ -363,14 +428,14 @@ export class WMProf {
                     WMProf.install(instance, wmprof);
                     return instance;
                 }
-                const buffer = await instrumentModule(bufferSource, WebAssembly);
+                const buffer = await instrumentModule(bufferSource, WebAssembly, options);
                 const result = await WebAssembly.instantiate(buffer, instrumentImportObject(importObject, wmprof));
                 WMProf.install(result.instance, wmprof);
                 return result;
             },
             instantiateStreaming: async (response, importObject) => {
                 const buffer = await (await response).arrayBuffer();
-                const instrumented = await instrumentModule(buffer, WebAssembly);
+                const instrumented = await instrumentModule(buffer, WebAssembly, options);
                 const wmprof = new WMProf(options);
                 const result = await WebAssembly.instantiate(instrumented, instrumentImportObject(importObject, wmprof));
                 WMProf.install(result.instance, wmprof);