From f5adce2309ce65cbca8c27c744e8b44012a9d821 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Tue, 3 Dec 2024 18:15:49 +0900 Subject: [PATCH] Offload the instrumentor to a Web Worker if possible --- src/wasm-memprof.ts | 91 ++++++++++++++++++++++++++++++++++++++------- 1 file changed, 78 insertions(+), 13 deletions(-) diff --git a/src/wasm-memprof.ts b/src/wasm-memprof.ts index 62a6794..72c7051 100644 --- a/src/wasm-memprof.ts +++ b/src/wasm-memprof.ts @@ -1,20 +1,73 @@ -// @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 | undefined; +type Bindgen = { + instrument_allocator: (moduleBytes: Uint8Array) => Promise +} +let _initBindgen: Promise | 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 { +async function initBindgen(WebAssembly: typeof globalThis.WebAssembly, options: WMProfOptions): Promise { + // 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 { + 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 @@ -22,6 +75,11 @@ async function initBindgen(WebAssembly: typeof globalThis.WebAssembly): Promise< // work for us. const bgModule = await WebAssembly.compile(bindgenWasm); await init(bgModule); + return { + async instrument_allocator(moduleBytes: Uint8Array): Promise { + 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 { +async function instrumentModule(moduleBytes: BufferSource, WebAssembly: typeof globalThis.WebAssembly, options: WMProfOptions): Promise { 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);