@message-queue-toolkit/kafka/dist/utils/KafkaMessageBatchStream.js

Kafka adapter for message-queue-toolkit

import { Duplex } from 'node:stream';
/**
 * A Duplex stream that batches Kafka messages based on size and timeout constraints.
 *
 * Key features:
 * - Accumulates messages across all partitions up to `batchSize` for true memory control
 * - Groups messages by topic-partition when flushing
 * - Implements backpressure: pauses input when downstream consumers are overwhelmed
 * - Auto-flushes on timeout to prevent messages from waiting indefinitely
 *
 * @example
 * ```typescript
 * const batchStream = new KafkaMessageBatchStream({ batchSize: 100, timeoutMilliseconds: 1000 })
 * batchStream.on('data', (batch) => {
 *   console.log(`Received ${batch.length} messages from ${batch[0].topic}:${batch[0].partition}`)
 * })
 * ```
 */
// biome-ignore lint/suspicious/noUnsafeDeclarationMerging: merging interface with class to add strong typing for 'data' event
export class KafkaMessageBatchStream extends Duplex {
    batchSize;
    timeout;
    messages;
    existingTimeout;
    pendingCallback;
    isBackPressured;
    constructor(options) {
        super({ objectMode: true, readableHighWaterMark: options.readableHighWaterMark });
        this.batchSize = options.batchSize;
        this.timeout = options.timeoutMilliseconds;
        this.messages = [];
        this.isBackPressured = false;
    }
    /**
     * Called when the downstream consumer is ready to receive more data.
     * This is the backpressure release mechanism: we resume the writable side
     * by calling the pending callback that was held during backpressure.
     */
    _read() {
        this.isBackPressured = false;
        if (!this.pendingCallback)
            return;
        const cb = this.pendingCallback;
        this.pendingCallback = undefined;
        cb(); // Resume the writable side
    }
    /**
     * Writes a message to the stream.
     * Messages accumulate until batchSize is reached or timeout expires.
     * Implements backpressure by holding the callback when downstream cannot consume.
     */
    _write(message, _encoding, callback) {
        let canContinue = true;
        try {
            this.messages.push(message);
            if (this.messages.length >= this.batchSize) {
                // Batch is full, flush immediately
                canContinue = this.flushMessages();
            }
            else {
                // Start/continue the timeout for partial batches
                // Using ??= ensures we only set one timeout at a time
                this.existingTimeout ??= setTimeout(() => this.flushMessages(), this.timeout);
            }
        }
        finally {
            // Backpressure handling: hold the callback if push() returned false
            if (!canContinue)
                this.pendingCallback = callback;
            else
                callback();
        }
    }
    _final(callback) {
        // Clean timeout
        clearTimeout(this.existingTimeout);
        this.existingTimeout = undefined;
        // If there are remaining messages -> skip them
        // As they are not committed, the next consumer will process them
        this.messages = [];
        this.push(null);
        callback();
    }
    flushMessages() {
        clearTimeout(this.existingTimeout);
        this.existingTimeout = undefined;
        if (this.isBackPressured) {
            this.existingTimeout = setTimeout(() => this.flushMessages(), this.timeout);
            return false;
        }
        // Extract all accumulated messages and clear the array
        const messageBatch = this.messages.splice(0, this.messages.length);
        // Group by topic-partition to maintain commit guarantees
        const messagesByTopicPartition = {};
        for (const message of messageBatch) {
            const key = getTopicPartitionKey(message.topic, message.partition);
            if (!messagesByTopicPartition[key])
                messagesByTopicPartition[key] = [];
            messagesByTopicPartition[key].push(message);
        }
        // Push each topic-partition batch and track backpressure.
        // All batches must be pushed regardless: messages were already splice'd from the buffer,
        // so breaking early would lose them. Once push() returns false, subsequent calls in the
        // same tick also return false, so the last value correctly reflects backpressure.
        let canContinue = true;
        for (const messagesForKey of Object.values(messagesByTopicPartition)) {
            canContinue = this.push(messagesForKey);
        }
        if (!canContinue)
            this.isBackPressured = true;
        return canContinue;
    }
}
const getTopicPartitionKey = (topic, partition) => `${topic}:${partition}`;
//# sourceMappingURL=KafkaMessageBatchStream.js.map