Skip to content

refactor: small adjustments over WorkerInfo #860

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
metcoder95 merged 2 commits into from
Jul 29, 2025
Merged

refactor: small adjustments over WorkerInfo #860

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3bc9750
refactor: small adjustments over WorkerInfo
metcoder95 Jul 27, 2025
91742b8
refactor: Apply suggestion from @Copilot
metcoder95 Jul 29, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
108 changes: 49 additions & 59 deletions src/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand All @@ -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) {
Expand Down
90 changes: 59 additions & 31 deletions src/worker_pool/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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';
Expand All @@ -25,9 +25,10 @@ export type PiscinaWorker = {
}

type WorkerInfoParams = {
worker: Worker,
worker: {
filename: string,
} & WorkerOptions,
port: MessagePort,
onMessage: ResponseCallback,
enableHistogram: boolean,
}

Expand All @@ -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));
Expand All @@ -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;

Expand Down Expand Up @@ -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!));

Comment thread
metcoder95 marked this conversation as resolved.
this.onMessage(message);

Expand All @@ -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 = {
Expand All @@ -135,22 +166,20 @@ export class WorkerInfo extends AsynchronouslyCreatedResource {

try {
this.port.postMessage(message, taskInfo.transferList);
queueMicrotask(() => this.clearIdleTimeout())
Comment thread
metcoder95 marked this conversation as resolved.
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 () {
Expand All @@ -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);
}
}
Expand All @@ -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: string) : TaskInfo | null {
const task = this.taskInfos.get(taskId) ?? null;
this.taskInfos.delete(taskId);

if (task != null) this.taskInfos.delete(taskId);
Comment thread
metcoder95 marked this conversation as resolved.

return task;
}
Expand Down
Loading