more linear memory usage and recovery from failure of snapshot use

src/mono/wasm/runtime/assets.ts

@@ -8,7 +8,7 @@ import { mono_wasm_load_bytes_into_heap } from "./memory";
 import { endMeasure, MeasuredBlock, startMeasure } from "./profiler";
 import { createPromiseController, PromiseAndController } from "./promise-controller";
 import { delay } from "./promise-utils";
-import { abort_startup, beforeOnRuntimeInitialized } from "./startup";
+import { abort_startup, beforeOnRuntimeInitialized, memorySnapshotIsResolved } from "./startup";
 import { AssetBehaviours, AssetEntry, AssetEntryInternal, LoadingResource, mono_assert, ResourceRequest } from "./types";
 import { InstantiateWasmSuccessCallback, VoidPtr } from "./types/emscripten";
 
@@ -59,7 +59,10 @@ export function resolve_asset_path(behavior: AssetBehaviours) {
     return asset;
 }
 export async function mono_download_assets(): Promise<void> {
-    if (runtimeHelpers.memoryIsLoaded) {
+    // continue after we know if memory snapshot is available or not
+    await memorySnapshotIsResolved.promise;
+
+    if (runtimeHelpers.useMemorySnapshot) {
         allAssetsInMemory.promise_control.resolve();
         return;
     }

src/mono/wasm/runtime/startup.ts

@@ -27,7 +27,7 @@ import { mono_wasm_init_diagnostics } from "./diagnostics";
 import { preAllocatePThreadWorkerPool, instantiateWasmPThreadWorkerPool } from "./pthreads/browser";
 import { export_linker } from "./exports-linker";
 import { endMeasure, MeasuredBlock, startMeasure } from "./profiler";
-import { getMemory, storeMemory } from "./storage";
+import { getMemorySnapshot, storeMemorySnapshot, getMemorySnapshotSize } from "./storage";
 
 // legacy
 import { init_legacy_exports } from "./net6-legacy/corebindings";
@@ -39,6 +39,7 @@ let configLoaded = false;
 export const dotnetReady = createPromiseController<any>();
 export const afterConfigLoaded = createPromiseController<MonoConfig>();
 export const beforeInstantiateWasm = createPromiseController<void>();
+export const memorySnapshotIsResolved = createPromiseController<void>();
 export const afterInstantiateWasm = createPromiseController<void>();
 export const beforePreInit = createPromiseController<void>();
 export const afterPreInit = createPromiseController<void>();
@@ -108,7 +109,9 @@ function instantiateWasm(
 
     const mark = startMeasure();
     if (userInstantiateWasm) {
+        // user wasm doesn't support memory snapshots
         beforeInstantiateWasm.promise_control.resolve();
+        memorySnapshotIsResolved.promise_control.resolve();
         const exports = userInstantiateWasm(imports, (instance: WebAssembly.Instance, module: WebAssembly.Module | undefined) => {
             endMeasure(mark, MeasuredBlock.instantiateWasm);
             afterInstantiateWasm.promise_control.resolve();
@@ -179,7 +182,6 @@ async function preInitWorkerAsync() {
     const mark = startMeasure();
     try {
         if (runtimeHelpers.diagnosticTracing) console.debug("MONO_WASM: preInitWorker");
-        beforePreInit.promise_control.resolve();
         mono_wasm_pre_init_essential(true);
         await init_polyfills_async();
         afterPreInit.promise_control.resolve();
@@ -286,6 +288,7 @@ export function abort_startup(reason: any, should_exit: boolean): void {
     if (runtimeHelpers.diagnosticTracing) console.trace("MONO_WASM: abort_startup");
     dotnetReady.promise_control.reject(reason);
     beforeInstantiateWasm.promise_control.reject(reason);
+    memorySnapshotIsResolved.promise_control.reject(reason);
     afterInstantiateWasm.promise_control.reject(reason);
     beforePreInit.promise_control.reject(reason);
     afterPreInit.promise_control.reject(reason);
@@ -334,7 +337,9 @@ async function mono_wasm_pre_init_full(): Promise<void> {
     if (runtimeHelpers.diagnosticTracing) console.debug("MONO_WASM: mono_wasm_pre_init_full");
     Module.addRunDependency("mono_wasm_pre_init_full");
 
+    // continue after wasm is downloaded
     await beforeInstantiateWasm.promise;
+
     await mono_download_assets();
 
     Module.removeRunDependency("mono_wasm_pre_init_full");
@@ -348,8 +353,8 @@ async function mono_wasm_before_user_runtime_initialized(): Promise<void> {
         mono_wasm_globalization_init();
 
         if (!runtimeHelpers.mono_wasm_load_runtime_done) mono_wasm_load_runtime("unused", config.debugLevel);
-        if (config.cacheMemory && !runtimeHelpers.memoryIsLoaded) {
-            await storeMemory(Module.HEAP8.buffer);
+        if (config.cacheMemory && !runtimeHelpers.useMemorySnapshot) {
+            await storeMemorySnapshot(Module.HEAP8.buffer);
         }
         if (MonoWasmThreads) {
             preAllocatePThreadWorkerPool(MONO_PTHREAD_POOL_SIZE, config);
@@ -460,19 +465,23 @@ async function instantiate_wasm_module(
 ): Promise<void> {
     // this is called so early that even Module exports like addRunDependency don't exist yet
     try {
-        let memoryBytes: ArrayBuffer | undefined = undefined;
+        let memorySize: number | undefined = undefined;
         await mono_wasm_load_config(Module.configSrc);
         if (runtimeHelpers.diagnosticTracing) console.debug("MONO_WASM: instantiate_wasm_module");
         const assetToLoad = resolve_asset_path("dotnetwasm");
+        // FIXME: this would not apply re-try (on connection reset during download) for dotnet.wasm because we could not download the buffer before we pass it to instantiate_wasm_asset
+        await start_asset_download(assetToLoad);
+        beforeInstantiateWasm.promise_control.resolve();
 
         if (config.cacheMemory && config.assetsHash) {
-            memoryBytes = await getMemory();
-            runtimeHelpers.memoryIsLoaded = !!memoryBytes;
+            memorySize = await getMemorySnapshotSize();
+            runtimeHelpers.useMemorySnapshot = !!memorySize;
+        }
+        if (!runtimeHelpers.useMemorySnapshot) {
+            // we should start downloading DLLs etc as they are not in the snapshot
+            memorySnapshotIsResolved.promise_control.resolve();
         }
-        beforeInstantiateWasm.promise_control.resolve();
 
-        // FIXME: this would not apply re-try (on connection reset during download) for dotnet.wasm because we could not download the buffer before we pass it to instantiate_wasm_asset
-        await start_asset_download(assetToLoad);
         await beforePreInit.promise;
         Module.addRunDependency("instantiate_wasm_module");
 
@@ -483,13 +492,23 @@ async function instantiate_wasm_module(
         assetToLoad.buffer = null as any; // GC
         if (runtimeHelpers.diagnosticTracing) console.debug("MONO_WASM: instantiate_wasm_module done");
 
-        if (runtimeHelpers.memoryIsLoaded) {
-            const wasmMemory = anyModule.asm.memory;
-            // .grow() takes a delta compared to the previous size
-            wasmMemory.grow((memoryBytes!.byteLength - wasmMemory.buffer.byteLength + 65535) >>> 16);
-            runtimeHelpers.updateGlobalBufferAndViews(wasmMemory.buffer);
-            Module.HEAP8.set(new Int8Array(memoryBytes!), 0);
-            if (runtimeHelpers.diagnosticTracing) console.debug("MONO_WASM: Loaded memory from cache");
+        if (runtimeHelpers.useMemorySnapshot) {
+            try {
+                const wasmMemory = anyModule.asm.memory;
+                // .grow() takes a delta compared to the previous size
+                wasmMemory.grow((memorySize! - wasmMemory.buffer.byteLength + 65535) >>> 16);
+                runtimeHelpers.updateGlobalBufferAndViews(wasmMemory.buffer);
+
+                // get the bytes after we re-sized the memory, so that we don't have too much memory in use at the same time
+                const memoryBytes = await getMemorySnapshot();
+                Module.HEAP8.set(new Int8Array(memoryBytes!), 0);
+                if (runtimeHelpers.diagnosticTracing) console.debug("MONO_WASM: Loaded memory from cache");
+            } catch (err) {
+                console.warn("MONO_WASM: failed to load memory snapshot", err);
+                runtimeHelpers.useMemorySnapshot = false;
+            }
+            // now we know if the loading of memory succeeded or not, we can start loading the rest of the assets
+            memorySnapshotIsResolved.promise_control.resolve();
         }
         afterInstantiateWasm.promise_control.resolve();
     } catch (err) {
@@ -531,7 +550,7 @@ export function mono_wasm_load_runtime(unused?: string, debugLevel?: number): vo
     runtimeHelpers.mono_wasm_load_runtime_done = true;
     try {
         const mark = startMeasure();
-        if (!runtimeHelpers.memoryIsLoaded) {
+        if (!runtimeHelpers.useMemorySnapshot) {
             if (debugLevel == undefined) {
                 debugLevel = 0;
                 if (config && config.debugLevel) {

src/mono/wasm/runtime/storage.ts

@@ -8,7 +8,7 @@ import { runtimeHelpers } from "./imports";
 const memoryPrefix = "https://dotnet.generated.invalid/wasm-memory";
 
 // adapted from Blazor's WebAssemblyResourceLoader.ts
-async function open(): Promise<Cache | null> {
+async function openCache(): Promise<Cache | null> {
     // caches will be undefined if we're running on an insecure origin (secure means https or localhost)
     if (typeof globalThis.caches === "undefined") {
         return null;
@@ -43,14 +43,34 @@ async function open(): Promise<Cache | null> {
     }
 }
 
-export async function getMemory(): Promise<ArrayBuffer | undefined> {
+export async function getMemorySnapshotSize(): Promise<number | undefined> {
     try {
         const inputsHash = await getInputsHash();
         if (!inputsHash) {
             return undefined;
         }
         const cacheKey = `${memoryPrefix}-${ProductVersion}-${GitHash}-${inputsHash}`;
-        const cache = await open();
+        const cache = await openCache();
+        if (!cache) {
+            return undefined;
+        }
+        const res = await cache.match(cacheKey);
+        const contentLength = res?.headers.get("content-length");
+        return contentLength ? parseInt(contentLength) : undefined;
+    } catch (ex) {
+        console.warn("MONO_WASM: Failed find memory snapshot in the cache", ex);
+        return undefined;
+    }
+}
+
+export async function getMemorySnapshot(): Promise<ArrayBuffer | undefined> {
+    try {
+        const inputsHash = await getInputsHash();
+        if (!inputsHash) {
+            return undefined;
+        }
+        const cacheKey = `${memoryPrefix}-${ProductVersion}-${GitHash}-${inputsHash}`;
+        const cache = await openCache();
         if (!cache) {
             return undefined;
         }
@@ -60,18 +80,19 @@ export async function getMemory(): Promise<ArrayBuffer | undefined> {
         }
         return res.arrayBuffer();
     } catch (ex) {
+        console.warn("MONO_WASM: Failed load memory snapshot from the cache", ex);
         return undefined;
     }
 }
 
-export async function storeMemory(memory: ArrayBuffer) {
+export async function storeMemorySnapshot(memory: ArrayBuffer) {
     try {
         const inputsHash = await getInputsHash();
         if (!inputsHash) {
             return;
         }
         const cacheKey = `${memoryPrefix}-${ProductVersion}-${GitHash}-${inputsHash}`;
-        const cache = await open();
+        const cache = await openCache();
         if (!cache) {
             return;
         }
@@ -85,15 +106,16 @@ export async function storeMemory(memory: ArrayBuffer) {
 
         await cache.put(cacheKey, responseToCache);
 
-        cleanupMemory(cacheKey); // no await
+        cleanupMemorySnapshots(cacheKey); // no await
     } catch (ex) {
+        console.warn("MONO_WASM: Failed to store memory snapshot in the cache", ex);
         return;
     }
 }
 
-export async function cleanupMemory(protectKey: string) {
+export async function cleanupMemorySnapshots(protectKey: string) {
     try {
-        const cache = await open();
+        const cache = await openCache();
         if (!cache) {
             return;
         }

src/mono/wasm/runtime/types.ts

@@ -233,7 +233,7 @@ export type RuntimeHelpers = {
     locateFile: (path: string, prefix?: string) => string,
     javaScriptExports: JavaScriptExports,
     loadedFiles: string[],
-    memoryIsLoaded: boolean,
+    useMemorySnapshot: boolean,
     subtle: SubtleCrypto | null,
     preferredIcuAsset: string | null,
     timezone: string | null,