perf(hash): chunk message input to reduce WASM heap size (#1010)

hash/_wasm/build.ts

@@ -60,7 +60,8 @@ const generatedWasm = await Deno.readFile("./out/deno_hash_bg.wasm");
 
 // Replace the lines loading the WASM from an external file with our inlined
 // copy, to avoid the need for net or read permissions.
-const inlinedScript = `// deno-lint-ignore-file
+const inlinedScript = `\
+  // deno-lint-ignore-file
   import * as base64 from "../../encoding/base64.ts"; ${
   generatedScript.replace(
     /^const file =.*?;\nconst wasmFile =.*?;\nconst wasmModule =.*?;\n/sm,
@@ -69,7 +70,11 @@ const inlinedScript = `// deno-lint-ignore-file
       base64.encode(generatedWasm).replace(/.{78}/g, "$&\\\n")
     }"));`,
   )
-}`;
+}
+
+  // only exposed for testing
+  export const _wasm = wasm;
+`;
 
 await Deno.writeFile("wasm.js", new TextEncoder().encode(inlinedScript));
 

hash/_wasm/hash.ts

@@ -41,7 +41,20 @@ export class Hash implements Hasher {
       throw new Error(TYPE_ERROR_MSG);
     }
 
-    updateHash(this.#hash, msg);
+    // Messages will be split into chunks of this size to avoid unneccessarily
+    // increasing the size of the WASM heap.
+    const CHUNK_SIZE = 65_536;
+
+    for (let offset = 0; offset < msg.length; offset += CHUNK_SIZE) {
+      updateHash(
+        this.#hash,
+        new Uint8Array(
+          msg.buffer,
+          offset,
+          Math.min(CHUNK_SIZE, msg.length - offset),
+        ),
+      );
+    }
 
     return this;
   }

hash/_wasm/wasm.js

@@ -2504,3 +2504,6 @@ LjE5LjAMd2FzbS1iaW5kZ2VuBjAuMi43NA==",
 );
 const wasmInstance = new WebAssembly.Instance(wasmModule, imports);
 const wasm = wasmInstance.exports;
+
+// only exposed for testing
+export const _wasm = wasm;

hash/test.ts

@@ -1,6 +1,9 @@
 // Copyright 2018-2021 the Deno authors. All rights reserved. MIT license.
-import { assertEquals, assertThrows } from "../testing/asserts.ts";
+import { assert, assertEquals, assertThrows } from "../testing/asserts.ts";
 import { createHash, SupportedAlgorithm } from "./mod.ts";
+import { dirname, fromFileUrl } from "../path/mod.ts";
+
+const moduleDir = dirname(fromFileUrl(import.meta.url));
 
 const millionAs = "a".repeat(1000000);
 
@@ -355,6 +358,93 @@ Deno.test("[hash/all/base64] testAllBase64", () => {
   }
 });
 
+Deno.test("[hash/memory_use] testMemoryUse", async () => {
+  const process = Deno.run({
+    cmd: [Deno.execPath(), "--quiet", "run", "--no-check", "-"],
+    cwd: moduleDir,
+    stdout: "piped",
+    stdin: "piped",
+  });
+
+  await process.stdin.write(
+    new TextEncoder().encode(`
+      import { createHash } from "./mod.ts";
+      import { _wasm } from "./_wasm/wasm.js";
+
+      const { memory } = _wasm as { memory: WebAssembly.Memory };
+
+      const heapBytesInitial = memory.buffer.byteLength;
+
+      const smallData = new Uint8Array(64);
+      const smallHasher = createHash("md5");
+      smallHasher.update(smallData);
+      const smallDigest = smallHasher.toString();
+      const heapBytesAfterSmall = memory.buffer.byteLength;
+
+      const largeData = new Uint8Array(64_000_000);
+      const largeHasher = createHash("md5");
+      largeHasher.update(largeData);
+      const largeDigest = largeHasher.toString();
+      const heapBytesAfterLarge = memory.buffer.byteLength;
+
+      console.log(JSON.stringify(
+        {
+          heapBytesInitial,
+          smallDigest,
+          heapBytesAfterSmall,
+          largeDigest,
+          heapBytesAfterLarge,
+        },
+        null,
+        2,
+      ));
+    `),
+  );
+  process.stdin.close();
+
+  const stdout = new TextDecoder().decode(await process.output());
+  const status = await process.status();
+  process.close();
+
+  assertEquals(status.success, true);
+  const {
+    heapBytesInitial,
+    smallDigest,
+    heapBytesAfterSmall,
+    largeDigest,
+    heapBytesAfterLarge,
+  }: {
+    heapBytesInitial: number;
+    smallDigest: string;
+    heapBytesAfterSmall: number;
+    largeDigest: string;
+    heapBytesAfterLarge: number;
+  } = JSON.parse(stdout);
+
+  assertEquals(smallDigest, "3b5d3c7d207e37dceeedd301e35e2e58");
+  assertEquals(largeDigest, "e78585b8bfda6036cfd818710a210f23");
+
+  // Heap should stay under 2MB even though we provided a 64MB input.
+  assert(
+    heapBytesInitial < 2_000_000,
+    `WASM heap was too large initially: ${
+      (heapBytesInitial / 1_000_000).toFixed(1)
+    } MB`,
+  );
+  assert(
+    heapBytesAfterSmall < 2_000_000,
+    `WASM heap was too large after small input: ${
+      (heapBytesAfterSmall / 1_000_000).toFixed(1)
+    } MB`,
+  );
+  assert(
+    heapBytesAfterLarge < 2_000_000,
+    `WASM heap was too large after large input: ${
+      (heapBytesAfterLarge / 1_000_000).toFixed(1)
+    } MB`,
+  );
+});
+
 Deno.test("[hash/double_digest] testDoubleDigest", () => {
   assertThrows(
     (): void => {