Initial setup for concurrent workers (#46)

* Initial setup for concurrent workers * Add logs for timer * Prettier lint * Fix worker script issues and await in api * Lint * Fix worker in dev mode * Remove unneeded awaits * Lint and format * Ignore worker.js file in eslint * remove class from worker for readability --------- Co-authored-by: florianbgt <florian.bigot321@gmail.com>
nanoapi-io · Dec 11, 2024 · 4e01be7 · 4e01be7
1 parent 0cc7191
commit 4e01be7
Show file tree

Hide file tree

Showing 9 changed files with 163 additions and 97 deletions.
diff --git a/eslint.config.mjs b/eslint.config.mjs
@@ -1,8 +1,8 @@
 // @ts-check
 
 import eslint from "@eslint/js";
-import tseslint from "typescript-eslint";
 import eslintPluginPrettierRecommended from "eslint-plugin-prettier/recommended";
+import tseslint from "typescript-eslint";
 
 export default tseslint.config(
   eslint.configs.recommended,
@@ -16,6 +16,7 @@ export default tseslint.config(
       "examples/**",
       "packages/app/dist",
       "packages/cli/dist",
+      "packages/cli/src/**/worker.js",
     ],
   },
 );
diff --git a/package-lock.json b/package-lock.json
diff --git a/packages/cli/src/api/index.ts b/packages/cli/src/api/index.ts
@@ -1,10 +1,10 @@
 import { json, Router } from "express";
+import { z } from "zod";
+import { napiConfigSchema } from "../config";
+import { scanSchema, splitSchema, syncSchema } from "./helpers/validation";
 import { scan } from "./scan";
 import { split } from "./split";
 import { sync } from "./sync";
-import { scanSchema, splitSchema, syncSchema } from "./helpers/validation";
-import { napiConfigSchema } from "../config";
-import { z } from "zod";
 
 export function getApi(napiConfig: z.infer<typeof napiConfigSchema>) {
   const api = Router();
@@ -42,13 +42,13 @@ export function getApi(napiConfig: z.infer<typeof napiConfigSchema>) {
     res.status(200).json({ success: true });
   });
 
