diff --git a/packages/cli/src/lib/builder/buildTypeDefinitions.ts b/packages/cli/src/lib/builder/buildTypeDefinitions.ts index 85e54850e4..3bf3a4495f 100644 --- a/packages/cli/src/lib/builder/buildTypeDefinitions.ts +++ b/packages/cli/src/lib/builder/buildTypeDefinitions.ts @@ -16,10 +16,10 @@ import fs from 'fs-extra'; import chalk from 'chalk'; -import { Worker } from 'worker_threads'; import { relative as relativePath, resolve as resolvePath } from 'path'; import { paths } from '../paths'; import { buildTypeDefinitionsWorker } from './buildTypeDefinitionsWorker'; +import { runWorkerThreads } from '../parallel'; // These message types are ignored since we want to avoid duplicating the logic of // handling them correctly, and we already have the API Reports warning about them. @@ -84,59 +84,46 @@ export async function buildTypeDefinitions( const typescriptDir = paths.resolveTargetRoot('node_modules/typescript'); const hasTypescript = await fs.pathExists(typescriptDir); const typescriptCompilerFolder = hasTypescript ? typescriptDir : undefined; - - const worker = new Worker(`(${buildTypeDefinitionsWorker})()`, { - eval: true, + await runWorkerThreads({ + threadCount: 1, workerData: { entryPoints, workerConfigs, typescriptCompilerFolder, }, - }); - - await new Promise((resolve, reject) => { - worker.once('error', reject); - worker.once('exit', code => { - if (code) { - reject(new Error(`Worker exited with code ${code}`)); + worker: buildTypeDefinitionsWorker, + onMessage: ({ + message, + targetTypesDir, + }: { + message: any; + targetTypesDir: string; + }) => { + if (ignoredMessages.has(message.messageId)) { + return; } - }); - worker.on('message', data => { - if (data.type === 'done') { - if (data.error) { - reject(data.error); - } else { - resolve(); - } - } else if (data.type === 'message') { - const { message, targetTypesDir } = data; - if (ignoredMessages.has(message.messageId)) { - return; - } - - let text = `${message.text} (${message.messageId})`; - if (message.sourceFilePath) { - text += ' at '; - text += relativePath(targetTypesDir, message.sourceFilePath); - if (message.sourceFileLine) { - text += `:${message.sourceFileLine}`; - if (message.sourceFileColumn) { - text += `:${message.sourceFileColumn}`; - } + let text = `${message.text} (${message.messageId})`; + if (message.sourceFilePath) { + text += ' at '; + text += relativePath(targetTypesDir, message.sourceFilePath); + if (message.sourceFileLine) { + text += `:${message.sourceFileLine}`; + if (message.sourceFileColumn) { + text += `:${message.sourceFileColumn}`; } } - if (message.logLevel === 'error') { - console.error(chalk.red(`Error: ${text}`)); - } else if ( - message.logLevel === 'warning' || - message.category === 'Extractor' - ) { - console.warn(`Warning: ${text}`); - } else { - console.log(text); - } } - }); + if (message.logLevel === 'error') { + console.error(chalk.red(`Error: ${text}`)); + } else if ( + message.logLevel === 'warning' || + message.category === 'Extractor' + ) { + console.warn(`Warning: ${text}`); + } else { + console.log(text); + } + }, }); } diff --git a/packages/cli/src/lib/builder/buildTypeDefinitionsWorker.ts b/packages/cli/src/lib/builder/buildTypeDefinitionsWorker.ts index 8cb843f58e..43e5cbc526 100644 --- a/packages/cli/src/lib/builder/buildTypeDefinitionsWorker.ts +++ b/packages/cli/src/lib/builder/buildTypeDefinitionsWorker.ts @@ -20,7 +20,10 @@ * be self-contained. * Using TypeScript is fine as it is transpiled before being stringified. */ -export function buildTypeDefinitionsWorker() { +export async function buildTypeDefinitionsWorker( + workerData: any, + sendMessage: (message: any) => void, +) { try { require('@microsoft/api-extractor'); } catch (error) { @@ -31,7 +34,6 @@ export function buildTypeDefinitionsWorker() { } const { dirname } = require('path'); - const { workerData, parentPort } = require('worker_threads'); const { entryPoints, workerConfigs, typescriptCompilerFolder } = workerData; const apiExtractor = require('@microsoft/api-extractor'); @@ -63,7 +65,6 @@ export function buildTypeDefinitionsWorker() { return old.call(this, path); }; - let success = true; let compilerState; for (const { extractorOptions, targetTypesDir } of workerConfigs) { const extractorConfig = ExtractorConfig.prepare(extractorOptions); @@ -82,24 +83,15 @@ export function buildTypeDefinitionsWorker() { showDiagnostics: false, messageCallback: (message: any) => { message.handled = true; - parentPort.postMessage({ type: 'message', message, targetTypesDir }); + sendMessage({ message, targetTypesDir }); }, }); if (!extractorResult.succeeded) { - parentPort.postMessage({ - type: 'done', - error: new Error( - `Type definition build completed with ${extractorResult.errorCount} errors` + - ` and ${extractorResult.warningCount} warnings`, - ), - }); - success = false; - break; + throw new Error( + `Type definition build completed with ${extractorResult.errorCount} errors` + + ` and ${extractorResult.warningCount} warnings`, + ); } } - - if (success) { - parentPort.postMessage({ type: 'done' }); - } } diff --git a/packages/cli/src/lib/parallel.test.ts b/packages/cli/src/lib/parallel.test.ts index d851f33f82..824a83d335 100644 --- a/packages/cli/src/lib/parallel.test.ts +++ b/packages/cli/src/lib/parallel.test.ts @@ -18,6 +18,8 @@ import { parseParallelismOption, getEnvironmentParallelism, runParallelWorkers, + runWorkerQueueThreads, + runWorkerThreads, } from './parallel'; describe('parseParallelismOption', () => { @@ -142,3 +144,68 @@ describe('runParallelWorkers', () => { expect(done).toEqual([0, 1, 2, 3, 4]); }); }); + +describe('runWorkerQueueThreads', () => { + it('should execute work in parallel', async () => { + const sharedData = new SharedArrayBuffer(10); + const sharedView = new Uint8Array(sharedData); + + const results = await runWorkerQueueThreads({ + threadCount: 4, + workerData: sharedData, + items: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], + workerFactory: 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]); + }); +}); + +describe('runWorkerThreads', () => { + it('should run a single thread without items', async () => { + const [result] = await runWorkerThreads({ + threadCount: 1, + workerData: 'foo', + worker: async data => `${data}bar`, + }); + + expect(result).toBe('foobar'); + }); + + it('should run multiple threads without items', async () => { + const results = await runWorkerThreads({ + threadCount: 4, + worker: async () => 'foo', + }); + + expect(results).toEqual(['foo', 'foo', 'foo', 'foo']); + }); + + it('should send messages', async () => { + const messages = new Array(); + + await runWorkerThreads({ + threadCount: 2, + worker: async (_data, sendMessage) => { + sendMessage('foo'); + await new Promise(resolve => setTimeout(resolve, 50)); + sendMessage('bar'); + await new Promise(resolve => setTimeout(resolve, 50)); + sendMessage('baz'); + }, + onMessage: (message: string) => messages.push(message), + }); + + expect(messages).toEqual(['foo', 'foo', 'bar', 'bar', 'baz', 'baz']); + }); +}); diff --git a/packages/cli/src/lib/parallel.ts b/packages/cli/src/lib/parallel.ts index f8af967d22..246020fcc3 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,236 @@ 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; + } + | { + type: 'message'; + message: unknown; + }; + +export 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. + * + * 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`. + */ + 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 runWorkerQueueThreads( + options: WorkerQueueThreadsOptions, +): Promise { + const { + workerFactory, + workerData, + threadCount = getEnvironmentParallelism(), + } = options; + + const iterator = options.items[Symbol.iterator](); + const results = new Array(); + let itemIndex = 0; + + await Promise.all( + Array(threadCount) + .fill(0) + .map(async () => { + const thread = new Worker(`(${workerQueueThread})(${workerFactory})`, { + 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; +} + +function workerQueueThread( + 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 }); + } + } + }); + + parentPort.postMessage({ type: 'start' }); +} + +export 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, + sendMessage: (message: TMessage) => void, + ) => Promise; + /** Data supplied to each worker */ + workerData?: TData; + /** Number of threads, defaults to 1 */ + threadCount?: number; + /** An optional handler for messages posted from the worker thread */ + onMessage?: (message: TMessage) => void; +}; + +/** + * Spawns one or more worker threads using the `worker_threads` module. + */ +export async function runWorkerThreads( + options: WorkerThreadsOptions, +): Promise { + const { worker, workerData, threadCount = 1, onMessage } = 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); + } else if (message.type === 'message') { + onMessage?.(message.message as TMessage); + } + }); + + thread.on('error', reject); + thread.on('exit', (code: number) => { + reject( + new Error(`Unexpected worker thread exit with code ${code}`), + ); + }); + }); + }), + ); +} + +function workerThread( + workerFunc: ( + data: unknown, + sendMessage: (message: unknown) => void, + ) => Promise, +) { + const { parentPort, workerData } = require('worker_threads'); + + const sendMessage = (message: unknown) => { + parentPort.postMessage({ type: 'message', message }); + }; + + workerFunc(workerData, sendMessage).then( + result => { + parentPort.postMessage({ + type: 'result', + index: 0, + result, + }); + }, + error => { + parentPort.postMessage({ type: 'error', error }); + }, + ); +}