batchjs/dist/esm/streams/classes/ParallelStream.js

Batch processing framework for NodeJS

import { InternalBufferDuplex } from "../interfaces/_index";
/**
 * @class
 * Class that allows you to transform and stream data in parallel.
 * @extends InternalBufferDuplex
 * @template TInput The type of the input data.
 * @template TOutput The type of the output data.
 * @example
 * ```typescript
 * const stream:ParallelStream<string,string> = new ParallelStream({
 *     objectMode: true,
 *     maxConcurrent: 2,
 *     transform(chunk: string) {
 *         return Promise.resolve(chunk.toUpperCase());
 *     },
 * });
 *
 * stream.write("data1");
 * stream.write("data2");
 * stream.write("data3");
 * stream.end();
 *
 * stream.on("data", (chunk: string) => {
 *     console.log(``Pushed chunk: ${chunk}```);
 * });
 * ```
 * ```shell
 * >> Pushed chunk: DATA1
 * >> Pushed chunk: DATA2
 * >> Pushed chunk: DATA3
 * ```
 */
export class ParallelStream extends InternalBufferDuplex {
    queue = [];
    pool = new Set();
    maxConcurrent;
    transform;
    /**
     * @constructor
     * @param {ParallelStreamOptions<TInput, TOutput>} options - The options for the ParallelStream.
     * @param [options.maxConcurrent] {number} - The maximum number of concurrent promises.
     * @param [options.transform] {Function} - The function to transform the data returning a promise.
     */
    constructor(options) {
        super(options);
        this.maxConcurrent = options.maxConcurrent;
        this.transform = options.transform;
    }
    /**
     * A method to write data to the stream, push the chunk to the queue, transform it, and then execute the callback.
     *
     * @param {TInput} chunk - The data chunk to write to the stream.
     * @param {BufferEncoding} encoding - The encoding of the data.
     * @param {TransformCallback} callback - The callback function to be executed after writing the data.
     * @return {void} This function does not return anything.
     */
    _write(chunk, encoding, callback) {
        this.queue.push(chunk);
        this._transform();
        callback();
    }
    /**
     * Asynchronously finalizes the stream by draining the queue and buffer, pushing any remaining chunks to the stream,
     * and calling the provided callback when complete. If the stream is unable to push a chunk, the chunk is placed back
     * into the buffer and a PushError is passed to the callback.
     *
     * @override
     * @param {TransformCallback} callback - The callback to be called when the stream is finalized.
     * @return {Promise<void>} A promise that resolves when the stream is finalized.
     */
    _final(callback) {
        const awaitAllProcessed = () => {
            if (this.queue.length > 0 || this.pool.size > 0) {
                setImmediate(awaitAllProcessed);
            }
            else {
                super._final(callback);
            }
        };
        awaitAllProcessed();
    }
    /**
     * Loop through the pool and queue to process chunks, adding promises to the pool.
     */
    _transform() {
        while (this.pool.size < this.maxConcurrent && this.queue.length > 0) {
            const chunk = this.queue.shift(); // Get the next chunk
            const promise = this.transform(chunk)
                .then((result) => {
                this.buffer.push(result);
                this._flush();
            })
                .catch((err) => {
                this.emit("error", err);
            })
                .finally(() => {
                this.pool.delete(promise); // Remove promise from pool
                this._transform(); // Continue processing remaining chunks
            });
            this.pool.add(promise); // Add promise to pool
        }
    }
}