refactor: small adjustments over WorkerInfo (#860)

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
(cherry picked from commit 48bf5fa)

Author: 2 people

src/index.ts

@@ -231,31 +231,40 @@ class ThreadPool {
   }
 
   _addNewWorker () : void {
-    if (this.closingUp) return;
-
-    const pool = this;
-    const worker = new Worker(resolve(__dirname, 'worker.js'), {
-      env: this.options.env,
-      argv: this.options.argv,
-      execArgv: this.options.execArgv,
-      resourceLimits: this.options.resourceLimits,
-      workerData: this.options.workerData,
-      trackUnmanagedFds: this.options.trackUnmanagedFds
-    });
+    if (this.closingUp === true) return;
 
     const { port1, port2 } = new MessageChannel();
     const workerInfo = new WorkerInfo({
-      worker, 
-      onMessage,
+      worker: {
+        filename: resolve(__dirname, 'worker.js'),
+        env: this.options.env,
+        argv: this.options.argv,
+        execArgv: this.options.execArgv,
+        resourceLimits: this.options.resourceLimits,
+        workerData: this.options.workerData,
+        trackUnmanagedFds: this.options.trackUnmanagedFds
+      }, 
       port: port1, 
       enableHistogram: this.options.workerHistogram
-    });
-
+    }, onMessage.bind(this));
+    const message : StartupMessage = {
+      filename: this.options.filename,
+      name: this.options.name,
+      port: port2,
+      sharedBuffer: workerInfo.sharedBuffer,
+      atomics: this.options.atomics!,
+      niceIncrement: this.options.niceIncrement
+    };
+    
     workerInfo.onDestroy(() => {
       this.publicInterface.emit('workerDestroy', workerInfo.interface);
     });
+    workerInfo.onWorkerMessage(onWorkerMessage.bind(this));
+    workerInfo.onWorkerError(onWorkerError.bind(this));
+    workerInfo.onWorkerExit(onWorkerExit.bind(this));
+    workerInfo.onPortClose(() => { workerInfo.workerRef(); });
 
-    if (this.startingUp) {
+    if (this.startingUp === true) {
       // There is no point in waiting for the initial set of Workers to indicate
       // that they are ready, we just mark them as such from the start.
       workerInfo.markAsReady();
@@ -271,74 +280,55 @@ class ThreadPool {
         this._onWorkerReady(workerInfo);
       });
     }
+    
+    workerInfo.init(message, [port2]).workerUnref();
+    this.workers.add(workerInfo);
 
