@devteks/node-workers
Version:
Simple and easy to use worker pool implementation for Node.js
324 lines (317 loc) • 10.4 kB
JavaScript
/**
* @devteks/node-workers
* Simple and easy to use worker pool implementation for Node.js
* Version: 0.0.6
* Author: Mosa Muhana (https://github.com/mosamuhana)
* License: MIT
* Homepage: https://github.com/mosamuhana/node-workers#readme
*/
'use strict';
Object.defineProperty(exports, '__esModule', { value: true });
var os = require('os');
var async_hooks = require('async_hooks');
var events = require('events');
var worker_threads = require('worker_threads');
var path = require('path');
const UNLOCKED = 0;
const LOCKED = 1;
// copied from https://github.com/mosamuhana/node-atomics
class Mutex {
#array;
constructor(input) {
this.#array = new Int32Array(input || new SharedArrayBuffer(4));
}
get buffer() { return this.#array.buffer; }
lock() {
while (true) {
if (Atomics.compareExchange(this.#array, 0, UNLOCKED, LOCKED) === UNLOCKED)
return;
Atomics.wait(this.#array, 0, LOCKED);
}
}
unlock() {
if (Atomics.compareExchange(this.#array, 0, LOCKED, UNLOCKED) !== LOCKED) {
throw new Error("Inconsistent state: unlock on unlocked Mutex.");
}
Atomics.notify(this.#array, 0, 1);
}
sync(fn) {
try {
this.lock();
return fn();
}
finally {
this.unlock();
}
}
}
const CPUS = os.cpus().length;
const TERMINATE_PROMISE_SUPPORT = (() => {
const [major, minor] = process.version.replace("v", "").split(".").map((x) => parseInt(x, 10));
return major >= 12 && minor >= 5;
})();
const CLOSED_ERROR = new Error("WorkerPool is closed");
class WorkerError extends Error {
error;
task;
constructor(error, task) {
super((error?.message || error || "Unknown error").toString());
this.error = error;
this.task = task;
this.name = "WorkerError";
this.stack = error?.stack;
}
[Symbol.toStringTag]() { return "WorkerError"; }
}
class TaskInfo extends async_hooks.AsyncResource {
callback;
task;
constructor(callback, task) {
super("TaskInfo");
this.callback = callback;
this.task = task;
}
done(error, _result) {
let result;
if (error) {
error = new WorkerError(error, this.task);
}
else {
const res = _result;
if (res.status === 'fulfilled') {
result = res.value;
}
else if (res.status === 'rejected') {
error = new WorkerError(res.reason, this.task);
}
}
this.runInAsyncScope(this.callback, null, error, result);
this.emitDestroy();
}
}
class WorkerPool extends events.EventEmitter {
#scriptFile;
#eval = false;
#maxWorkers;
#workers = [];
#freeWorkers = [];
#tasks = [];
#eventMutex = new Mutex();
#mutex = new Mutex();
#closed = false;
#timeout = -1;
get maxWorkers() { return this.#maxWorkers; }
constructor(options) {
super();
const maxWorkers = options.maxWorkers;
if (maxWorkers == null) {
this.#maxWorkers = CPUS * 2;
}
else {
if (!Number.isInteger(maxWorkers) || maxWorkers < 1) {
throw new Error('maxWorkers must be a positive integer >= 1');
}
this.#maxWorkers = Math.min(CPUS * 2, maxWorkers);
}
if (options.workerFile) {
const filename = options.workerFile.replace(/\\/g, "/");
if (!/\.(c|m)?js|\.ts$/i.test(filename)) {
throw new Error("Worker file must be `.js`, `.mjs`, `.cjs` or `.ts`.");
}
this.#eval = path.extname(filename).toLowerCase() === '.ts';
this.#scriptFile = this.#eval ? `require('ts-node').register();require("${filename}");` : filename;
}
else if (options.workerScript) {
this.#scriptFile = options.workerScript;
this.#eval = true;
}
else {
throw new Error('workerFile or workerScript must be one specified');
}
const timeout = options.timeout;
if (timeout != null) {
if (!Number.isInteger(timeout) || timeout <= 0) {
throw new Error('timeout must be a positive integer');
}
this.#timeout = timeout;
}
}
#createWorker() {
const { port1, port2 } = new worker_threads.MessageChannel();
const worker = new worker_threads.Worker(this.#scriptFile, {
eval: this.#eval,
workerData: { lock: this.#eventMutex.buffer, port: port1, timeout: this.#timeout },
transferList: [port1],
});
worker.port = port2;
return worker;
}
#addNewWorker() {
const worker = this.#createWorker();
worker.port.on("message", ({ event, message }) => this.emit(event, message));
worker.on("message", (result) => {
const taskInfo = worker.taskInfo;
worker.taskInfo = undefined;
taskInfo.done(undefined, result);
this.#mutex.sync(() => this.#freeWorkers.push(worker));
this.#runNext();
});
worker.on("error", (error) => {
if (worker.taskInfo) {
worker.taskInfo.done(error);
}
else {
this.emit("error", error);
}
worker.port.removeAllListeners();
worker.removeAllListeners();
this.#mutex.sync(() => this.#workers.splice(this.#workers.indexOf(worker), 1));
this.#runNext();
});
this.#mutex.sync(() => {
this.#workers.push(worker);
this.#freeWorkers.push(worker);
});
}
#runNext() {
const task = this.#tasks.shift();
if (task) {
this.#run(task.task, task.callback);
}
}
#run(task, callback) {
if (this.#closed)
throw CLOSED_ERROR;
const worker = this.#mutex.sync(() => this.#freeWorkers.shift());
if (worker) {
worker.taskInfo = new TaskInfo(callback, task);
worker.postMessage(task ?? {});
}
else {
this.#tasks.push({ task, callback });
if (this.#workers.length < this.#maxWorkers) {
this.#addNewWorker();
this.#runNext();
}
}
}
#runOne(task) {
return new Promise((resolve, reject) => {
this.#run(task, (error, result) => {
if (error)
return reject(error);
resolve(result);
});
});
}
async #runMany(tasks) {
const all = await Promise.allSettled(tasks.map(task => this.#runOne(task)));
const errors = all.filter(result => result.status === "rejected")
.map(x => x.reason);
const results = all.filter(result => result.status === "fulfilled")
.map(x => x.value);
return { results, errors };
}
run(input, callback) {
if (typeof callback === "function") {
this.#run(input, callback);
}
else {
return Array.isArray(input) ?
this.#runMany(input) :
this.#runOne(input);
}
}
static async run(options, input, emit) {
if (input == null)
throw new Error("task must be defined");
if (typeof options.maxWorkers === 'undefined') {
options.maxWorkers = Array.isArray(input) ? input.length : 1;
}
const pool = new WorkerPool(options);
if (typeof emit === 'function') {
pool.on('message', message => emit(message));
}
try {
if (Array.isArray(input)) {
return await pool.run(input);
}
else {
return await pool.run(input);
}
}
finally {
await pool.close();
}
}
async close() {
this.#mutex.lock();
try {
if (!this.#closed) {
this.#closed = true;
await Promise.allSettled(this.#workers.map(closeWorker));
this.#workers = [];
this.#freeWorkers = [];
}
}
finally {
this.#mutex.unlock();
}
}
}
function startWorker(fn) {
if (worker_threads.isMainThread) {
throw new Error("startWorker can only be used in a worker thread.");
}
const timeout = worker_threads.workerData.timeout ?? 0;
const mutex = new Mutex(worker_threads.workerData.lock);
const port = worker_threads.workerData.port;
const emit = (event, message) => {
mutex.sync(() => port.postMessage({ event, message }));
};
worker_threads.parentPort.on("message", async (request) => {
try {
const value = timeout > 0 ?
await timedFn(() => fn(request, emit), timeout) :
await fn(request, emit);
worker_threads.parentPort.postMessage({ status: "fulfilled", value });
}
catch (reason) {
worker_threads.parentPort.postMessage({ status: "rejected", reason });
}
});
}
async function timedFn(fn, timeout) {
let t;
const result = await Promise.race([
new Promise((_, reject) => {
t = setTimeout(() => reject(new Error("timeout")), timeout);
}),
fn().then(res => {
clearTimeout(t);
return res;
}),
]);
return result;
}
async function closeWorker(worker) {
try {
worker.port.close();
}
catch (ex) { }
if (TERMINATE_PROMISE_SUPPORT) {
try {
await worker.terminate();
}
catch (ex) { }
}
else {
try {
await new Promise((resolve) => worker.terminate(resolve));
}
catch (ex) { }
}
}
exports.WorkerError = WorkerError;
exports.WorkerPool = WorkerPool;
exports.startWorker = startWorker;