cli: initial worker thread helper

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

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

@@ -18,6 +18,7 @@ import {
   parseParallelismOption,
   getEnvironmentParallelism,
   runParallelWorkers,
+  runWorkerThreads,
 } from './parallel';
 
 describe('parseParallelismOption', () => {
@@ -142,3 +143,48 @@ describe('runParallelWorkers', () => {
     expect(done).toEqual([0, 1, 2, 3, 4]);
   });
 });
+
+describe('runWorkerThreads', () => {
+  it('should execute work in parallel', async () => {
+    const sharedData = new SharedArrayBuffer(10);
+    const sharedView = new Uint8Array(sharedData);
+
+    const results = await runWorkerThreads({
+      threadCount: 4,
+      workerData: sharedData,
+      items: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
+      workerFactorySource: data => {
+        const view = new Uint8Array(data);
+
+        return async (i: number) => {
+          view[i] = 10 + i;
+          return 20 + i;
+        };
+      },
+    });
+
+    expect(Array.from(sharedView)).toEqual([
+      10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
+    ]);
+    expect(results).toEqual([20, 21, 22, 23, 24, 25, 26, 27, 28, 29]);
+  });
+
+  it('should run a single thread without items', async () => {
+    const [result] = await runWorkerThreads({
+      threadCount: 1,
+      workerData: 'foo',
+      workerFactorySource: data => async () => `${data}bar`,
+    });
+
+    expect(result).toBe('foobar');
+  });
+
+  it('should run multiple threads without items', async () => {
+    const results = await runWorkerThreads({
+      threadCount: 4,
+      workerFactorySource: () => async () => 'foo',
+    });
+
+    expect(results).toEqual(['foo', 'foo', 'foo', 'foo']);
+  });
+});

packages/cli/src/lib/parallel.ts

@@ -14,6 +14,9 @@
  * limitations under the License.
  */
 
+import { ErrorLike } from '@backstage/errors';
+import { Worker } from 'worker_threads';
+
 export const DEFAULT_PARALLELISM = 4;
 
 export const PARALLEL_ENV_VAR = 'BACKSTAGE_CLI_BUILD_PARALLEL';
@@ -88,3 +91,151 @@ export async function runParallelWorkers<TItem>(
       }),
   );
 }
+
+type WorkerThreadMessage =
+  | {
+      type: 'done';
+    }
+  | {
+      type: 'item';
+      index: number;
+      item: unknown;
+    }
+  | {
+      type: 'start';
+    }
+  | {
+      type: 'result';
+      index: number;
+      result: unknown;
+    }
+  | {
+      type: 'error';
+      error: ErrorLike;
+    };
+
+function workerThread(
+  workerFuncFactory: (data: unknown) => (item: unknown) => Promise<void>,
+) {
+  const { parentPort, workerData } = require('worker_threads');
+  const workerFunc = workerFuncFactory(workerData);
+
+  parentPort.on('message', async (message: WorkerThreadMessage) => {
+    if (message.type === 'done') {
+      parentPort.close();
+      return;
+    }
+    if (message.type === 'item') {
+      try {
+        const result = await workerFunc(message.item);
+        parentPort.postMessage({
+          type: 'result',
+          index: message.index,
+          result,
+        });
+      } catch (error) {
+        parentPort.postMessage({
+          type: 'error',
+          error: {
+            name: error.name,
+            message: error.message,
+            stack: error.stack,
+          },
+        });
+      }
+    }
+  });
+
+  parentPort.postMessage({ type: 'start' });
+}
+
+type WorkerThreadsOptions<TItem, TResult, TData> = {
+  items?: Iterable<TItem>;
+  workerData?: TData;
+  /**
+   * 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.
+   *
+   * 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 and `items`, 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`.
+   */
+  workerFactorySource: (data: TData) => (item: TItem) => Promise<TResult>;
+  /** Number of threads, defaults to the environment parallelism */
+  threadCount?: number;
+};
+
+/**
+ * Spawns one or more worker threads using the `worker_threads` module.
+ */
+export async function runWorkerThreads<TItem, TResult, TData>(
+  options: WorkerThreadsOptions<TItem, TResult, TData>,
+): Promise<TResult[]> {
+  const {
+    workerFactorySource,
+    workerData,
+    threadCount = getEnvironmentParallelism(),
+  } = options;
+
+  const iterator = (options.items ?? Array(threadCount).fill(undefined))[
+    Symbol.iterator
+  ]();
+  const results = new Array<TResult>();
+  let itemIndex = 0;
+
+  await Promise.all(
+    Array(threadCount)
+      .fill(0)
+      .map(async () => {
+        const thread = new Worker(
+          `(${workerThread})((${workerFactorySource}))`,
+          {
+            eval: true,
+            workerData,
+          },
+        );
+
+        return new Promise<void>((resolve, reject) => {
+          thread.on('message', (message: WorkerThreadMessage) => {
+            if (message.type === 'start' || message.type === 'result') {
+              if (message.type === 'result') {
+                results[message.index] = message.result as TResult;
+              }
+              const { value, done } = iterator.next();
+              if (done) {
+                thread.postMessage({ type: 'done' });
+              } else {
+                thread.postMessage({
+                  type: 'item',
+                  index: itemIndex,
+                  item: value,
+                });
+                itemIndex += 1;
+              }
+            } else if (message.type === 'error') {
+              const error = new Error(message.error.message);
+              error.name = message.error.name;
+              error.stack = message.error.stack;
+              reject(error);
+            }
+          });
+
+          thread.on('error', reject);
+          thread.on('exit', (code: number) => {
+            if (code !== 0) {
+              reject(new Error(`Worker thread exited with code ${code}`));
+            } else {
+              resolve();
+            }
+          });
+        });
+      }),
+  );
+
+  return results;
+}