Merge pull request #4 from thinknathan/thread-fix

Thread fix

Author: thinknathan

package.json

@@ -1,6 +1,6 @@
 {
 	"name": "node-simple-pixelate",
-	"version": "1.1.5",
+	"version": "1.1.6",
 	"description": "Pixelizes images to create pixel art.",
 	"repository": {
 		"type": "git",

px.cjs

@@ -5,9 +5,10 @@ const yargs = require('yargs');
 const os = require('os');
 const processImage_1 = require('./utils/processImage');
 const processPath_1 = require('./utils/processPath');
-function main() {
+async function main() {
+	// console.time('Done in');
 	// Parse command line arguments
-	const options = yargs
+	const options = await yargs(process.argv.slice(2))
 		.option('f', {
 			alias: 'filename',
 			describe: 'Input image filename',
@@ -250,7 +251,8 @@ function main() {
 				}
 				return Math.round(value);
 			},
-		}).argv;
+		})
+		.parse();
 	if (options.filename) {
 		// Process a single image
 		(0, processImage_1.processImage)(options);
@@ -263,12 +265,13 @@ function main() {
 			console.error(err);
 		}
 		numCores = Math.max(numCores - 1, 1); // Min 1
-		numCores = Math.min(numCores, 16); // Max 16
-		(0, processPath_1.processPath)(options.folderPath, options, numCores);
+		numCores = Math.min(numCores, 32); // Max 32
+		await (0, processPath_1.processPath)(options.folderPath, options, numCores);
 	} else {
 		console.error(
 			'Error: Requires either `filename` or `folderPath`. Run `px --help` for help.',
 		);
 	}
+	// console.timeEnd('Done in');
 }
 main();

src/px.ts

@@ -6,9 +6,11 @@ import * as os from 'os';
 import { processImage } from './utils/processImage';
 import { processPath } from './utils/processPath';
 
-function main() {
+async function main() {
+	// console.time('Done in');
+
 	// Parse command line arguments
-	const options = yargs
+	const options = (await yargs(process.argv.slice(2))
 		.option('f', {
 			alias: 'filename',
 			describe: 'Input image filename',
@@ -252,7 +254,8 @@ function main() {
 				}
 				return Math.round(value);
 			},
-		}).argv as unknown as Options;
+		})
+		.parse()) as unknown as Options;
 
 	if (options.filename) {
 		// Process a single image
@@ -266,13 +269,14 @@ function main() {
 			console.error(err);
 		}
 		numCores = Math.max(numCores - 1, 1); // Min 1
-		numCores = Math.min(numCores, 16); // Max 16
-		processPath(options.folderPath, options, numCores);
+		numCores = Math.min(numCores, 32); // Max 32
+		await processPath(options.folderPath, options, numCores);
 	} else {
 		console.error(
 			'Error: Requires either `filename` or `folderPath`. Run `px --help` for help.',
 		);
 	}
+	// console.timeEnd('Done in');
 }
 
 main();

src/utils/processImage.ts

@@ -1,7 +1,7 @@
 import * as Jimp from 'jimp';
 import * as fs from 'fs';
 import * as path from 'path';
-import { workerData, isMainThread } from 'worker_threads';
+import { parentPort, isMainThread } from 'worker_threads';
 
 import { definedPalettes } from './definedPalettes';
 import { applyAlphaThreshold } from './applyAlphaThreshold';
@@ -18,14 +18,18 @@ function errorCallback(err: unknown) {
 	}
 }
 
