feat: introduce parallelization worker for chunked generation processing

Author: ovflowd

bin/commands/generate.mjs

@@ -16,6 +16,10 @@ import { loadAndParse } from '../utils.mjs';
 
 const availableGenerators = Object.keys(publicGenerators);
 
+// Half of available logical CPUs guarantees in general all physical CPUs are being used
+// which in most scenarios is the best way to maximize performance
+const optimalThreads = Math.floor(cpus().length / 2) - 1;
+
 /**
  * @typedef {Object} Options
  * @property {Array<string>|string} input - Specifies the glob/path for input files.
@@ -26,6 +30,7 @@ const availableGenerators = Object.keys(publicGenerators);
  * @property {string} typeMap - Specifies the path to the Node.js Type Map.
  * @property {string} [gitRef] - Git ref/commit URL.
  * @property {number} [threads] - Number of threads to allow.
+ * @property {number} [chunkSize] - Number of items to process per worker thread.
  */
 
 /**
@@ -61,10 +66,20 @@ export default {
     },
     threads: {
       flags: ['-p', '--threads <number>'],
+      desc: 'Number of worker threads to use',
       prompt: {
         type: 'text',
         message: 'How many threads to allow',
-        initialValue: String(Math.max(cpus().length, 1)),
+        initialValue: String(Math.max(optimalThreads, 1)),
+      },
+    },
+    chunkSize: {
+      flags: ['--chunk-size <number>'],
+      desc: 'Number of items to process per worker thread (default: auto)',
+      prompt: {
+        type: 'text',
+        message: 'Items per worker thread',
+        initialValue: '20',
       },
     },
     version: {
@@ -149,6 +164,7 @@ export default {
       releases,
       gitRef: opts.gitRef,
       threads: parseInt(opts.threads, 10),
+      chunkSize: parseInt(opts.chunkSize, 10),
       index,
       typeMap,
     });

src/generators.mjs

@@ -2,6 +2,7 @@
 
 import { allGenerators } from './generators/index.mjs';
 import WorkerPool from './threading/index.mjs';
+import createParallelWorker from './threading/parallel.mjs';
 
 /**
  * This method creates a system that allows you to register generators
@@ -31,14 +32,26 @@ const createGenerator = input => {
    */
   const cachedGenerators = { ast: Promise.resolve(input) };
 
-  const threadPool = new WorkerPool();
-
   /**
    * Runs the Generator engine with the provided top-level input and the given generator options
    *
    * @param {GeneratorOptions} options The options for the generator runtime
    */
