@naturalcycles/nodejs-lib/dist/stream/transform/worker/transformMultiThreaded.js

Standard library for Node.js

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.transformMultiThreaded = void 0;
const worker_threads_1 = require("worker_threads");
const js_lib_1 = require("@naturalcycles/js-lib");
const through2Concurrent = require("through2-concurrent");
const workerProxyFilePath = `${__dirname}/workerClassProxy.js`;
/**
 * Spawns a pool of Workers (threads).
 * Distributes (using round-robin, equally) all inputs over Workers.
 * Workers emit 1 output for each 1 input.
 * Output of Workers is passed down the stream. Order is RANDOM (since it's a multi-threaded environment).
 */
function transformMultiThreaded(opt) {
    const { workerFile, poolSize = 2, workerData } = opt;
    const maxConcurrency = opt.concurrency || poolSize;
    const highWaterMark = Math.max(16, maxConcurrency);
    console.log({
        poolSize,
        maxConcurrency,
        highWaterMark,
    });
    const workerDonePromises = [];
    const messageDonePromises = {};
    let index = -1; // input chunk index, will start from 0
    const workers = (0, js_lib_1._range)(0, poolSize).map(workerIndex => {
        workerDonePromises.push((0, js_lib_1.pDefer)());
        const worker = new worker_threads_1.Worker(workerProxyFilePath, {
            workerData: {
                workerIndex,
                workerFile,
                ...workerData,
            },
        });
        // const {threadId} = worker
        // console.log({threadId})
        worker.on('error', err => {
            console.error(`Worker ${workerIndex} error`, err);
            workerDonePromises[workerIndex].reject(err);
        });
        worker.on('exit', _exitCode => {
            // console.log(`Worker ${index} exit: ${exitCode}`)
            workerDonePromises[workerIndex].resolve(undefined);
        });
        worker.on('message', (out) => {
            // console.log(`Message from Worker ${workerIndex}:`, out)
            // console.log(Object.keys(messageDonePromises))
            // tr.push(out.payload)
            if (out.error) {
                messageDonePromises[out.index].reject(out.error);
            }
            else {
                messageDonePromises[out.index].resolve(out.payload);
            }
        });
        return worker;
    });
    return through2Concurrent.obj({
        maxConcurrency,
        highWaterMark,
        async final(cb) {
            try {
                // Push null (complete) to all sub-streams
                workers.forEach(worker => worker.postMessage(null));
                console.log(`transformMultiThreaded.final is waiting for all chains to be done`);
                await Promise.all(workerDonePromises);
                console.log(`transformMultiThreaded.final all chains done`);
                cb();
            }
            catch (err) {
                cb(err);
            }
        },
    }, async function transformMapFn(chunk, _, cb) {
        // Freezing the index, because it may change due to concurrency
        const currentIndex = ++index;
        // Create the unresolved promise (to avait)
        messageDonePromises[currentIndex] = (0, js_lib_1.pDefer)();
        const worker = workers[currentIndex % poolSize]; // round-robin
        worker.postMessage({
            index: currentIndex,
            payload: chunk,
        });
        try {
            // awaiting for result
            const out = await messageDonePromises[currentIndex];
            // console.log('awaited!')
            // return the result
            cb(null, out);
        }
        catch (err) {
            // Currently we only support ErrorMode.SUPPRESS
            // Error is logged and output continues
            console.error(err);
            cb(); // emit nothing in case of an error
        }
        // clean up
        delete messageDonePromises[currentIndex];
    });
}
exports.transformMultiThreaded = transformMultiThreaded;