Use worker_threads to index files in parallel.

ShukantPal · ShukantPal · commit c199f64558c5 · 2021-05-25T23:22:45.000-04:00
diff --git a/.idea/misc.xml b/.idea/misc.xml
diff --git a/packages/webdoc-cli/src/index.js b/packages/webdoc-cli/src/index.js
@@ -107,7 +107,7 @@ async function main(argv: yargs.Argv) {
   }
 
   try {
-    parse(sourceFiles, documentTree);
+    await parse(sourceFiles, documentTree);
   } catch (e) {
     // Make sure we get that API structure out so the user can debug the problem!
     if (config.opts && config.opts.export) {
diff --git a/packages/webdoc-parser/src/indexer/IndexerWorkerPool.js b/packages/webdoc-parser/src/indexer/IndexerWorkerPool.js
@@ -0,0 +1,100 @@
+// @flow
+
+import EventEmitter from "events";
+import type {LanguageConfig} from "../types/LanguageIntegration";
+import {type SourceFile} from "@webdoc/types";
+import {type Symbol} from "../types/Symbol";
+import {Worker} from "worker_threads";
+import os from "os";
+import path from "path";
+
+export type JobData = {
+  jobId: number,
+  jobLanguageIntegrationModule: string,
+  jobFile: SourceFile,
+  jobConfig: LanguageConfig,
+};
+
+export type JobResult = {
+  jobId: number,
+  jobResult?: Symbol,
+  jobError?: any,
+};
+
+class IndexerWorker extends EventEmitter {
+  onError: (ev: "error", ...args: any[]) => void;
+  onMessage: (...args: any[]) => void;
+  worker: Worker;
+
+  constructor(worker: Worker) {
+    super();
+
+    this.onError = this.emit.bind(this, "error");
+    this.onMessage = this.onMessageImpl.bind(this);
+
+    this.worker = worker;
+    this.worker.on("error", this.onError);
+    this.worker.on("message", this.onMessage);
+  }
+
+  send(
+    jobId: number,
+    jobLanguageIntegrationModule: string,
+    jobFile: SourceFile,
+    jobConfig: LanguageConfig,
+  ): number {
+    this.worker.postMessage(({
+      jobId,
+      jobLanguageIntegrationModule,
+      jobFile,
+      jobConfig,
+    }: JobData));
+
+    return jobId;
+  }
+
+  onMessageImpl({jobId, jobResult, jobError}: JobResult) {
+    this.emit(`response-${jobId}`, jobResult, jobError);
+  }
+
+  destroy() {
+    this.worker.terminate();
+  }
+}
+
+export class IndexerWorkerPool {
+  workers: IndexerWorker[];
+  ptr: number = 0;
+
+  constructor(limit?: number) {
+    const workerPoolSize = Math.min(os.cpus().length, limit || Infinity);
+    const workers = new Array<Worker>(workerPoolSize);
+    const workerPath = path.resolve(__dirname, "./worker.js");
+
+    for (let i = 0; i < workerPoolSize; i++) {
+      workers[i] = new IndexerWorker(new Worker(workerPath));
+    }
+
+    this.workers = workers;
+  }
+
+  index(langIntegrationModule: string, file: SourceFile, config: LanguageConfig): Promise<Symbol> {
+    return new Promise<Symbol>((resolve, reject) => {
+      const jobId = this.ptr++;
+      const worker = this.workers[jobId % this.workers.length];
+
+      worker.send(jobId, langIntegrationModule, file, config);
+      worker.on(`response-${jobId}`, (jobResult: Symbol, jobError: any): void => {
+        if (jobResult) resolve(jobResult);
+        else reject(jobError);
+      });
+    });
+  }
+
+  destroy() {
+    for (const worker of this.workers) {
+      worker.destroy();
+    }
+    this.workers.length = 0;
+  }
+}
diff --git a/packages/webdoc-parser/src/indexer/index.js b/packages/webdoc-parser/src/indexer/index.js
@@ -0,0 +1,81 @@
+// @flow
+
+import type {LanguageConfig, LanguageIntegration} from "../types/LanguageIntegration";
+import {IndexerWorkerPool} from "./IndexerWorkerPool";
+import {type SourceFile} from "@webdoc/types";
+import {type Symbol} from "../types/Symbol";
+import {parserLogger} from "../Logger";
+import path from "path";
+
+declare var globalThis: any;
+
+// File-extension -> LanguageIntegration mapping
+const languages: { [id: string]: LanguageIntegration } = {};
+
+// Register a language-integration that will be used to generate a symbol-tree for files with its
+// file-extensions.
+export function register(lang: LanguageIntegration): void {
+  for (const ext of lang.extensions) {
+    if (languages[ext]) {
+      parserLogger.warn("LanguageIntegration",
+        `.${ext} file extension has already been registered`);
+    }
+
+    languages[ext] = lang;
+  }
+}
+
+export async function run(files: SourceFile[], config: LanguageConfig): Promise<Symbol[]> {
+  const symbolTrees: Array<Symbol> = new Array(files.length);
+  const symbolIndexingOperations: Array<Promise<void>> = new Array(files.length);
+  const workerPool = new IndexerWorkerPool();
+
+  for (let i = 0; i < files.length; i++) {
+    const fileName = files[i].path;
+    const extension = fileName.substring(fileName.lastIndexOf(".") + 1, fileName.length);
+
+    if (!(extension in languages)) {
+      throw new Error(`.${extension} language is not registered with the Indexer!`);
+    }
+
+    const file = {
+      ...files[i],
+      path: path.resolve(globalThis.process.cwd(), files[i].path),
+    };
+    const lang = languages[extension].module;
+
+    symbolIndexingOperations[i] = workerPool.index(lang, file, config).then(
+      (symbolTree: Symbol): void => {
+        symbolTrees[i] = symbolTree;
+      },
+    );
+  }
+
+  await Promise.all(symbolIndexingOperations);
+
+  workerPool.destroy();
+
+  return symbolTrees;
+}
+
+export function process(
+  file: string,
+  source: SourceFile,
+  config: LanguageConfig,
+): Symbol {
+  const fileName = source.path;
+  const lang = languages[fileName.substring(fileName.lastIndexOf(".") + 1, fileName.length)];
+
+  if (!lang) {
+    throw new Error(`.${lang} file language is not registered`);
+  }
+
+  return lang.parse(file, source, config);
+}
+
+export function lang(module: string): LanguageIntegration {
+  module = module.replace("@webdoc/parser", "../../");
+
+  // $FlowFixMe
+  return ((require(module): any).default: LanguageIntegration);
+}
diff --git a/packages/webdoc-parser/src/indexer/worker.js b/packages/webdoc-parser/src/indexer/worker.js
@@ -0,0 +1,39 @@
+// @flow
+
+import * as Indexer from "./";
+import type {JobData, JobResult} from "./IndexerWorkerPool";
+import {isMainThread, parentPort} from "worker_threads";
+import fs from "fs";
+
+function onMessage(data: JobData) {
+  function onFileRead(err: ?Error, contents: string) {
+    if (err) {
+      return parentPort.postMessage(({
+        jobId: data.jobId,
+        jobError: err,
+      }: JobResult));
+    }
+
+    const lang = Indexer.lang(data.jobLanguageIntegrationModule);
+    const symbolTree = lang.parse(
+      contents,
+      data.jobFile,
+      data.jobConfig,
+    );
+
+    parentPort.postMessage(({
+      jobId: data.jobId,
+      jobResult: symbolTree,
+    }: JobResult));
+  }
+
+  if (!data.jobFile.content) {
+    fs.readFile(data.jobFile.path, "utf8", onFileRead);
+  } else {
+    onFileRead(null, data.jobFile.content);
+  }
+}
+
+if (!isMainThread) {
+  parentPort.on("message", onMessage);
+}
diff --git a/packages/webdoc-parser/src/parse.js b/packages/webdoc-parser/src/parse.js
@@ -1,38 +1,20 @@
 // @flow
 
-import type {LanguageConfig, LanguageIntegration} from "./types/LanguageIntegration";
+import * as Indexer from "./indexer";
 import type {RootDoc, SourceFile} from "@webdoc/types";
 import {createPackageDoc, createRootDoc} from "@webdoc/model";
 import {langJS, langTS} from "./symbols-babel";
+import type {LanguageConfig} from "./types/LanguageIntegration";
 import type {Symbol} from "./types/Symbol";
 import assemble from "./assembler";
-import fs from "fs";
 import mod from "./transformer/document-tree-modifiers";
-import {parserLogger} from "./Logger";
-import path from "path";
 import transform from "./transformer";
 
 declare var Webdoc: any;
 
-// File-extension -> LanguageIntegration mapping
-const languages: { [id: string]: LanguageIntegration } = {};
-
-// Register a language-integration that will be used to generate a symbol-tree for files with its
-// file-extensions.
-export function registerLanguage(lang: LanguageIntegration): void {
-  for (const ext of lang.extensions) {
-    if (languages[ext]) {
-      parserLogger.warn("LanguageIntegration",
-        `.${ext} file extension has already been registered`);
-    }
-
-    languages[ext] = lang;
-  }
-}
-
 // Register built-in languages
-registerLanguage(langJS);
-registerLanguage(langTS);
+Indexer.register(langJS);
+Indexer.register(langTS);
 
 // Default language-config for parsing documentation
 const DEFAULT_LANG_CONFIG: LanguageConfig = {
@@ -79,14 +61,7 @@ export function buildSymbolTree(
     };
   }
 
-  const fileName = source.path;
-  const lang = languages[fileName.substring(fileName.lastIndexOf(".") + 1, fileName.length)];
-
-  if (!lang) {
-    throw new Error(`.${lang} file language is not registered`);
-  }
-
-  return lang.parse(file, source, config);
+  return Indexer.process(file, source, config);
 }
 
 // TODO: Asynchronous API for parsing
@@ -107,7 +82,10 @@ export function buildSymbolTree(
  * @param {RootDoc} root
  * @return {RootDoc}
  */
-export function parse(target: string | SourceFile[], root?: RootDoc = createRootDoc()): RootDoc {
+export async function parse(
+  target: string | SourceFile[],
+  root?: RootDoc = createRootDoc(),
+): Promise<RootDoc> {
   if (typeof target === "string") {
     target = [{
       content: target,
@@ -116,23 +94,12 @@ export function parse(target: string | SourceFile[], root?: RootDoc = createRoot
     }];
   }
 
-  const partialDoctrees = new Array(Array.isArray(target) ? target.length : target.size);
-
-  // Build a symbol-tree for all the files
-  for (let i = 0; i < target.length; i++) {
-    const {content, path: source} = target[i];
-
-    partialDoctrees[i] = buildSymbolTree(
-      content || fs.readFileSync(path.join(process.cwd(), source), "utf8"),
-      target[i],
-    );
-  }
-
-  const rsym = assemble(partialDoctrees);
+  const perFileSymbolTrees = await Indexer.run(target, Webdoc.DEFAULT_LANG_CONFIG);
+  const fullSymbolTree = assemble(perFileSymbolTrees);
 
   root.children = root.members;
 
-  transform(rsym, root);
+  transform(fullSymbolTree, root);
   mod(root);
 
   return root;
diff --git a/packages/webdoc-parser/src/symbols-babel/index.js b/packages/webdoc-parser/src/symbols-babel/index.js
@@ -3,7 +3,7 @@
 // This file provides langJS and langTS integration, both of which use Babel to generate
 // symbols trees. These are registered in parse.js!
 
-import type {LanguageConfig} from "../types/LanguageIntegration";
+import type {LanguageConfig, LanguageIntegration} from "../types/LanguageIntegration";
 import type {SourceFile} from "@webdoc/types";
 import type {Symbol} from "../types/Symbol";
 import buildSymbolTree from "./build-symbol-tree";
@@ -46,8 +46,9 @@ const tsPreset = [
   "typescript",
 ];
 
-export const langJS = {
+export const langJS: LanguageIntegration = {
   extensions: ["js", "jsx", "jsdoc"],
+  module: "@webdoc/parser/lib/symbols-babel/langJS",
   parse(file: string, source: SourceFile, config: LanguageConfig): Symbol {
     mode.current = "flow";
 
@@ -56,8 +57,9 @@ export const langJS = {
   },
 };
 
-export const langTS = {
+export const langTS: LanguageIntegration = {
   extensions: ["ts", "tsx"],
+  module: "@webdoc/parser/lib/symbols-babel/langTS",
   parse(file: string, source: SourceFile, config: LanguageConfig): Symbol {
     mode.current = "typescript";
 
diff --git a/packages/webdoc-parser/src/symbols-babel/langJS.js b/packages/webdoc-parser/src/symbols-babel/langJS.js
@@ -0,0 +1 @@
+export {langJS as default} from "./";
diff --git a/packages/webdoc-parser/src/symbols-babel/langTS.js b/packages/webdoc-parser/src/symbols-babel/langTS.js
@@ -0,0 +1 @@
+export {langTS as default} from "./";
diff --git a/packages/webdoc-parser/src/types/LanguageIntegration.js b/packages/webdoc-parser/src/types/LanguageIntegration.js
@@ -6,6 +6,7 @@ import type {Symbol} from "./Symbol";
 // Language-integration definition - used for parsing code into a symbol-tree
 export type LanguageIntegration = {
   extensions: string[],
+  module: string,
   parse(file: string, source: SourceFile, config: LanguageConfig): Symbol
 };
 
diff --git a/packages/webdoc-parser/test/parse-es5-class.js b/packages/webdoc-parser/test/parse-es5-class.js
@@ -4,8 +4,8 @@ const {parse} = require("../lib/parse");
 const expect = require("chai").expect;
 
 describe("@webdoc/parser.parse (ES5 classes)", function() {
-  it("should correctly infer ES5 classes", function() {
-    const docTree = parse(`      
+  it("should correctly infer ES5 classes", async function() {
+    const docTree = await parse(`      
       /**
        * ES5 constructor
        * @classdesc
diff --git a/packages/webdoc-parser/test/parse-tag-variations.js b/packages/webdoc-parser/test/parse-tag-variations.js
diff --git a/packages/webdoc-parser/test/parse-ts.js b/packages/webdoc-parser/test/parse-ts.js
diff --git a/packages/webdoc-parser/test/parse.js b/packages/webdoc-parser/test/parse.js

Original file line number	Diff line number	Diff line change
`@@ -107,7 +107,7 @@ async function main(argv: yargs.Argv) {`
`107`	`107`	`}`
`108`	`108`
`109`	`109`	`try {`
`110`		`- parse(sourceFiles, documentTree);`
	`110`	`+ await parse(sourceFiles, documentTree);`
`111`	`111`	`} catch (e) {`
`112`	`112`	`// Make sure we get that API structure out so the user can debug the problem!`
`113`	`113`	`if (config.opts && config.opts.export) {`