repomix/lib/shared/processConcurrency.js

A tool to pack repository contents to single file for AI consumption

import os from 'node:os';
import { Tinypool } from 'tinypool';
import { logger } from './logger.js';
const getWorkerPath = (workerType) => {
    if (process.env.REPOMIX_WORKER_PATH) {
        return process.env.REPOMIX_WORKER_PATH;
    }
    switch (workerType) {
        case 'fileProcess':
            return new URL('../core/file/workers/fileProcessWorker.js', import.meta.url).href;
        case 'securityCheck':
            return new URL('../core/security/workers/securityCheckWorker.js', import.meta.url).href;
        case 'calculateMetrics':
            return new URL('../core/metrics/workers/calculateMetricsWorker.js', import.meta.url).href;
        default:
            throw new Error(`Unknown worker type: ${workerType}`);
    }
};
const TASKS_PER_THREAD = 100;
export const getProcessConcurrency = () => {
    return typeof os.availableParallelism === 'function' ? os.availableParallelism() : os.cpus().length;
};
export const getWorkerThreadCount = (numOfTasks, maxWorkerThreads) => {
    const processConcurrency = getProcessConcurrency();
    const minThreads = 1;
    const effectiveConcurrency = maxWorkerThreads != null ? Math.min(processConcurrency, maxWorkerThreads) : processConcurrency;
    const maxThreads = Math.max(minThreads, Math.min(effectiveConcurrency, Math.ceil(numOfTasks / TASKS_PER_THREAD)));
    return {
        minThreads,
        maxThreads,
    };
};
export const createWorkerPool = (options) => {
    const { numOfTasks, workerType, runtime = 'child_process', maxWorkerThreads } = options;
    const { minThreads, maxThreads } = getWorkerThreadCount(numOfTasks, maxWorkerThreads);
    const workerPath = getWorkerPath(workerType);
    logger.trace(`Initializing worker pool with min=${minThreads}, max=${maxThreads} threads, runtime=${runtime}. Worker type: ${workerType}`);
    const startTime = process.hrtime.bigint();
    const pool = new Tinypool({
        filename: workerPath,
        runtime,
        minThreads,
        maxThreads,
        idleTimeout: 5000,
        teardown: 'onWorkerTermination',
        workerData: {
            workerType,
            logLevel: logger.getLogLevel(),
        },
        ...(runtime === 'child_process' && {
            env: {
                ...process.env,
                REPOMIX_WORKER_TYPE: workerType,
                REPOMIX_LOG_LEVEL: logger.getLogLevel().toString(),
                FORCE_COLOR: process.env.FORCE_COLOR || (process.stdout.isTTY ? '1' : '0'),
                TERM: process.env.TERM || 'xterm-256color',
                COLUMNS: process.env.COLUMNS || process.stdout.columns?.toString() || '',
            },
        }),
    });
    const endTime = process.hrtime.bigint();
    const initTime = Number(endTime - startTime) / 1e6;
    logger.debug(`Tinypool initialization took ${initTime.toFixed(2)}ms`);
    return pool;
};
export const cleanupWorkerPool = async (pool) => {
    try {
        logger.debug('Cleaning up worker pool...');
        const isBun = process.versions?.bun;
        if (isBun) {
            logger.debug('Running in Bun environment, skipping Tinypool destroy method');
        }
        else {
            await pool.destroy();
        }
        logger.debug('Worker pool cleaned up successfully');
    }
    catch (error) {
        logger.debug('Error during worker pool cleanup:', error);
    }
};
export const initTaskRunner = (options) => {
    const pool = createWorkerPool(options);
    return {
        run: (task) => pool.run(task),
        cleanup: () => cleanupWorkerPool(pool),
    };
};