-  api.post("/api/split", (req, res) => {
+  api.post("/api/split", async (req, res) => {
     const result = splitSchema.safeParse(req.body);
     if (!result.success) {
       res.status(400).json(result.error.issues);
       return;
     }
-    const splitResult = split(result.data);
+    const splitResult = await split(result.data);
     res.status(200).json(splitResult);
   });
 

diff --git a/packages/cli/src/api/split.ts b/packages/cli/src/api/split.ts
@@ -1,15 +1,12 @@
-import fs from "fs";
 import path from "path";
+import { z } from "zod";
 import DependencyTreeManager from "../dependencyManager/dependencyManager";
 import { cleanupOutputDir, createOutputDir } from "../helper/file";
-import SplitRunner from "../splitRunner/splitRunner";
+import { runWithWorker, writeSplitsToDisk } from "../splitRunner/splitRunner";
 import { splitSchema } from "./helpers/validation";
-import { z } from "zod";
-import { Group } from "../dependencyManager/types";
 
-export function split(payload: z.infer<typeof splitSchema>) {
+export async function split(payload: z.infer<typeof splitSchema>) {
   console.time("split command");
-  const groupMap: Record<number, Group> = {};
 
   // Get the dependency tree
   const dependencyTreeManager = new DependencyTreeManager(
@@ -26,30 +23,19 @@ export function split(payload: z.infer<typeof splitSchema>) {
   const groups = dependencyTreeManager.getGroups();
 
   // Process each group for splitting
-  groups.forEach((group, index) => {
-    const splitRunner = new SplitRunner(dependencyTreeManager, group);
-    const files = splitRunner.run();
-
-    const targetDir = path.dirname(payload.entrypointPath);
-    const annotationDirectory = path.join(outputDir, index.toString());
-
-    files.forEach((file) => {
-      const relativeFileNamePath = path.relative(targetDir, file.path);
-      const destinationPath = path.join(
-        annotationDirectory,
-        relativeFileNamePath,
-      );
-      fs.mkdirSync(path.dirname(destinationPath), { recursive: true });
-      fs.writeFileSync(destinationPath, file.sourceCode, "utf8");
-    });
-  });
-
-  // Store the processed annotations in the output directory
-  groups.forEach((group, index) => {
-    groupMap[index] = group;
-  });
-  const annotationFilePath = path.join(outputDir, "annotations.json");
-  fs.writeFileSync(annotationFilePath, JSON.stringify(groupMap, null, 2));
+  const splits = groups.map((group, index) =>
+    runWithWorker(
+      index,
+      group,
+      dependencyTreeManager.entryPointPath,
+      dependencyTreeManager.getFiles(),
+    ),
+  );
+
+  // Wait for all splits to be processed
+  const processedSplits = await Promise.all(splits.map(async (split) => split));
+
+  writeSplitsToDisk(outputDir, payload.entrypointPath, processedSplits);
 
   console.timeEnd("split command");
   return { groups, success: true };

diff --git a/packages/cli/src/commands/split.ts b/packages/cli/src/commands/split.ts
@@ -1,16 +1,12 @@
-import path from "path";
-import fs from "fs";
 import DependencyTreeManager from "../dependencyManager/dependencyManager";
 import { cleanupOutputDir, createOutputDir } from "../helper/file";
-import SplitRunner from "../splitRunner/splitRunner";
-import { Group } from "../dependencyManager/types";
+import { runWithWorker, writeSplitsToDisk } from "../splitRunner/splitRunner";
 
-export default function splitCommandHandler(
+export default async function splitCommandHandler(
   entrypointPath: string, // Path to the entrypoint file
   outputDir: string, // Path to the output directory
 ) {
-  const groupMap: Record<number, Group> = {};
-
+  console.time("split command");
   const dependencyTreeManager = new DependencyTreeManager(entrypointPath);
 
   cleanupOutputDir(outputDir);
@@ -20,28 +16,19 @@ export default function splitCommandHandler(
   const groups = dependencyTreeManager.getGroups();
 
   // Process each group for splitting
-  groups.forEach((group, index) => {
-    const splitRunner = new SplitRunner(dependencyTreeManager, group);
-    const files = splitRunner.run();
+  const splits = groups.map((group, index) =>
+    runWithWorker(
+      index,
+      group,
+      dependencyTreeManager.entryPointPath,
+      dependencyTreeManager.getFiles(),
+    ),
+  );
 
-    const targetDir = path.dirname(entrypointPath);
-    const annotationDirectory = path.join(outputDir, index.toString());
+  // Wait for all splits to be processed
+  const processedSplits = await Promise.all(splits.map(async (split) => split));
 
-    files.forEach((file) => {
-      const relativeFileNamePath = path.relative(targetDir, file.path);
-      const destinationPath = path.join(
-        annotationDirectory,
-        relativeFileNamePath,
-      );
-      fs.mkdirSync(path.dirname(destinationPath), { recursive: true });
-      fs.writeFileSync(destinationPath, file.sourceCode, "utf8");
-    });
-  });
+  writeSplitsToDisk(outputDir, entrypointPath, processedSplits);
 
-  // Store the processed annotations in the output directory
-  groups.forEach((group, index) => {
-    groupMap[index] = group;
-  });
-  const annotationFilePath = path.join(outputDir, "annotations.json");
-  fs.writeFileSync(annotationFilePath, JSON.stringify(groupMap, null, 2));
+  console.timeEnd("split command");
 }
diff --git a/packages/cli/src/languagesPlugins/types.ts b/packages/cli/src/languagesPlugins/types.ts
@@ -1,5 +1,5 @@
-import { Group } from "../dependencyManager/types";
 import Parser from "tree-sitter";
+import { Group } from "../dependencyManager/types";
 
 export interface DepImportIdentifier {
   // Specific to each programing languages. Used by the language plugins.

diff --git a/packages/cli/src/splitRunner/splitRunner.ts b/packages/cli/src/splitRunner/splitRunner.ts
@@ -1,23 +1,30 @@
+import path from "path";
+import { Worker } from "worker_threads";
 import { Group } from "../dependencyManager/types";
-import { removeIndexesFromSourceCode } from "../helper/file";
-import DependencyTreeManager from "../dependencyManager/dependencyManager";
 import { File } from "./types";
-import Parser from "tree-sitter";
-import assert from "assert";
 import { getLanguagePlugin } from "../languagesPlugins";
 import { DepExport } from "../languagesPlugins/types";
+import { removeIndexesFromSourceCode } from "../helper/file";
+import assert from "assert";
+import Parser from "tree-sitter";
+import fs from "fs";
 
-class SplitRunner {
-  private dependencyTreeManager: DependencyTreeManager;
-  private entrypointPath: string;
+export class SplitRunner {
+  private index: number;
   private group: Group;
+  private entrypointPath: string;
   private files: File[];
 
-  constructor(dependencyTreeManager: DependencyTreeManager, group: Group) {
-    this.dependencyTreeManager = dependencyTreeManager;
-    this.entrypointPath = dependencyTreeManager.dependencyTree.path;
+  constructor(
+    index: number,
+    group: Group,
+    entrypointPath: string,
+    files: File[],
+  ) {
+    this.index = index;
+    this.entrypointPath = entrypointPath;
     this.group = group;
-    this.files = dependencyTreeManager.getFiles();
+    this.files = files;
   }
 
   #removeAnnotationFromOtherGroups() {
@@ -83,7 +90,7 @@ class SplitRunner {
       // We always want to keep the entrypoint file.
       // It will never be imported anywhere, so we add it now.
       const filesToKeep = new Set<string>();
-      filesToKeep.add(this.dependencyTreeManager.dependencyTree.path);
+      filesToKeep.add(this.entrypointPath);
 
       this.files.forEach((file) => {
         const languagePlugin = getLanguagePlugin(
@@ -191,41 +198,100 @@ class SplitRunner {
   }
 
   run() {
-    console.info("\n");
-    console.time("Splitting");
+    console.time(`Splitting-${this.index}`);
 
-    console.time("remove annotation from other groups");
+    console.time(`remove annotation from other groups-${this.index}`);
     this.#removeAnnotationFromOtherGroups();
-    console.timeEnd("remove annotation from other groups");
+    console.timeEnd(`remove annotation from other groups-${this.index}`);
 
-    console.time("Get export map");
+    console.time(`Get export map-${this.index}`);
     const exportMap = this.#getExportMap();
-    console.timeEnd("Get export map");
+    console.timeEnd(`Get export map-${this.index}`);
 
-    console.time("Remove invalid imports and usages");
+    console.time(`Remove invalid imports and usages-${this.index}`);
     this.#removeInvalidImportsAndUsages(exportMap);
-    console.timeEnd("Remove invalid imports and usages");
+    console.timeEnd(`Remove invalid imports and usages-${this.index}`);
 
-    console.time("Remove unused imports");
+    console.time(`Remove unused imports-${this.index}`);
     this.#removeUnusedImports();
-    console.timeEnd("Remove unused imports");
+    console.timeEnd(`Remove unused imports-${this.index}`);
 
-    console.time("Remove unused files");
+    console.time(`Remove unused files-${this.index}`);
     this.#removeUnusedFiles();
-    console.timeEnd("Remove unused files");
+    console.timeEnd(`Remove unused files-${this.index}`);
 
-    console.time("Remove unused exports");
+    console.time(`Remove unused exports-${this.index}`);
     this.#removeUnusedExports(exportMap);
-    console.timeEnd("Remove unused exports");
+    console.timeEnd(`Remove unused exports-${this.index}`);
 
-    console.time("Remove errors");
+    console.time(`Remove errors-${this.index}`);
     this.#removeErrors();
-    console.timeEnd("Remove errors");
+    console.timeEnd(`Remove errors-${this.index}`);
 
-    console.timeEnd("Splitting");
+    console.timeEnd(`Splitting-${this.index}`);
 
-    return this.files;
+    return { index: this.index, group: this.group, files: this.files };
   }
 }
 
-export default SplitRunner;
+export function runWithWorker(
+  index: number,
+  group: Group,
+  entryPointPath: string,
+  files: File[],
+) {
+  const worker = new Worker(path.resolve(__dirname, "worker"), {
+    workerData: {
+      index,
+      group,
+      entryPointPath,
+      files,
+    },
+  });
+
+  return new Promise<{ index: number; group: Group; files: File[] }>(
+    (resolve, reject) => {
+      worker.on(
+        "message",
+        (split: { index: number; group: Group; files: File[] }) => {
+          resolve(split);
+        },
+      );
+
+      worker.on("error", reject);
+      worker.on("exit", (code) => {
+        if (code !== 0) {
+          reject(new Error(`Worker stopped with exit code ${code}`));
+        }
+      });
+    },
+  );
+}
+
+export function writeSplitsToDisk(
+  outputDir: string,
+  entrypointPath: string,
+  splits: { index: number; group: Group; files: File[] }[],
+) {
+  const targetDir = path.dirname(entrypointPath);
+  const groupMap: Record<number, Group> = {};
+
+  splits.forEach((split) => {
+    const annotationDirectory = path.join(outputDir, split.index.toString());
+
+    split.files.forEach((file) => {
+      const relativeFileNamePath = path.relative(targetDir, file.path);
+      const destinationPath = path.join(
+        annotationDirectory,
+        relativeFileNamePath,
+      );
+      fs.mkdirSync(path.dirname(destinationPath), { recursive: true });
+      fs.writeFileSync(destinationPath, file.sourceCode, "utf8");
+    });
+
+    groupMap[split.index] = split.group;
+  });
+
+  const annotationFilePath = path.join(outputDir, "annotations.json");
+  fs.writeFileSync(annotationFilePath, JSON.stringify(groupMap, null, 2));
+}
diff --git a/packages/cli/src/splitRunner/worker.js b/packages/cli/src/splitRunner/worker.js
@@ -0,0 +1,3 @@
+const path = require("path");
+require("ts-node").register();
+require(path.resolve(__dirname, "worker.ts"));
diff --git a/packages/cli/src/splitRunner/worker.ts b/packages/cli/src/splitRunner/worker.ts
@@ -0,0 +1,23 @@
+import { parentPort, workerData } from "worker_threads";
+import { SplitRunner } from "./splitRunner";
+import { Group } from "../dependencyManager/types";
+import { File } from "./types";
+
+const {
+  index,
+  group,
+  entryPointPath,
+  files,
+}: {
+  index: number;
+  group: Group;
+  entryPointPath: string;
+  files: File[];
+} = workerData;
+
+(() => {
+  const splitRunner = new SplitRunner(index, group, entryPointPath, files);
+  const updatedFiled = splitRunner.run();
+  // Send updated files back to the parent
+  parentPort?.postMessage(updatedFiled);
+})();