@naturalcycles/nodejs-lib/dist/stream/transform/transformThrottleByRSS.js

Standard library for Node.js

import { Transform } from 'node:stream';
import { _mb } from '@naturalcycles/js-lib';
import { _ms, localTime } from '@naturalcycles/js-lib/datetime';
import { createCommonLoggerAtLevel } from '@naturalcycles/js-lib/log';
import { pDefer } from '@naturalcycles/js-lib/promise/pDefer.js';
/**
 * Throttles the stream based on process memory (RSS) usage.
 * When RSS exceeds `maxRSS` (in megabytes), the stream pauses
 * and periodically re-checks until RSS drops below the threshold.
 *
 * Useful for pipelines that process large amounts of data and
 * may cause memory pressure (e.g. database imports, file processing).
 *
 * @experimental
 */
export function transformThrottleByRSS(opt) {
    const { maxRSS, pollInterval = 5000, pollTimeout = 30 * 60_000, // 30 min
    onPollTimeout = 'open-the-floodgates', objectMode = true, highWaterMark = 1, } = opt;
    const maxRSSBytes = maxRSS * 1024 * 1024;
    let lock;
    let pollTimer;
    let rssCheckTimer;
    let lastRSS = 0;
    let pausedSince = 0;
    let disabled = false;
    const logger = createCommonLoggerAtLevel(opt.logger, opt.logLevel);
    return new Transform({
        objectMode,
        highWaterMark,
        async transform(item, _, cb) {
            if (lock) {
                try {
                    await lock;
                }
                catch (err) {
                    cb(err);
                    return;
                }
            }
            if (!disabled && lastRSS > maxRSSBytes && !lock) {
                lock = pDefer();
                pausedSince = Date.now();
                logger.log(`${localTime.now().toPretty()} transformThrottleByRSS paused: RSS ${_mb(lastRSS)} > ${maxRSS} MB`);
                pollTimer = setTimeout(() => pollRSS(), pollInterval);
            }
            cb(null, item);
        },
        construct(cb) {
            // Start periodic RSS checking
            checkRSS();
            cb();
        },
        final(cb) {
            clearTimeout(pollTimer);
            clearTimeout(rssCheckTimer);
            cb();
        },
    });
    function checkRSS() {
        lastRSS = process.memoryUsage.rss();
        rssCheckTimer = setTimeout(() => checkRSS(), pollInterval);
    }
    function pollRSS() {
        const rss = lastRSS;
        if (rss <= maxRSSBytes) {
            logger.log(`${localTime.now().toPretty()} transformThrottleByRSS resumed: RSS ${_mb(rss)} <= ${maxRSS} MB`);
            lock.resolve();
            lock = undefined;
        }
        else if (pollTimeout && Date.now() - pausedSince >= pollTimeout) {
            clearTimeout(rssCheckTimer);
            if (onPollTimeout === 'throw') {
                lock.reject(new Error(`transformThrottleByRSS pollTimeout of ${_ms(pollTimeout)} reached, RSS ${_mb(rss)} still > ${maxRSS} MB`));
                lock = undefined;
            }
            else {
                // open-the-floodgates
                logger.error(`${localTime.now().toPretty()} transformThrottleByRSS: pollTimeout of ${_ms(pollTimeout)} reached, RSS ${_mb(rss)} still > ${maxRSS} MB — DISABLING THROTTLE`);
                disabled = true;
                lock.resolve();
                lock = undefined;
            }
        }
        else {
            logger.log(`${localTime.now().toPretty()} transformThrottleByRSS still paused: RSS ${_mb(rss)} > ${maxRSS} MB, rechecking in ${_ms(pollInterval)}`);
            pollTimer = setTimeout(() => pollRSS(), pollInterval);
        }
    }
}