cli: refactor worker thread helper to separate queue and plain threads

Signed-off-by: Patrik Oldsberg <poldsberg@gmail.com>

packages/cli/src/lib/parallel.test.ts

@@ -18,6 +18,7 @@ import {
   parseParallelismOption,
   getEnvironmentParallelism,
   runParallelWorkers,
+  runWorkerQueueThreads,
   runWorkerThreads,
 } from './parallel';
 
@@ -144,16 +145,16 @@ describe('runParallelWorkers', () => {
   });
 });
 
-describe('runWorkerThreads', () => {
+describe('runWorkerQueueThreads', () => {
   it('should execute work in parallel', async () => {
     const sharedData = new SharedArrayBuffer(10);
     const sharedView = new Uint8Array(sharedData);
 
-    const results = await runWorkerThreads({
+    const results = await runWorkerQueueThreads({
       threadCount: 4,
       workerData: sharedData,
       items: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
-      workerFactorySource: data => {
+      workerFactory: data => {
         const view = new Uint8Array(data);
 
         return async (i: number) => {
@@ -173,7 +174,7 @@ describe('runWorkerThreads', () => {
     const [result] = await runWorkerThreads({
       threadCount: 1,
       workerData: 'foo',
-      workerFactorySource: data => async () => `${data}bar`,
+      worker: async data => `${data}bar`,
     });
 
     expect(result).toBe('foobar');
@@ -182,7 +183,7 @@ describe('runWorkerThreads', () => {
   it('should run multiple threads without items', async () => {
     const results = await runWorkerThreads({
       threadCount: 4,
-      workerFactorySource: () => async () => 'foo',
+      worker: async () => 'foo',
     });
 
     expect(results).toEqual(['foo', 'foo', 'foo', 'foo']);

packages/cli/src/lib/parallel.ts

@@ -114,8 +114,8 @@ type WorkerThreadMessage =
       error: ErrorLike;
     };
 
-function workerThread(
-  workerFuncFactory: (data: unknown) => (item: unknown) => Promise<void>,
+function workerQueueThread(
+  workerFuncFactory: (data: unknown) => (item: unknown) => Promise<unknown>,
 ) {
   const { parentPort, workerData } = require('worker_threads');
   const workerFunc = workerFuncFactory(workerData);
@@ -134,14 +134,7 @@ function workerThread(
           result,
         });
       } catch (error) {
-        parentPort.postMessage({
-          type: 'error',
-          error: {
-            name: error.name,
-            message: error.message,
-            stack: error.stack,
-          },
-        });
+        parentPort.postMessage({ type: 'error', error });
       }
     }
   });
@@ -149,9 +142,26 @@ function workerThread(
   parentPort.postMessage({ type: 'start' });
 }
 
-type WorkerThreadsOptions<TItem, TResult, TData> = {
-  items?: Iterable<TItem>;
-  workerData?: TData;
+function workerThread(workerFunc: (data: unknown) => Promise<unknown>) {
+  const { parentPort, workerData } = require('worker_threads');
+
+  workerFunc(workerData).then(
+    result => {
+      parentPort.postMessage({
+        type: 'result',
+        index: 0,
+        result,
+      });
+    },
+    error => {
+      parentPort.postMessage({ type: 'error', error });
+    },
+  );
+}
+
+type WorkerQueueThreadsOptions<TItem, TResult, TData> = {
+  /** The items to process */
+  items: Iterable<TItem>;
   /**
    * A function that will be called within each worker thread at startup,
    * which should return the worker function that will be called for each item.
@@ -165,26 +175,27 @@ type WorkerThreadsOptions<TItem, TResult, TData> = {
    * note that they are both copied by value into the worker thread, except for
    * types that are explicitly shareable across threads, such as `SharedArrayBuffer`.
    */
-  workerFactorySource: (data: TData) => (item: TItem) => Promise<TResult>;
+  workerFactory: (data: TData) => (item: TItem) => Promise<TResult>;
+  /** Data supplied to each worker factory */
+  workerData?: TData;
   /** Number of threads, defaults to the environment parallelism */
   threadCount?: number;
 };
 
 /**
  * Spawns one or more worker threads using the `worker_threads` module.
+ * Each thread processes one item at a time from the provided `options.items`.
  */
-export async function runWorkerThreads<TItem, TResult, TData>(
-  options: WorkerThreadsOptions<TItem, TResult, TData>,
+export async function runWorkerQueueThreads<TItem, TResult, TData>(
+  options: WorkerQueueThreadsOptions<TItem, TResult, TData>,
 ): Promise<TResult[]> {
   const {
-    workerFactorySource,
+    workerFactory,
     workerData,
     threadCount = getEnvironmentParallelism(),
   } = options;
 
-  const iterator = (options.items ?? Array(threadCount).fill(undefined))[
-    Symbol.iterator
-  ]();
+  const iterator = options.items[Symbol.iterator]();
   const results = new Array<TResult>();
   let itemIndex = 0;
 
@@ -193,7 +204,7 @@ export async function runWorkerThreads<TItem, TResult, TData>(
       .fill(0)
       .map(async () => {
         const thread = new Worker(
-          `(${workerThread})((${workerFactorySource}))`,
+          `(${workerQueueThread})((${workerFactory}))`,
           {
             eval: true,
             workerData,
@@ -239,3 +250,60 @@ export async function runWorkerThreads<TItem, TResult, TData>(
 
   return results;
 }
+
+type WorkerThreadsOptions<TResult, TData> = {
+  /**
+   * A function that is called by each worker thread to produce a result.
+   *
+   * This function must be defined as an arrow function or using the
+   * function keyword, and must be entirely self contained, not referencing
+   * any variables outside of its scope. This is because the function source
+   * is stringified and evaluated in the worker thread.
+   *
+   * To pass data to the worker, use the `workerData` option, but
+   * note that they are both copied by value into the worker thread, except for
+   * types that are explicitly shareable across threads, such as `SharedArrayBuffer`.
+   */
+  worker: (data: TData) => Promise<TResult>;
+  /** Data supplied to each worker */
+  workerData?: TData;
+  /** Number of threads, defaults to 1 */
+  threadCount?: number;
+};
+
+/**
+ * Spawns one or more worker threads using the `worker_threads` module.
+ */
+export async function runWorkerThreads<TResult, TData>(
+  options: WorkerThreadsOptions<TResult, TData>,
+): Promise<TResult[]> {
+  const { worker, workerData, threadCount = 1 } = options;
+
+  return Promise.all(
+    Array(threadCount)
+      .fill(0)
+      .map(async () => {
+        const thread = new Worker(`(${workerThread})((${worker}))`, {
+          eval: true,
+          workerData,
+        });
+
+        return new Promise<TResult>((resolve, reject) => {
+          thread.on('message', (message: WorkerThreadMessage) => {
+            if (message.type === 'result') {
+              resolve(message.result as TResult);
+            } else if (message.type === 'error') {
+              reject(message.error);
+            }
+          });
+
+          thread.on('error', reject);
+          thread.on('exit', (code: number) => {
+            reject(
+              new Error(`Unexpected worker thread exit with code ${code}`),
+            );
+          });
+        });
+      }),
+  );
+}