-    const message : StartupMessage = {
-      filename: this.options.filename,
-      name: this.options.name,
-      port: port2,
-      sharedBuffer: workerInfo.sharedBuffer,
-      atomics: this.options.atomics!,
-      niceIncrement: this.options.niceIncrement
-    };
-    worker.postMessage(message, [port2]);
-
-    function onMessage (message : ResponseMessage) {
+    function onMessage (this: ThreadPool, message : ResponseMessage) {
       const { taskId, result } = message;
       // In case of success: Call the callback that was passed to `runTask`,
       // remove the `TaskInfo` associated with the Worker, which marks it as
       // free again.
       const taskInfo = workerInfo.popTask(taskId);
-      pool.workers.taskDone(workerInfo);
+      this.workers.taskDone(workerInfo);
 
-      /* istanbul ignore if */
-      if (taskInfo == null) {
+      
+      if (taskInfo == null) { /* c8 ignore next */
         const err = new Error(
           `Unexpected message from Worker: ${inspect(message)}`);
-        pool.publicInterface.emit('error', err);
+        this.publicInterface.emit('error', err);
       } else {
         taskInfo.done(message.error, result);
       }
 
-      pool._processPendingMessages();
+      this._processPendingMessages();
     }
 
-    function onReady () {
+    function onWorkerReady () {
       workerInfo.currentUsage() === 0 && workerInfo.unref();
       workerInfo.isReady() === false && workerInfo.markAsReady();
     }
 
-    function onEventMessage (message: any) {
-      pool.publicInterface.emit('message', message);
+    function onEventMessage (this: ThreadPool, message: any) {
+      this.publicInterface.emit('message', message);
     }
 
-    worker.on('message', (message : any) => {
-      message instanceof Object && READY in message ? onReady() : onEventMessage(message);
-    });
+    function onWorkerMessage (this: ThreadPool, message: any) {
+      message instanceof Object && READY in message ? onWorkerReady() : onEventMessage.call(this, message);
+    }
 
-    worker.on('error', (err : Error) => {
+    function onWorkerError (this: ThreadPool, err: Error) {
       this._onError(workerInfo, err, false);
-    });
+    }
 
-    worker.on('exit', (exitCode : number) => {
-      if (this.destroying) {
-        return;
+    function onWorkerExit (this: ThreadPool, code: number) {
+      if (this.destroying === false) {
+        const err = new Error(`worker exited with code: ${code}`);
+        // Only error unfinished tasks on process exit, since there are legitimate
+        // reasons to exit workers and we want to handle that gracefully when possible.
+        this._onError(workerInfo, err, true);
       }
-
-      const err = new Error(`worker exited with code: ${exitCode}`);
-      // Only error unfinished tasks on process exit, since there are legitimate
-      // reasons to exit workers and we want to handle that gracefully when possible.
-      this._onError(workerInfo, err, true);
-    });
-
-    worker.unref();
-    port1.on('close', () => {
-      // The port is only closed if the Worker stops for some reason, but we
-      // always .unref() the Worker itself. We want to receive e.g. 'error'
-      // events on it, so we ref it once we know it's going to exit anyway.
-      worker.ref();
-    });
-
-    this.workers.add(workerInfo);
+    }
   }
 
   _onError (workerInfo: WorkerInfo, err: Error, onlyErrorUnfinishedTasks: boolean) {

src/worker_pool/index.ts

@@ -1,8 +1,8 @@
-import { Worker, MessagePort, receiveMessageOnPort } from 'node:worker_threads';
+import { Worker, MessagePort, receiveMessageOnPort, WorkerOptions, Transferable } from 'node:worker_threads';
 import { createHistogram, RecordableHistogram } from 'node:perf_hooks';
 import assert from 'node:assert';
 
-import { RequestMessage, ResponseMessage } from '../types';
+import { RequestMessage, ResponseMessage, StartupMessage } from '../types';
 import { Errors } from '../errors';
 
 import { TaskInfo } from '../task_queue';
@@ -25,9 +25,10 @@ export type PiscinaWorker = {
 }
 
 type WorkerInfoParams = {
-  worker: Worker,
+  worker: {
+    filename: string,
+  } & WorkerOptions,
   port: MessagePort,
-  onMessage: ResponseCallback,
   enableHistogram: boolean,
 }
 
@@ -47,16 +48,16 @@ export class WorkerInfo extends AsynchronouslyCreatedResource {
       {
         worker,
         port,
-        onMessage,
         enableHistogram
-      }: WorkerInfoParams
+      }: WorkerInfoParams,
+      onMessage: ResponseCallback
     ) {
       super();
-      this.worker = worker;
+      const { filename, ...workerOpts } = worker;
+      this.worker = new Worker(filename, workerOpts);
       this.port = port;
-      this.port.on('message',
-        (message : ResponseMessage) => this._handleResponse(message));
       this.onMessage = onMessage;
+      this.port.on('message', this._handleResponse.bind(this));
       this.taskInfos = new Map();
       this.sharedBuffer = new Int32Array(
         new SharedArrayBuffer(kFieldCount * Int32Array.BYTES_PER_ELEMENT));
@@ -67,6 +68,37 @@ export class WorkerInfo extends AsynchronouslyCreatedResource {
       return this.worker.threadId;
     }
 
+    onWorkerMessage(handler: (msg: any) => void): void {
+      this.worker.on('message', handler);
+    }
+
+    onWorkerError(handler: (err: Error) => void): void {
+      this.worker.on('error', handler);
+    }
+
+    onWorkerExit(handler: (code: number) => void): void {
+      this.worker.on('exit', handler);
+    }
+
+    onPortClose(handler: () => void): void {
+      this.port.on('close', handler);
+    }
+
+    init(msg: StartupMessage, toTransfer: Transferable[]): WorkerInfo {
+      this.worker.postMessage(msg, toTransfer);
+      return this;
+    }
+
+    workerRef(): WorkerInfo {
+      this.worker.ref();
+      return this;
+    }
+
+    workerUnref(): WorkerInfo {
+      this.worker.unref();
+      return this;
+    }
+
     destroy () : void {
       if (this.terminating || this.destroyed) return;
 
@@ -101,16 +133,13 @@ export class WorkerInfo extends AsynchronouslyCreatedResource {
     }
 
     unref () : WorkerInfo {
-      // Note: Do not call ref()/unref() on the Worker itself since that may cause
-      // a hard crash, see https://github.com/nodejs/node/pull/33394.
       this.port.unref();
       return this;
     }
 
     _handleResponse (message : ResponseMessage) : void {
-      if (message.time != null) {
-        this.histogram?.record(PiscinaHistogramHandler.toHistogramIntegerNano(message.time));
-      }
+      // Both cannot be in different state if histogram enabled.
+      this.histogram?.record(PiscinaHistogramHandler.toHistogramIntegerNano(message?.time!));
 
       this.onMessage(message);
 
@@ -122,7 +151,9 @@ export class WorkerInfo extends AsynchronouslyCreatedResource {
     }
 
     postTask (taskInfo : TaskInfo) {
+      // Avoid duplicates
       assert(!this.taskInfos.has(taskInfo.taskId));
+      // Avoid posting when pool is shutting down or worker already destroyed
       assert(!this.terminating && !this.destroyed);
 
       const message : RequestMessage = {
@@ -135,22 +166,20 @@ export class WorkerInfo extends AsynchronouslyCreatedResource {
 
       try {
         this.port.postMessage(message, taskInfo.transferList);
+        queueMicrotask(() => this.clearIdleTimeout())
+        taskInfo.workerInfo = this;
+        this.taskInfos.set(taskInfo.taskId, taskInfo);
+        this.ref();
+
+        // Inform the worker that there are new messages posted, and wake it up
+        // if it is waiting for one.
+        Atomics.add(this.sharedBuffer, kRequestCountField, 1);
+        Atomics.notify(this.sharedBuffer, kRequestCountField, 1);
       } catch (err) {
         // This would mostly happen if e.g. message contains unserializable data
         // or transferList is invalid.
         taskInfo.done(<Error>err);
-        return;
       }
-
-      taskInfo.workerInfo = this;
-      this.taskInfos.set(taskInfo.taskId, taskInfo);
-      queueMicrotask(() => this.clearIdleTimeout())
-      this.ref();
-
-      // Inform the worker that there are new messages posted, and wake it up
-      // if it is waiting for one.
-      Atomics.add(this.sharedBuffer, kRequestCountField, 1);
-      Atomics.notify(this.sharedBuffer, kRequestCountField, 1);
     }
 
     processPendingMessages () {
@@ -166,7 +195,7 @@ export class WorkerInfo extends AsynchronouslyCreatedResource {
         this.lastSeenResponseCount = actualResponseCount;
 
         let entry;
-        while ((entry = receiveMessageOnPort(this.port)) !== undefined) {
+        while ((entry = receiveMessageOnPort(this.port)) != null) {
           this._handleResponse(entry.message);
         }
       }
@@ -176,18 +205,17 @@ export class WorkerInfo extends AsynchronouslyCreatedResource {
       // If there are abortable tasks, we are running one at most per Worker.
       if (this.taskInfos.size !== 1) return false;
       const [[, task]] = this.taskInfos;
-      return task.abortSignal !== null;
+      return task.abortSignal != null;
     }
 
     currentUsage () : number {
-      if (this.isRunningAbortableTask()) return Infinity;
-      return this.taskInfos.size;
+      return this.isRunningAbortableTask() ? Infinity : this.taskInfos.size;
     }
 
     popTask (taskId: number) : TaskInfo | null {
       const task = this.taskInfos.get(taskId) ?? null;
-      
-      this.taskInfos.delete(taskId);
+
+      if (task != null) this.taskInfos.delete(taskId);
 
       return task;
     }