Zip Back-pressure Handling #5128

import { appendFileSync, createReadStream, mkdirSync } from "node:fs";
import { Transform } from "node:stream";
import { pipeline } from "node:stream/promises";
import { Gzip, createGzip } from "node:zlib";

async function streamTodo(csvFilePath: string): Promise<void> {

    type TSplitStreamContext = {
        zipper?: Gzip
    };

    const tagsStreamContext = new Map<string, TSplitStreamContext>();

    const inputFileStream = createReadStream(csvFilePath, { encoding: 'utf8' });

    let chunkCounter = 0;

    const zipStream = new Transform({
        objectMode: true,
        transform(tagWiseCSVChunks: string, encoding, callback) {
            //const tagIdVSChunkRecords = Object.entries(tagWiseCSVChunks);
            chunkCounter++;
            let tagIdVSChunkRecords = new Map<string, string>();
            if (chunkCounter % 2 === 0) {
                tagIdVSChunkRecords.set("even", tagWiseCSVChunks);
            }
            else {
                tagIdVSChunkRecords.set("odd", tagWiseCSVChunks);
            }

            let upstreamPressure = 0;
            let downstreamBackPressure = 0;
            let invokeCallback = true;
            const invocationName = chunkCounter;

            for (const [tagId, csvChunk] of tagIdVSChunkRecords) {
                const existingContext = tagsStreamContext.get(tagId) ?? {} as TSplitStreamContext;

                if (existingContext.zipper == null || existingContext.zipper === undefined) {
                    existingContext.zipper = createGzip();

                    existingContext.zipper.on('data', (chunk: Buffer) => {
                        this.push({ [tagId]: chunk });
                        // if (this.push({ [tagId]: chunk }) === false) {
                        //     downstreamBackPressure++;
                        //     existingContext.zipper!.pause();
                        //     this.once('drain', () => {
                        //         downstreamBackPressure--;
                        //         if (downstreamBackPressure === 0) {
                        //             existingContext.zipper!.resume();
                        //         }
                        //     });
                        // }
                    });

                    tagsStreamContext.set(tagId, existingContext);
                }
                //existingContext.zipper.write(csvChunk);
                if (existingContext.zipper.write(csvChunk) === false) {
                    invokeCallback = false;
                    upstreamPressure++;
                    existingContext.zipper.once('drain', () => {
                        upstreamPressure--;
                        if (upstreamPressure === 0) {
                            if (chunkCounter < 10) callback();//This controls upstream flow
                            console.log(`[${invocationName}] Callback invoked after drain for invocation ${invocationName}`);
                        }
                    });
                }
            }
            if (invokeCallback === true) {
                callback();//This controls upstream flow
            }
            else {
                console.log(`[${invocationName}] Callback delayed due to upstream pressure in invocation ${invocationName}`);
            }

        },

        final(callback) {
            const promiseHandles = [];

            for (const [tagId, context] of tagsStreamContext) {
                if (context.zipper !== null && context.zipper !== undefined) {
                    promiseHandles.push(new Promise<void>((resolve, reject) => {
                        context.zipper!.once('end', resolve);
                        context.zipper!.once('error', reject);
                        context.zipper!.end();
                    }));
                }
            }

            Promise.all(promiseHandles)
                .then(() => { callback(); console.log('All zippers ended.'); })
                .catch(err => callback(err));
        }
    });

    const fileWriter = new Transform({
        objectMode: true,
        transform(chunk: Record<string, Buffer>, encoding, callback) {
            for (const [tagId, zippedBufferChunk] of Object.entries(chunk)) {
                appendFileSync(`./zipped/${tagId}.gz`, zippedBufferChunk, { flag: 'a' });
            }
            callback();
        }
    });

    await pipeline(
        inputFileStream,
        zipStream,
        fileWriter
    );

    tagsStreamContext.clear();
}

async function mainModule() {
    mkdirSync('./zipped', { recursive: true })
    console.log('Computing...');
    await streamTodo('./1.csv');
}


mainModule()
    .then(() => console.log('All operations completed successfully.'))
    .catch(err => console.error('Error during operations:', err));
//node --experimental-strip-types ./hello.mts