From e7bc27d10ddaf20f269954eba667d85053998447 Mon Sep 17 00:00:00 2001
From: Sondre Aasemoen <sondre@eons.io>
Date: Wed, 27 Nov 2024 23:44:56 +0100
Subject: [PATCH] Parallelize and batch process files

---
 src/compress.ts | 69 ++++++++++++++++++++++++++++++++++++-------------
 src/index.ts    |  7 +++--
 2 files changed, 56 insertions(+), 20 deletions(-)

diff --git a/src/compress.ts b/src/compress.ts
index 086f400..efe31b2 100644
--- a/src/compress.ts
+++ b/src/compress.ts
@@ -7,6 +7,13 @@ import { createBrotliCompress, createGzip } from "node:zlib";
 
 import * as logger from "./logger.js";
 
+interface CompressionOptions {
+	dir: string;
+	extensions: Array<string>;
+	enabled?: boolean;
+	batchSize?: number;
+}
+
 async function* walkDir(dir: string, extensions: Array<string>): AsyncGenerator<string> {
 	const entries = await readdir(dir, { withFileTypes: true });
 	for (const entry of entries) {
@@ -23,44 +30,70 @@ const filterFile = (file: string, extensions: Array<string>): boolean => {
 	return extensions.some((ext) => extname(file) === ext);
 };
 
-export const gzip = async (dir: string, extensions: Array<string>, enabled?: boolean): Promise<void> => {
+// const compress = async <T>(name: string, compressor: () => T, opts: CompressionOptions): Promise<void> => {};
+
+export const gzip = async (
+	dir: string,
+	extensions: Array<string>,
+	enabled?: boolean,
+	batchSize = 10,
+): Promise<void> => {
 	if (!enabled) {
 		logger.warn("gzip compression disabled, skipping...");
 		return;
 	}
 
 	const start = hrtime.bigint();
-
-	let counter = 0;
+	const files = [];
 	for await (const file of walkDir(dir, extensions)) {
-		counter += 1;
-		const source = createReadStream(file);
-		const destination = createWriteStream(`${file}.gz`);
-		const gzip = createGzip({ level: 9 });
-		await stream.pipeline(source, gzip, destination);
+		files.push(file);
+	}
+
+	for (let i = 0; i < files.length; i += batchSize) {
+		const batch = files.slice(i, i + batchSize);
+		await Promise.all(
+			batch.map(async (path) => {
+				const source = createReadStream(path);
+				const destination = createWriteStream(`${path}.br`);
+				const brotli = createGzip({ level: 9 });
+				await stream.pipeline(source, brotli, destination);
+			}),
+		);
 	}
 
 	const end = hrtime.bigint();
-	logger.success(`finished gzip of ${counter} files in ${(end - start) / BigInt(1000000)}ms`);
+	logger.success(`finished gzip of ${files.length} files in ${(end - start) / BigInt(1000000)}ms`);
 };
 
-export const brotli = async (dir: string, extensions: Array<string>, enabled?: boolean): Promise<void> => {
+export const brotli = async (
+	dir: string,
+	extensions: Array<string>,
+	enabled?: boolean,
+	batchSize = 10,
+): Promise<void> => {
 	if (!enabled) {
 		logger.warn("brotli compression disabled, skipping...");
 		return;
 	}
 
 	const start = hrtime.bigint();
-
-	let counter = 0;
+	const files = [];
 	for await (const file of walkDir(dir, extensions)) {
-		counter += 1;
-		const source = createReadStream(file);
-		const destination = createWriteStream(`${file}.br`);
-		const brotli = createBrotliCompress();
-		await stream.pipeline(source, brotli, destination);
+		files.push(file);
+	}
+
+	for (let i = 0; i < files.length; i += batchSize) {
+		const batch = files.slice(i, i + batchSize);
+		await Promise.all(
+			batch.map(async (path) => {
+				const source = createReadStream(path);
+				const destination = createWriteStream(`${path}.br`);
+				const brotli = createBrotliCompress();
+				await stream.pipeline(source, brotli, destination);
+			}),
+		);
 	}
 
 	const end = hrtime.bigint();
-	logger.success(`finished brotli of ${counter} files in ${(end - start) / BigInt(1000000)}ms`);
+	logger.success(`finished brotli of ${files.length} files in ${(end - start) / BigInt(1000000)}ms`);
 };
diff --git a/src/index.ts b/src/index.ts
index fdee1c2..eb93df8 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -14,12 +14,15 @@ interface Options {
 	brotli?: boolean;
 	/** Extensions to compress, must be in the format `.html`, `.css` etc */
 	fileExtensions?: Array<string>;
+	/** Number of files to batch process */
+	batchSize?: number;
 }
 
 const defaultOptions: Required<Options> = {
 	gzip: true,
 	brotli: true,
 	fileExtensions: defaultFileExtensions,
+	batchSize: 10,
 };
 
 export default function (opts: Options = defaultOptions): AstroIntegration {
@@ -31,8 +34,8 @@ export default function (opts: Options = defaultOptions): AstroIntegration {
 			"astro:build:done": async ({ dir }) => {
 				const path = fileURLToPath(dir);
 				await Promise.allSettled([
-					gzip(path, options.fileExtensions, options.gzip),
-					brotli(path, options.fileExtensions, options.brotli),
+					gzip(path, options.fileExtensions, options.gzip, options.batchSize),
+					brotli(path, options.fileExtensions, options.brotli, options.batchSize),
 				]);
 				logger.success("Compression finished\n");
 			},