Reduce peak memory in getFileHashes

node-src/lib/getFileHashes.test.ts

@@ -1,9 +1,12 @@
-import { describe } from 'node:test';
-
-import { expect, it } from 'vitest';
+import { mkdtempSync, rmSync, writeFileSync } from 'fs';
+import { tmpdir } from 'os';
+import path from 'path';
+import { describe, expect, it, vi } from 'vitest';
 
 import { getFileHashes } from './getFileHashes';
 
+const BUFFER_BYTE_LENGTH = 64 * 1024;
+
 describe('getFileHashes', () => {
   it('should return a map of file paths to hashes', async () => {
     const hashes = await getFileHashes(['iframe.html', 'index.html'], 'node-src/__mocks__', 2);
@@ -13,4 +16,53 @@ describe('getFileHashes', () => {
       'index.html': '0e98fd69b0b01605',
     });
   });
+
+  it('returns an empty object when there are no files', async () => {
+    await expect(getFileHashes([], 'node-src/__mocks__', 2)).resolves.toEqual({});
+  });
+
+  it('allocates one 64KiB buffer per pool slot (min concurrency, file count), not one per file', async () => {
+    const fixtureDirectory = mkdtempSync(path.join(tmpdir(), 'chromatic-hash-pool-'));
+    const fileCount = 80;
+    const concurrency = 5;
+    const files: string[] = [];
+
+    for (let index = 0; index < fileCount; index++) {
+      const name = `f${index}.txt`;
+      files.push(name);
+      writeFileSync(path.join(fixtureDirectory, name), `row-${index}\n`);
+    }
+
+    const spy = vi.spyOn(Buffer, 'allocUnsafe');
+
+    await getFileHashes(files, fixtureDirectory, concurrency);
+
+    const poolAllocations = spy.mock.calls.filter((call) => call[0] === BUFFER_BYTE_LENGTH);
+    expect(poolAllocations.length).toBe(concurrency);
+
+    spy.mockRestore();
+    rmSync(fixtureDirectory, { recursive: true, force: true });
+  });
+
+  it('matches sequential hashing when reusing the buffer pool at higher concurrency', async () => {
+    const fixtureDirectory = mkdtempSync(path.join(tmpdir(), 'chromatic-hash-pool-'));
+    const fileCount = 40;
+    const files: string[] = [];
+
+    for (let index = 0; index < fileCount; index++) {
+      const name = `f${index}.txt`;
+      files.push(name);
+      // Mix short rows and a payload over 64KiB so incremental reads share pool buffers.
+      const line = `row-${index}\n`;
+      const body = index === 7 ? `${line}${'x'.repeat(70 * 1024)}` : line;
+      writeFileSync(path.join(fixtureDirectory, name), body);
+    }
+
+    const sequential = await getFileHashes(files, fixtureDirectory, 1);
+    const pooled = await getFileHashes(files, fixtureDirectory, 8);
+
+    expect(pooled).toEqual(sequential);
+
+    rmSync(fixtureDirectory, { recursive: true, force: true });
+  });
 });

node-src/lib/getFileHashes.ts

@@ -3,6 +3,8 @@ import pLimit from 'p-limit';
 import path from 'path';
 import xxHashWasm, { XXHash, XXHashAPI } from 'xxhash-wasm';
 
+const BUFFER_BYTE_LENGTH = 64 * 1024;
+
 const hashFile = (buffer: Buffer, path: string, xxhash: XXHashAPI): Promise<string> => {
   const BUFFER_SIZE = buffer.length;
 
@@ -54,20 +56,67 @@ const hashFile = (buffer: Buffer, path: string, xxhash: XXHashAPI): Promise<stri
   });
 };
 
+function createBufferPool(poolSize: number) {
+  const pool: Buffer[] = Array.from({ length: poolSize }, () =>
+    Buffer.allocUnsafe(BUFFER_BYTE_LENGTH)
+  );
+  const waitQueue: ((buffer: Buffer) => void)[] = [];
+
+  const acquire = (): Promise<Buffer> => {
+    const buffer = pool.pop();
+    if (buffer !== undefined) {
+      return Promise.resolve(buffer);
+    }
+    return new Promise<Buffer>((resolve) => {
+      waitQueue.push(resolve);
+    });
+  };
+
+  const release = (buffer: Buffer): void => {
+    const resolveNext = waitQueue.shift();
+    if (resolveNext) {
+      resolveNext(buffer);
+    } else {
+      pool.push(buffer);
+    }
+  };
+
+  return { acquire, release };
+}
+
+function effectiveConcurrency(requested: number): number {
+  if (Number.isFinite(requested) && requested >= 1) {
+    return Math.min(Math.max(Math.floor(requested), 1), 512);
+  }
+  return 48;
+}
+
 export const getFileHashes = async (files: string[], directory: string, concurrency: number) => {
-  // Limit the number of concurrent file reads and hashing operations.
-  const limit = pLimit(concurrency);
+  if (files.length === 0) {
+    return {};
+  }
+
+  const concurrentHashes = effectiveConcurrency(concurrency);
+  const limit = pLimit(concurrentHashes);
   const xxhash = await xxHashWasm();
 
-  // Pre-allocate a 64K buffer for each file, matching WASM memory page size.
-  const buffers = files.map((file) => [Buffer.allocUnsafe(64 * 1024), file] as const);
+  // Reuse a small pool of 64K buffers (one per in-flight hash) instead of pre-allocating one per
+  // file, which peaked at O(files * 64KB) heap for large Storybook outputs.
+  const poolSize = Math.min(concurrentHashes, files.length);
+  const { acquire, release } = createBufferPool(poolSize);
 
   const hashes = await Promise.all(
-    buffers.map(([buffer, file]) =>
-      limit(async () => [file, await hashFile(buffer, path.join(directory, file), xxhash)] as const)
+    files.map((file) =>
+      limit(async () => {
+        const buffer = await acquire();
+        try {
+          return [file, await hashFile(buffer, path.join(directory, file), xxhash)] as const;
+        } finally {
+          release(buffer);
+        }
+      })
     )
   );
 
-  // Path -> hash mapping
   return Object.fromEntries(hashes);
 };