+/**
+ * Called on a worker thread to signal current work is complete
+ */
+const workerIsDone = () => parentPort?.postMessage('complete');
+
 /**
  * Processes the given image with various image manipulation options.
  *
  * @param options - Image processing options.
  * @param skipExtCheck - (Optional) Skips extension check if set to true.
  */
 export function processImage(options: Options, skipExtCheck?: boolean): void {
-	console.time('Done in');
 	const { filename } = options;
 	Jimp.read(filename!)
 		.then((image) => {
@@ -204,14 +208,26 @@ function continueProcessing(image: Jimp, options: Options): void {
 	}
 	outputFilename = `${outputFilename}.png`;
 
-	image.write(outputFilename, errorCallback);
-	console.log(`Image saved: ${outputFilename}`);
-	console.timeEnd('Done in');
+	image
+		.writeAsync(outputFilename)
+		.then(() => {
+			console.log(`Image saved: ${outputFilename}`);
+
+			if (!isMainThread) {
+				workerIsDone();
+			}
+		})
+		.catch(errorCallback);
 }
 
 // If used as a worker thread, get file name from message
 if (!isMainThread) {
-	const { filePath, options } = workerData;
-	options.filename = filePath;
-	processImage(options, true);
+	parentPort?.on(
+		'message',
+		async (message: { filePath: string; options: Options }) => {
+			const { filePath, options } = message;
+			options.filename = filePath;
+			processImage(options, true);
+		},
+	);
 }

src/utils/processPath.ts

@@ -9,7 +9,7 @@ export async function processPath(
 	directoryPath: string,
 	options: Options,
 	maxWorkers: number,
-): Promise<void> {
+): Promise<boolean> {
 	const workerPool = new WorkerPool(maxWorkers);
 
 	try {
@@ -24,10 +24,13 @@ export async function processPath(
 				workerPool.addTask(filePath, options);
 			}
 		}
-
-		// Wait for all tasks to complete before exiting
-		workerPool.waitForCompletion();
 	} catch (err) {
 		console.error(`Error reading directory: ${directoryPath}`, err);
 	}
+
+	await workerPool.allComplete();
+
+	workerPool.exitAll();
+
+	return true;
 }

src/utils/workerPool.ts

@@ -1,13 +1,48 @@
 import { Worker } from 'worker_threads';
 import * as path from 'path';
 
+type TWorker = Worker & { isIdle: boolean };
+
 /**
  * Manages a pool of worker threads for parallel processing of image files.
  */
 export class WorkerPool {
-	private workers: Worker[] = [];
+	private workers: TWorker[] = [];
 	private taskQueue: { filePath: string; options: Options }[] = [];
 	private maxWorkers: number;
+	private completePromise?: Promise<void>;
+	private completeResolve?: () => void;
+	private isComplete(): boolean {
+		return (
+			this.taskQueue.length === 0 &&
+			this.workers.every((worker) => worker.isIdle)
+		);
+	}
+
+	/**
+	 * Terminate all workers in the pool.
+	 */
+	public exitAll(): void {
+		this.workers.forEach((worker) => worker.terminate());
+		this.workers = [];
+	}
+
+	/**
+	 * Returns a promise that resolves when all work is done.
+	 */
+	public async allComplete(): Promise<void> {
+		if (this.isComplete()) {
+			return Promise.resolve();
+		}
+
+		if (!this.completePromise) {
+			this.completePromise = new Promise<void>((resolve) => {
+				this.completeResolve = resolve;
+			});
+		}
+
+		return this.completePromise;
+	}
 
 	/**
 	 * Creates a new WorkerPool instance.
@@ -25,18 +60,22 @@ export class WorkerPool {
 	 * @param options - Image processing options for the file.
 	 */
 	private createWorker(filePath: string, options: Options): void {
-		const worker = new Worker(path.join(__dirname, 'processImage.js'), {
-			workerData: { filePath, options },
-		});
+		const worker = new Worker(
+			path.join(__dirname, 'processImage.js'),
+		) as TWorker;
+
+		worker.isIdle = false;
+		worker.postMessage({ filePath, options });
 
 		// Listen for messages and errors from the worker
-		worker.on('message', (message) => {
-			console.log(message);
+		worker.on('message', () => {
+			worker.isIdle = true;
 			this.processNextTask();
 		});
 
 		worker.on('error', (err) => {
 			console.error(`Error in worker for file ${filePath}:`, err);
+			worker.isIdle = true;
 			this.processNextTask();
 		});
 
@@ -49,7 +88,22 @@ export class WorkerPool {
 	private processNextTask(): void {
 		const nextTask = this.taskQueue.shift();
 		if (nextTask) {
-			this.createWorker(nextTask.filePath, nextTask.options);
+			if (this.workers.length < this.maxWorkers) {
+				this.createWorker(nextTask.filePath, nextTask.options);
+			} else {
+				const worker = this.workers.find((w) => w.isIdle);
+				if (worker) {
+					worker.isIdle = false;
+					worker.postMessage(nextTask);
+				} else {
+					// Something went wrong, there are no idle workers somehow
+					throw Error('Could not find an idle worker.');
+				}
+			}
+		} else if (this.isComplete() && this.completeResolve) {
+			this.completeResolve();
+			this.completePromise = undefined;
+			this.completeResolve = undefined;
 		}
 	}
 
@@ -66,15 +120,4 @@ export class WorkerPool {
 			this.taskQueue.push({ filePath, options });
 		}
 	}
-
-	/**
-	 * Waits for all tasks to complete before exiting.
-	 */
-	public waitForCompletion(): void {
-		this.workers.forEach((worker) => {
-			worker.on('exit', () => {
-				this.processNextTask();
-			});
-		});
-	}
 }

utils/processImage.js

@@ -18,14 +18,17 @@ function errorCallback(err) {
         console.error(err);
     }
 }
+/**
+ * Called on a worker thread to signal current work is complete
+ */
+const workerIsDone = () => worker_threads_1.parentPort?.postMessage('complete');
 /**
  * Processes the given image with various image manipulation options.
  *
  * @param options - Image processing options.
  * @param skipExtCheck - (Optional) Skips extension check if set to true.
  */
 function processImage(options, skipExtCheck) {
-    console.time('Done in');
     const { filename } = options;
     Jimp.read(filename)
         .then((image) => {
@@ -169,13 +172,21 @@ function continueProcessing(image, options) {
         outputFilename = `${outputFilename}-z_${pixelSize}`;
     }
     outputFilename = `${outputFilename}.png`;
-    image.write(outputFilename, errorCallback);
-    console.log(`Image saved: ${outputFilename}`);
-    console.timeEnd('Done in');
+    image
+        .writeAsync(outputFilename)
+        .then(() => {
+        console.log(`Image saved: ${outputFilename}`);
+        if (!worker_threads_1.isMainThread) {
+            workerIsDone();
+        }
+    })
+        .catch(errorCallback);
 }
 // If used as a worker thread, get file name from message
 if (!worker_threads_1.isMainThread) {
-    const { filePath, options } = worker_threads_1.workerData;
-    options.filename = filePath;
-    processImage(options, true);
+    worker_threads_1.parentPort?.on('message', async (message) => {
+        const { filePath, options } = message;
+        options.filename = filePath;
+        processImage(options, true);
+    });
 }

utils/processPath.js

@@ -19,11 +19,12 @@ async function processPath(directoryPath, options, maxWorkers) {
                 workerPool.addTask(filePath, options);
             }
         }
-        // Wait for all tasks to complete before exiting
-        workerPool.waitForCompletion();
     }
     catch (err) {
         console.error(`Error reading directory: ${directoryPath}`, err);
     }
+    await workerPool.allComplete();
+    workerPool.exitAll();
+    return true;
 }
 exports.processPath = processPath;