From 1ac6871673a40ce37715051f04258f15c80a56e7 Mon Sep 17 00:00:00 2001 From: Patrik Oldsberg Date: Tue, 8 Feb 2022 13:11:08 +0100 Subject: [PATCH] cli: refactor worker thread helper to separate queue and plain threads Signed-off-by: Patrik Oldsberg --- packages/cli/src/lib/parallel.test.ts | 11 +-- packages/cli/src/lib/parallel.ts | 110 +++++++++++++++++++++----- 2 files changed, 95 insertions(+), 26 deletions(-) diff --git a/packages/cli/src/lib/parallel.test.ts b/packages/cli/src/lib/parallel.test.ts index 872738bd60..6c7e7387d7 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, + 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']); diff --git a/packages/cli/src/lib/parallel.ts b/packages/cli/src/lib/parallel.ts index f838d202a7..3e9479407d 100644 --- a/packages/cli/src/lib/parallel.ts +++ b/packages/cli/src/lib/parallel.ts @@ -114,8 +114,8 @@ type WorkerThreadMessage = error: ErrorLike; }; -function workerThread( - workerFuncFactory: (data: unknown) => (item: unknown) => Promise, +function workerQueueThread( + workerFuncFactory: (data: unknown) => (item: unknown) => Promise, ) { 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 = { - items?: Iterable; - workerData?: TData; +function workerThread(workerFunc: (data: unknown) => Promise) { + 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 = { + /** The items to process */ + items: Iterable; /** * 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 = { * 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; + workerFactory: (data: TData) => (item: TItem) => Promise; + /** 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( - options: WorkerThreadsOptions, +export async function runWorkerQueueThreads( + options: WorkerQueueThreadsOptions, ): Promise { 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(); let itemIndex = 0; @@ -193,7 +204,7 @@ export async function runWorkerThreads( .fill(0) .map(async () => { const thread = new Worker( - `(${workerThread})((${workerFactorySource}))`, + `(${workerQueueThread})((${workerFactory}))`, { eval: true, workerData, @@ -239,3 +250,60 @@ export async function runWorkerThreads( return results; } + +type WorkerThreadsOptions = { + /** + * 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; + /** 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( + options: WorkerThreadsOptions, +): Promise { + 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((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}`), + ); + }); + }); + }), + ); +}