From bd7824e8b822d4db0a582c3ba2c4914d306786ad Mon Sep 17 00:00:00 2001
From: Tony Sullivan <tony.f.sullivan@outlook.com>
Date: Thu, 15 Sep 2022 10:01:01 -0500
Subject: [PATCH] perf: allow a thread for each output encoding format

---
 .../image/src/vendor/squoosh/image-pool.ts    | 88 +++++++++++++------
 1 file changed, 60 insertions(+), 28 deletions(-)

diff --git a/packages/integrations/image/src/vendor/squoosh/image-pool.ts b/packages/integrations/image/src/vendor/squoosh/image-pool.ts
index a266128db5..440dd1a5b4 100644
--- a/packages/integrations/image/src/vendor/squoosh/image-pool.ts
+++ b/packages/integrations/image/src/vendor/squoosh/image-pool.ts
@@ -1,4 +1,4 @@
-
+import { cpus } from 'os'
 import { isMainThread } from 'node:worker_threads';
 import WorkerPool from '../../utils/workerPool.js';
 import * as impl from './impl.js';
@@ -16,9 +16,15 @@ type ResizeOperation = {
 }
 export type Operation = RotateOperation | ResizeOperation
 
-// HACK!  Find the right import assuming this works
 const getWorker = execOnce(
-  () => new WorkerPool(6, './node_modules/@astrojs/image/dist/vendor/squoosh/image-pool.js')
+  () => {
+		return new WorkerPool(
+			// There will be at most 7 workers needed since each worker will take
+      // at least 1 operation type.
+      Math.max(1, Math.min(cpus().length - 1, 7)),
+			'./node_modules/@astrojs/image/dist/vendor/squoosh/image-pool.js'
+		);
+	}
 )
 
 type DecodeParams = {
@@ -36,13 +42,26 @@ type RotateParams = {
 	imageData: ImageData,
 	numRotations: number
 };
-type EncodeParams = {
-	operation: 'encode',
+type EncodeAvifParams = {
+	operation: 'encodeavif',
 	imageData: ImageData,
-	format: OutputFormat,
-	quality?: number
+	quality: number
 }
-type JobMessage = DecodeParams | ResizeParams | RotateParams | EncodeParams
+type EncodeJpegParams = {
+	operation: 'encodejpeg',
+	imageData: ImageData,
+	quality: number
+}
+type EncodePngParams = {
+	operation: 'encodepng',
+	imageData: ImageData
+}
+type EncodeWebpParams = {
+	operation: 'encodewebp',
+	imageData: ImageData,
+	quality: number
+}
+type JobMessage = DecodeParams | ResizeParams | RotateParams | EncodeAvifParams | EncodeJpegParams | EncodePngParams | EncodeWebpParams
 
 function handleJob(params: JobMessage) {
   switch (params.operation) {
@@ -52,20 +71,14 @@ function handleJob(params: JobMessage) {
 			return impl.resize({ image: params.imageData as any, width: params.width, height: params.height })
 		case 'rotate':
 			return impl.rotate(params.imageData as any, params.numRotations);
-		case 'encode':
-			switch (params.format) {
-				case 'jpeg':
-				case 'jpg':
-					return impl.encodeJpeg(params.imageData as any, { quality: params.quality || 100 })
-				case 'webp':
-					return impl.encodeWebp(params.imageData as any, { quality: params.quality || 100 })
-				case 'avif':
-					return impl.encodeAvif(params.imageData as any, { quality: params.quality || 100 })
-				case 'png':
-					return impl.encodePng(params.imageData as any)
-				default:
-					throw Error(`Unsupported encoding format`)
-			}
+		case 'encodeavif':
+			return impl.encodeAvif(params.imageData as any, { quality: params.quality })
+		case 'encodejpeg':
+			return impl.encodeJpeg(params.imageData as any, { quality: params.quality })
+		case 'encodepng':
+			return impl.encodePng(params.imageData as any)
+		case 'encodewebp':
+			return impl.encodeWebp(params.imageData as any, { quality: params.quality })
     default:
       throw Error(`Invalid job "${(params as any).operation}"`);
   }
@@ -77,6 +90,7 @@ export async function processBuffer(
   encoding: OutputFormat,
   quality: number
 ): Promise<Buffer> {
+	// @ts-ignore
 	const worker = await getWorker()
 
   let imageData = await worker.dispatchJob({
@@ -100,12 +114,30 @@ export async function processBuffer(
     }
   }
 
-	return await worker.dispatchJob({
-		operation: 'encode',
-		format: encoding,
-		imageData,
-		quality
-	}) as Buffer
+	switch (encoding) {
+		case 'avif':
+			return await worker.dispatchJob({ operation: 'encodeavif', imageData, quality: quality || 100 }) as Buffer;
+		case 'jpeg':
+		case 'jpg':
+			return await worker.dispatchJob({
+				operation: 'encodejpeg',
+				imageData,
+				quality: quality || 100,
+			}) as Buffer;
+		case 'png':
+			return await worker.dispatchJob({
+				operation: 'encodepng',
+				imageData,
+			}) as Buffer;
+		case 'webp':
+			return await worker.dispatchJob({
+				operation: 'encodejpeg',
+				imageData,
+				quality: quality || 100,
+			}) as Buffer;
+		default:
+			throw Error(`Unsupported encoding format`)
+	}
 }
 
 if (!isMainThread) {