From 00f6bcda2a06db6d14fa43dd73658f1dabef856b Mon Sep 17 00:00:00 2001 From: Patrik Oldsberg Date: Tue, 8 Feb 2022 12:33:58 +0100 Subject: [PATCH] cli: initial worker thread helper Signed-off-by: Patrik Oldsberg --- packages/cli/src/lib/parallel.test.ts | 46 ++++++++ packages/cli/src/lib/parallel.ts | 151 ++++++++++++++++++++++++++ 2 files changed, 197 insertions(+) diff --git a/packages/cli/src/lib/parallel.test.ts b/packages/cli/src/lib/parallel.test.ts index d851f33f82..872738bd60 100644 --- a/packages/cli/src/lib/parallel.test.ts +++ b/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']); + }); +}); diff --git a/packages/cli/src/lib/parallel.ts b/packages/cli/src/lib/parallel.ts index f8af967d22..f838d202a7 100644 --- a/packages/cli/src/lib/parallel.ts +++ b/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( }), ); } + +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, +) { + 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 = { + items?: Iterable; + 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; + /** 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( + options: WorkerThreadsOptions, +): Promise { + const { + workerFactorySource, + workerData, + threadCount = getEnvironmentParallelism(), + } = options; + + const iterator = (options.items ?? Array(threadCount).fill(undefined))[ + Symbol.iterator + ](); + const results = new Array(); + let itemIndex = 0; + + await Promise.all( + Array(threadCount) + .fill(0) + .map(async () => { + const thread = new Worker( + `(${workerThread})((${workerFactorySource}))`, + { + eval: true, + workerData, + }, + ); + + return new Promise((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; +}