Merge pull request #9410 from backstage/rugvip/threads

cli: refactoring and more utilities for working with worker threads

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<void>((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);
+      }
+    },
   });
 }

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' });
-  }
 }

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<string>();
+
+    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']);
+  });
+});

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<TItem>(
       }),
   );
 }
+
+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<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.
+   *
+   * 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<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 runWorkerQueueThreads<TItem, TResult, TData>(
+  options: WorkerQueueThreadsOptions<TItem, TResult, TData>,
+): Promise<TResult[]> {
+  const {
+    workerFactory,
+    workerData,
+    threadCount = getEnvironmentParallelism(),
+  } = options;
+
+  const iterator = options.items[Symbol.iterator]();
+  const results = new Array<TResult>();
+  let itemIndex = 0;
+
+  await Promise.all(
+    Array(threadCount)
+      .fill(0)
+      .map(async () => {
+        const thread = new Worker(`(${workerQueueThread})(${workerFactory})`, {
+          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;
+}
+
+function workerQueueThread(
+  workerFuncFactory: (data: unknown) => (item: unknown) => Promise<unknown>,
+) {
+  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<TResult, TData, TMessage> = {
+  /**
+   * 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<TResult>;
+  /** 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<TResult, TData, TMessage>(
+  options: WorkerThreadsOptions<TResult, TData, TMessage>,
+): Promise<TResult[]> {
+  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<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);
+            } 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<unknown>,
+) {
+  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 });
+    },
+  );
+}