-  const runGenerators = async ({ generators, threads, ...extra }) => {
+  const runGenerators = async ({
+    generators,
+    threads,
+    chunkSize,
+    ...extra
+  }) => {
+    // WorkerPool for running full generators in worker threads
+    const generatorPool = new WorkerPool('./generator-worker.mjs', threads);
+
+    // WorkerPool for chunk-level parallelization within generators
+    const chunkPool = new WorkerPool('./chunk-worker.mjs', threads);
+
+    // Options including threading config
+    const threadingOptions = { threads, chunkSize };
+
     // Note that this method is blocking, and will only execute one generator per-time
     // but it ensures all dependencies are resolved, and that multiple bottom-level generators
     // can reuse the already parsed content from the top-level/dependency generators
@@ -50,20 +63,32 @@ const createGenerator = input => {
       if (dependsOn && dependsOn in cachedGenerators === false) {
         await runGenerators({
           ...extra,
-          threads,
+          ...threadingOptions,
           generators: [dependsOn],
         });
       }
 
       // Ensures that the dependency output gets resolved before we run the current
       // generator with its dependency output as the input
-      const dependencyOutput = await cachedGenerators[dependsOn];
+      const input = await cachedGenerators[dependsOn];
+
+      // Create a ParallelWorker for this generator to use for item-level parallelization
+      const worker = createParallelWorker(generatorName, chunkPool, {
+        ...extra,
+        ...threadingOptions,
+      });
+
+      // Generator options with worker instance
+      const generatorOptions = { ...extra, ...threadingOptions, worker };
+
+      // Worker options for the worker thread
+      const workerOptions = { ...extra, ...threadingOptions };
 
       // Adds the current generator execution Promise to the cache
       cachedGenerators[generatorName] =
         threads < 2
-          ? generate(dependencyOutput, extra) // Run in main thread
-          : threadPool.run(generatorName, dependencyOutput, threads, extra); // Offload to worker thread
+          ? generate(input, generatorOptions) // Run in main thread
+          : generatorPool.run({ generatorName, input, options: workerOptions }); // Offload to worker thread
     }
 
     // Returns the value of the last generator of the current pipeline

src/generators/ast-js/index.mjs

@@ -23,21 +23,42 @@ export default {
   dependsOn: 'metadata',
 
   /**
-   * @param {Input} _
+   * Process a chunk of JavaScript files in a worker thread.
+   * Called by chunk-worker.mjs for parallel processing.
+   *
+   * @param {unknown} _ - Unused (we use options.input instead)
+   * @param {number[]} itemIndices - Indices of source files to process
    * @param {Partial<GeneratorOptions>} options
    */
-  async generate(_, options) {
+  async processChunk(_, itemIndices, { input }) {
     const { loadFiles } = createJsLoader();
 
-    // Load all of the Javascript sources into memory
-    const sourceFiles = loadFiles(options.input ?? []);
+    const sourceFiles = loadFiles(input ?? []);
+
+    const { parseJsSource } = createJsParser();
+
+    const results = [];
 
-    const { parseJsSources } = createJsParser();
+    for (const idx of itemIndices) {
+      results.push(await parseJsSource(sourceFiles[idx]));
+    }
 
-    // Parse the Javascript sources into ASTs
-    const parsedJsFiles = await parseJsSources(sourceFiles);
+    return results;
+  },
+
+  /**
+   * @param {Input} _
+   * @param {Partial<GeneratorOptions>} options
+   */
+  async generate(_, { input, worker }) {
+    const { loadFiles } = createJsLoader();
+
+    // Load all of the Javascript sources into memory
+    const sourceFiles = loadFiles(input ?? []);
 
-    // Return the ASTs so they can be used in another generator
-    return parsedJsFiles;
+    // Parse the Javascript sources into ASTs in parallel using worker threads
+    // Note: We pass sourceFiles as items but _ (empty) as fullInput since
+    // processChunk reloads files from options.input
+    return worker.map(sourceFiles, _, { input });
   },
 };

src/generators/jsx-ast/index.mjs

@@ -1,7 +1,10 @@
-import { OVERRIDDEN_POSITIONS } from './constants.mjs';
 import { buildSideBarProps } from './utils/buildBarProps.mjs';
 import buildContent from './utils/buildContent.mjs';
-import { groupNodesByModule } from '../../utils/generators.mjs';
+import {
+  buildDocPages,
+  getSortedHeadNodes,
+  groupNodesByModule,
+} from '../../utils/generators.mjs';
 import { getRemarkRecma } from '../../utils/remark.mjs';
 
 /**
@@ -10,78 +13,53 @@ import { getRemarkRecma } from '../../utils/remark.mjs';
  * @typedef {Array<ApiDocMetadataEntry>} Input
  * @type {GeneratorMetadata<Input, string>}
  */
-
-/**
- * Sorts entries by OVERRIDDEN_POSITIONS and then heading name.
- * @param {Array<ApiDocMetadataEntry>} entries
- */
-const getSortedHeadNodes = entries => {
-  return entries
-    .filter(node => node.heading.depth === 1)
-    .sort((a, b) => {
-      const ai = OVERRIDDEN_POSITIONS.indexOf(a.api);
-      const bi = OVERRIDDEN_POSITIONS.indexOf(b.api);
-
-      if (ai !== -1 && bi !== -1) {
-        return ai - bi;
-      }
-
-      if (ai !== -1) {
-        return -1;
-      }
-
-      if (bi !== -1) {
-        return 1;
-      }
-
-      return a.heading.data.name.localeCompare(b.heading.data.name);
-    });
-};
-
 export default {
   name: 'jsx-ast',
   version: '1.0.0',
   description: 'Generates JSX AST from the input MDAST',
   dependsOn: 'metadata',
 
   /**
-   * Generates a JSX AST
-   *
-   * @param {Input} entries
+   * Process a chunk of items in a worker thread.
+   * @param {Input} fullInput
+   * @param {number[]} itemIndices
    * @param {Partial<GeneratorOptions>} options
-   * @returns {Promise<Array<string>>} Array of generated content
    */
-  async generate(entries, { index, releases, version }) {
-    const remarkRecma = getRemarkRecma();
-    const groupedModules = groupNodesByModule(entries);
-    const headNodes = getSortedHeadNodes(entries);
-
-    // Generate table of contents
-    const docPages = index
-      ? index.map(({ section, api }) => [section, `${api}.html`])
-      : headNodes.map(node => [node.heading.data.name, `${node.api}.html`]);
+  async processChunk(fullInput, itemIndices, { index, releases, version }) {
+    const processor = getRemarkRecma();
+    const groupedModules = groupNodesByModule(fullInput);
+    const headNodes = getSortedHeadNodes(fullInput);
+    const docPages = buildDocPages(headNodes, index);
 
-    // Process each head node and build content
-    const results = [];
+    return Promise.all(
+      itemIndices.map(async idx => {
+        const entry = headNodes[idx];
 
-    for (const entry of headNodes) {
-      const sideBarProps = buildSideBarProps(
-        entry,
-        releases,
-        version,
-        docPages
-      );
+        const sideBarProps = buildSideBarProps(
+          entry,
+          releases,
+          version,
+          docPages
+        );
 
-      results.push(
-        await buildContent(
+        return buildContent(
           groupedModules.get(entry.api),
           entry,
           sideBarProps,
-          remarkRecma
-        )
-      );
-    }
+          processor
+        );
+      })
+    );
+  },
+
+  /**
+   * Generates a JSX AST
+   * @param {Input} entries
+   * @param {Partial<GeneratorOptions>} options
+   */
+  async generate(entries, { index, releases, version, worker }) {
+    const headNodes = getSortedHeadNodes(entries);
 
-    return results;
+    return worker.map(headNodes, entries, { index, releases, version });
   },
 };

src/generators/legacy-html-all/index.mjs

@@ -5,7 +5,7 @@ import { join, resolve } from 'node:path';
 
 import HTMLMinifier from '@minify-html/node';
 
-import { getRemarkRehypeWithShiki } from '../../utils/remark.mjs';
+import { getRemarkRehype } from '../../utils/remark.mjs';
 import dropdowns from '../legacy-html/utils/buildDropdowns.mjs';
 import tableOfContents from '../legacy-html/utils/tableOfContents.mjs';
 
@@ -49,7 +49,7 @@ export default {
     const inputWithoutIndex = input.filter(entry => entry.api !== 'index');
 
     // Gets a Remark Processor that parses Markdown to minified HTML
-    const remarkWithRehype = getRemarkRehypeWithShiki();
+    const remarkWithRehype = getRemarkRehype();
 
     // Current directory path relative to the `index.mjs` file
     // from the `legacy-html` generator, as all the assets are there