@hirosystems/api-toolkit/dist/helpers/worker-thread-manager.js

API development toolkit

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.WorkerThreadManager = void 0;
const WorkerThreads = require("node:worker_threads");
const os = require("node:os");
const node_events_1 = require("node:events");
const time_1 = require("./time");
const serialize_error_1 = require("./serialize-error");
const worker_thread_init_1 = require("./worker-thread-init");
class WorkerThreadManager {
    workers = new Set();
    idleWorkers = [];
    jobQueue = [];
    msgRequests = new Map();
    lastMsgId = 0;
    workerCount;
    workerFile;
    abortControlller = new AbortController();
    events = new node_events_1.EventEmitter();
    get idleWorkerCount() {
        return this.idleWorkers.length;
    }
    get busyWorkerCount() {
        return this.workerCount - this.idleWorkers.length;
    }
    get queuedJobCount() {
        return this.jobQueue.length;
    }
    static init(workerModule, opts = {}) {
        const workerManager = new WorkerThreadManager(workerModule, opts);
        return new Promise(resolve => {
            workerManager.events.once('workersReady', () => {
                resolve(workerManager);
            });
        });
    }
    constructor(workerModule, opts = {}) {
        if (!WorkerThreads.isMainThread) {
            throw new Error(`${this.constructor.name} must be instantiated in the main thread`);
        }
        if ('default' in workerModule) {
            this.workerFile = workerModule.default.workerModule.filename;
        }
        else {
            this.workerFile = workerModule.workerModule.filename;
        }
        this.workerCount = opts.workerCount ?? os.cpus().length;
        this.createWorkerPool();
    }
    exec(...args) {
        this.abortControlller.signal.throwIfAborted();
        if (this.lastMsgId >= Number.MAX_SAFE_INTEGER) {
            this.lastMsgId = 0;
        }
        const msgId = this.lastMsgId++;
        const replyWaiter = (0, time_1.waiter)();
        this.msgRequests.set(msgId, replyWaiter);
        const reqMsg = {
            msgId,
            req: args,
        };
        this.jobQueue.push(reqMsg);
        this.assignJobs();
        return replyWaiter;
    }
    createWorkerPool() {
        let workersReady = 0;
        for (let i = 0; i < this.workerCount; i++) {
            const workerData = {
                workerFile: this.workerFile,
            };
            const workerOpt = {
                workerData,
            };
            const hasTsSource = worker_thread_init_1.filename.endsWith('.ts') || this.workerFile.endsWith('.ts');
            if (hasTsSource) {
                if (process.env.NODE_ENV !== 'test') {
                    throw new Error('Worker threads are being created with ts-node outside of a test environment');
                }
                workerOpt.execArgv = ['-r', 'ts-node/register/transpile-only'];
            }
            const worker = new WorkerThreads.Worker(worker_thread_init_1.filename, workerOpt);
            worker.unref();
            this.workers.add(worker);
            worker.on('error', err => {
                console.error(`Worker error`, err);
            });
            worker.on('messageerror', err => {
                console.error(`Worker message error`, err);
            });
            worker.once('message', (message) => {
                if (message !== 'ready') {
                    throw new Error(`Unexpected first msg from worker thread: ${JSON.stringify(message)}`);
                }
                this.setupWorkerHandler(worker);
                this.idleWorkers.push(worker);
                this.assignJobs();
                workersReady++;
                if (workersReady === this.workerCount) {
                    this.events.emit('workersReady');
                }
            });
        }
        (0, node_events_1.addAbortListener)(this.abortControlller.signal, () => {
            for (const replyWaiter of this.msgRequests.values()) {
                // eslint-disable-next-line @typescript-eslint/no-unsafe-argument
                replyWaiter.reject(this.abortControlller.signal.reason);
            }
            this.msgRequests.clear();
        });
    }
    setupWorkerHandler(worker) {
        worker.on('message', (message) => {
            this.idleWorkers.push(worker);
            this.assignJobs();
            const msg = message;
            const replyWaiter = this.msgRequests.get(msg.msgId);
            if (replyWaiter) {
                if (msg.error) {
                    const error = (0, serialize_error_1.isErrorLike)(msg.error) ? (0, serialize_error_1.deserializeError)(msg.error) : msg.error;
                    replyWaiter.reject(error);
                }
                else if (msg.resp) {
                    replyWaiter.resolve(msg.resp);
                }
                this.msgRequests.delete(msg.msgId);
            }
            else {
                console.error('Received unexpected message from worker', msg);
            }
        });
    }
    assignJobs() {
        while (this.idleWorkers.length > 0 && this.jobQueue.length > 0) {
            // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
            const worker = this.idleWorkers.shift();
            const job = this.jobQueue.shift();
            worker.postMessage(job);
        }
    }
    async close() {
        this.abortControlller.abort();
        await Promise.all([...this.workers].map(worker => worker.terminate()));
        this.workers.clear();
    }
}
exports.WorkerThreadManager = WorkerThreadManager;
//# sourceMappingURL=worker-thread-manager.js.map