Worker migration to worker_threads

Author: Daniel Del Core

packages/core/lib/Worker.js

@@ -1,43 +1,28 @@
-/* eslint-disable @typescript-eslint/no-var-requires */
-'use strict';
-
-const path = require('path');
-const { EventEmitter } = require('events');
-const async = require('neo-async');
-const fs = require('graceful-fs');
-const writeFileAtomic = require('write-file-atomic');
-const { DEFAULT_EXTENSIONS } = require('@babel/core');
-
-const getParser = require('jscodeshift/src/getParser');
-const jscodeshift = require('jscodeshift/src/core');
-
-let presetEnv;
-try {
-  presetEnv = require('@babel/preset-env');
-} catch (_) {}
-
-let emitter;
-let finish;
-let notify;
+import path from 'path';
+import async from 'neo-async';
+import fs from 'graceful-fs';
+import writeFileAtomic from 'write-file-atomic';
+import { fileURLToPath } from 'url';
+import { register } from 'tsx/esm/api';
+
+import getParser from 'jscodeshift/src/getParser.js';
+import jscodeshift from 'jscodeshift/src/core.js';
+import { workerData, isMainThread, parentPort } from 'worker_threads';
+
+/**
+ * Register the TSX plugin to allow import TS(X) files.
+ */
+register();
+
 let transform;
 let parserFromTransform;
 
-if (module.parent) {
-  emitter = new EventEmitter();
-  // @ts-expect-error
-  emitter.send = data => run(data);
-  finish = () => emitter.emit('disconnect');
-  notify = data => emitter.emit('message', data);
-
-  module.exports = args => {
-    setup(args[0], args[1]);
-    return emitter;
-  };
-} else {
-  finish = () => setImmediate(() => process.disconnect());
-  notify = data => process.send(data);
-  process.on('message', data => run(data));
-  setup(process.argv[2], process.argv[3]);
+// Get the __filename and __dirname equivalents for ESM
+const __filename = fileURLToPath(import.meta.url);
+
+if (!isMainThread) {
+  await setup(workerData.entrypointPath);
+  parentPort.on('message', data => run(data));
 }
 
 function prepareJscodeshift(options) {
@@ -102,12 +87,12 @@ async function setup(entryPath) {
 
 function updateStatus(status, file, msg) {
   msg = msg ? file + ' ' + msg : file;
-  notify({ action: 'status', status, msg });
+  parentPort.postMessage({ action: 'status', status, msg });
 }
 
 function stats(name, quantity) {
   quantity = typeof quantity !== 'number' ? 1 : quantity;
-  notify({ action: 'update', name: name, quantity: quantity });
+  parentPort.postMessage({ action: 'update', name: name, quantity: quantity });
 }
 
 function trimStackTrace(trace) {
@@ -130,9 +115,10 @@ function run(data) {
   const options = data.options || {};
 
   if (!files.length) {
-    finish();
+    parentPort.close();
     return;
   }
+
   async.each(
     files,
     function (file, callback) {
@@ -153,7 +139,8 @@ function run(data) {
               jscodeshift: jscodeshift,
               // eslint-disable-next-line @typescript-eslint/no-empty-function
               stats: options.dry ? stats : () => {},
-              report: msg => notify({ action: 'report', file, msg }),
+              report: msg =>
+                parentPort.postMessage({ action: 'report', file, msg }),
             },
             options,
           );
@@ -196,7 +183,7 @@ function run(data) {
       if (err) {
         updateStatus('error', '', 'This should never be shown!');
       }
-      notify({ action: 'free' });
+      parentPort.postMessage({ action: 'free' });
     },
   );
 }

packages/core/src/runner.ts

@@ -1,12 +1,15 @@
-import child_process from 'child_process';
+import { Worker } from 'worker_threads';
 import chalk from 'chalk';
 import fs from 'graceful-fs';
 import path from 'path';
 import os from 'os';
-// @ts-expect-error
-import ignores from 'jscodeshift/src/ignoreFiles';
 
-import { Message, Flags, Statuses } from './types';
+import { Message, Flags, Statuses } from './types.js';
+
+// @ts-expect-error
+const ignores = await import('jscodeshift/src/ignoreFiles.js');
+// @ts-expect-error
+const __dirname = import.meta.dirname;
 
 type FileCounters = Record<Statuses, number>;
 type Stats = Record<string, number>;
@@ -230,25 +233,23 @@ export function run(
         }
       }
 
-      const args = [entrypointPath, options.babel ? 'babel' : 'no-babel'];
-
-      const workers = [];
+      const workers: Worker[] = [];
 
       for (let i = 0; i < processes; i++) {
         workers.push(
-          options.runInBand
-            ? // eslint-disable-next-line @typescript-eslint/no-var-requires
-              require('../lib/Worker')(args)
-            : child_process.fork(
-                path.join(__dirname, '..', 'lib', 'Worker.js'),
-                args,
-              ),
+          new Worker(path.join(__dirname, '..', 'lib', 'Worker.js'), {
+            workerData: {
+              entrypointPath,
+              babel: options.babel ? 'babel' : 'no-babel',
+            },
+          }),
         );
       }
 
-      return workers.map(child => {
-        child.send({ files: next(), options });
-        child.on('message', (message: Message) => {
+      return workers.map(worker => {
+        worker.postMessage({ files: next(), options });
+        worker.on('error', (message: Message) => console.error(message));
+        worker.on('message', (message: Message) => {
           switch (message.action) {
             case 'status':
               fileCounters[message.status] += 1;
@@ -261,7 +262,7 @@ export function run(
               statsCounter[message.name] += message.quantity;
               break;
             case 'free':
-              child.send({ files: next(), options });
+              worker.postMessage({ files: next(), options });
               break;
             case 'report':
               bufferedWrite(
@@ -274,7 +275,7 @@ export function run(
               break;
           }
         });
-        return new Promise(resolve => child.on('disconnect', resolve));
+        return new Promise(resolve => worker.on('close', resolve));
       });
     })
     .then(pendingWorkers =>