[WIP] Fix add batch processing

packages/webpack-stats-differ/src/diff.ts

@@ -19,7 +19,12 @@ import {
   GenerateDiffAsyncResult,
   GetFileDiffOptions,
   getFileDiffResult,
+  processPairsInBatches,
 } from "./generateDiffsAsync";
+import { parser } from "stream-json";
+import { streamArray } from "stream-json/streamers/StreamArray";
+import { chain } from "stream-chain";
+import { createReadStream } from "fs";
 
 /**
  * Calculates the diff between two sets of bundle stats
@@ -44,7 +49,7 @@ export async function diff(
     candidate: string | RemoteArtifact[];
     hostUrl: string;
   },
-  bundleStatsOwners?:Map<string, string[]> | undefined,
+  bundleStatsOwners?: Map<string, string[]> | undefined,
   useWorkers?: boolean
 ): Promise<FileDiffResults> {
   const [filesA, filesB] = await Promise.all(
@@ -94,44 +99,53 @@ const generateDiffs = async (
 ): Promise<FileDiffResults> => {
   let results: GenerateDiffAsyncResult[] = [];
   const numOfCpus = os.cpus().length;
-  const numOfWorkers = useWorkers ? Math.ceil(numOfCpus * 0.75) : 1;
+  // Limit workers to prevent memory issues
+  const numOfWorkers = useWorkers ? Math.min(8, Math.ceil(numOfCpus * 0.5)) : 1;
   const start = Date.now();
-  console.log(
-    `Started diffing for ${filePairs.length} bundle pairs using ${numOfWorkers} workers`
-  );
+
   if (numOfWorkers <= 1) {
-    for (const pair of filePairs) {
-      const result = await getFileDiffResult({ pair, filter });
-      results.push(result);
-    }
+    console.log(`Processing ${filePairs.length} bundle pairs in batches without workers`);
+    results = await processPairsInBatches(filePairs, filter);
   } else {
-    let pairIndex = 0;
-    const diffImageWorker = async () => {
-      const worker = new Worker(`${__dirname}/generateDiffsAsync.js`);
+    console.log(`Processing ${filePairs.length} bundle pairs in batches using ${numOfWorkers} workers`);
+
+    // Split pairs into batches for each worker
+    const batchesPerWorker = Math.ceil(filePairs.length / numOfWorkers);
+    const workerBatches = Array.from({ length: numOfWorkers }, (_, i) => 
+      filePairs.slice(i * batchesPerWorker, (i + 1) * batchesPerWorker)
+    ).filter(batch => batch.length > 0);
 
-      const generateDiffOptions: GetFileDiffOptions = {
-        pair: null,
-        filter,
-      };
+    const diffWorker = async (workerPairs: FilePair[]) => {
+      const worker = new Worker(`${__dirname}/generateDiffsAsync.js`);
+      const workerResults: GenerateDiffAsyncResult[] = [];
 
-      while ((generateDiffOptions.pair = filePairs[pairIndex++])) {
-        await new Promise((resolve) => {
-          worker.once("message", (result: GenerateDiffAsyncResult) => {
-            results.push(result);
-            resolve(undefined);
+      for (let i = 0; i < workerPairs.length; i += 10) {
+        const batch = workerPairs.slice(i, i + 10);
+        console.log(`Worker processing batch ${Math.floor(i/10) + 1} of ${Math.ceil(workerPairs.length/10)} (${batch.length} pairs)`);
+
+        for (const pair of batch) {
+          const generateDiffOptions: GetFileDiffOptions = { pair, filter };
+          const result = await new Promise<GenerateDiffAsyncResult>((resolve) => {
+            worker.once("message", (result: GenerateDiffAsyncResult) => {
+              resolve(result);
+            });
+            worker.postMessage(generateDiffOptions);
           });
-
-          worker.postMessage(generateDiffOptions);
-        });
+          workerResults.push(result);
+        }
+
+        // Small delay between batches to allow for GC
+        await new Promise(resolve => setTimeout(resolve, 100));
       }
 
       worker.terminate();
-      return;
+      return workerResults;
     };
 
-    const buckets = new Array(numOfWorkers).fill(1);
-    await Promise.all(buckets.map(diffImageWorker));
+    const workerResults = await Promise.all(workerBatches.map(diffWorker));
+    results = workerResults.flat();
   }
+
   console.log(`Diffing done in ${Date.now() - start} ms`);
   return {
     withDifferences: results
@@ -151,13 +165,28 @@ const generateDiffs = async (
 const getRemoteArtifactsManifest = async (
   filePath: string
 ): Promise<RemoteArtifact[]> => {
-  const data = await fs.promises.readFile(filePath, { encoding: "utf-8" });
-  try {
-    const parsed: RemoteArtifact[] = JSON.parse(data);
-    return parsed;
-  } catch (e: unknown) {
-    throw new Error(`Cannot parse remote artifact manifest ${filePath}: ${e}`);
-  }
+  return new Promise((resolve, reject) => {
+    const artifacts: RemoteArtifact[] = [];
+    const pipeline = chain([
+      createReadStream(filePath, { encoding: "utf8" }),
+      parser(),
+      streamArray(),
+    ]);
+
+    pipeline.on("data", (data) => {
+      artifacts.push(data.value);
+    });
+
+    pipeline.on("end", () => {
+      resolve(artifacts);
+    });
+
+    pipeline.on("error", (err) => {
+      reject(
+        new Error(`Cannot parse remote artifact manifest ${filePath}: ${err}`)
+      );
+    });
+  });
 };
 
 const instanceOfRemoteArtifact = (obj: any): obj is RemoteArtifact[] =>

packages/webpack-stats-differ/src/generateDiffsAsync.ts

@@ -13,7 +13,7 @@ import { chain } from "stream-chain";
 import { parser } from "stream-json";
 import { streamObject } from "stream-json/streamers/StreamObject";
 
-const getWebpackStatJSON = async (
+export const getWebpackStatJSON = async (
   filePath: string
 ): Promise<WebpackStatsJson> => {
   try {
@@ -22,24 +22,7 @@ const getWebpackStatJSON = async (
     throw new Error(`Cannot access webpack stats file at ${filePath}: ${e}`);
   }
 
-  // Get file size
-  const stats = await fs.promises.stat(filePath);
-  const fileSizeInMB = stats.size / (1024 * 1024);
-
-  // For small files, use the regular JSON.parse method
-  if (fileSizeInMB < 100) {
-    // Adjust threshold as needed
-    try {
-      const parsed: WebpackStatsJson = JSON.parse(
-        await fs.promises.readFile(filePath, { encoding: "utf-8" })
-      );
-      return parsed;
-    } catch (e: unknown) {
-      throw new Error(`Cannot parse webpack state file ${filePath}: ${e}`);
-    }
-  }
-
-  // For large files, use streaming approach
+  // Always use streaming approach for webpack stats files since they can be very large
   return new Promise<WebpackStatsJson>((resolve, reject) => {
     const result: WebpackStatsJson = {} as WebpackStatsJson;
 
@@ -141,6 +124,43 @@ export const getFileDiffResult = async ({
   }
 };
 
+// Add batch processing utilities
+const BATCH_SIZE = 10; // Process 10 pairs at a time
+const MAX_WORKERS = 8; // Limit number of workers to prevent memory issues
+
+export const processBatch = async (
+  pairs: FilePair[],
+  filter?: string | string[]
+): Promise<GenerateDiffAsyncResult[]> => {
+  const results: GenerateDiffAsyncResult[] = [];
+  for (const pair of pairs) {
+    const result = await getFileDiffResult({ pair, filter });
+    results.push(result);
+  }
+  return results;
+};
+
+export const processPairsInBatches = async (
+  pairs: FilePair[],
+  filter?: string | string[]
+): Promise<GenerateDiffAsyncResult[]> => {
+  const results: GenerateDiffAsyncResult[] = [];
+
+  for (let i = 0; i < pairs.length; i += BATCH_SIZE) {
+    const batch = pairs.slice(i, i + BATCH_SIZE);
+    console.log(`Processing batch ${Math.floor(i/BATCH_SIZE) + 1} of ${Math.ceil(pairs.length/BATCH_SIZE)} (${batch.length} pairs)`);
+    const batchResults = await processBatch(batch, filter);
+    results.push(...batchResults);
+
+    // Force garbage collection between batches if available
+    if (global.gc) {
+      global.gc();
+    }
+  }
+
+  return results;
+};
+
 if (!isMainThread) {
   (async () => {
     return new Promise(() => {

packages/webpack-stats-differ/src/index.ts

@@ -8,3 +8,4 @@ export {
   FileDiffResultWithComparisonToolUrl,
 } from "./diff";
 export { getFriendlyAssetName } from "./getFriendlyAssetName";
+export { getWebpackStatJSON } from "./generateDiffsAsync";