寻找Node.js工作线程结果回传主线程流的设计模式

核心思路是通过自定义Transform流作为线程池和下游流的粘合剂，将readline读取的每行数据提交到Piscina线程池，任务完成后立即将结果推送到下一个流，同时利用背压机制控制任务提交速度，保证线程池始终忙碌，无需手动等待单个任务结果。

实现代码 #

1. 主进程代码 #

const fs = require('fs');
const readline = require('readline');
const { Transform, pipeline, Readable, Writable } = require('stream');
const Piscina = require('piscina');

// 初始化16线程的Piscina线程池
const piscina = new Piscina({
  filename: './worker.js',
  maxThreads: 16,
});

// 自定义Transform流：对接线程池和下游流
class WorkerPoolTransform extends Transform {
  constructor(options) {
    super({ ...options, objectMode: true });
    this.activeTasks = 0;
    this.maxActiveTasks = piscina.options.maxThreads;
    this.pendingCallbacks = [];
  }

  _transform(line, encoding, callback) {
    // 达到线程池上限时，暂存回调等待空闲线程
    if (this.activeTasks >= this.maxActiveTasks) {
      this.pendingCallbacks.push(callback);
      return;
    }

    this.activeTasks++;
    // 提交任务到线程池，不阻塞当前流程
    piscina.run(line)
      .then(result => this.push(result)) // 任务完成立即推结果到下游
      .catch(err => this.emit('error', err))
      .finally(() => {
        this.activeTasks--;
        // 有等待的回调则触发，允许处理下一行数据
        if (this.pendingCallbacks.length) {
          this.pendingCallbacks.shift()();
        } else {
          callback();
        }
      });
  }

  // 等待所有活跃任务完成后再关闭流
  _flush(callback) {
    const waitForTasks = () => {
      if (this.activeTasks === 0) callback();
      else setTimeout(waitForTasks, 100);
    };
    waitForTasks();
  }
}

// 将readline转换为可读取流
const readStream = fs.createReadStream('large.ndjson');
const rl = readline.createInterface({ input: readStream, crlfDelay: Infinity });
const lineStream = new Readable({ objectMode: true, read() {} });

rl.on('line', line => {
  if (!lineStream.push(line)) rl.pause(); // 背压触发时暂停读取
});
rl.on('close', () => lineStream.push(null));
lineStream.on('drain', () => rl.resume()); // 流恢复时继续读取

// 自定义可写流：将结果分发到多个文件
const multiFileWriter = new Writable({
  objectMode: true,
  write(result, _, callback) {
    // 示例：根据结果字段决定写入目标文件
    const targetFile = `output-${result.group}.txt`;
    fs.appendFile(targetFile, JSON.stringify(result) + '\n', callback);
  }
});

// 串联整个处理流程
pipeline(
  lineStream,
  new WorkerPoolTransform(),
  multiFileWriter,
  err => {
    if (err) console.error('处理失败:', err);
    else console.log('所有数据处理完成');
  }
);

2. 工作线程脚本（worker.js） #

module.exports = async (line) => {
  // 这里替换为你的实际数据处理逻辑
  const data = JSON.parse(line);
  // 模拟耗时操作（比如数据清洗、计算）
  await new Promise(resolve => setTimeout(resolve, 50));
  // 返回处理后的结果（可携带用于分发的字段）
  return { ...data, processedAt: Date.now(), group: data.category };
};

关键逻辑说明 #

• 线程池任务调度：自定义Transform流通过activeTasks跟踪正在运行的任务数，达到线程池上限时暂存回调，避免任务积压；线程空闲时自动取出暂存的回调，继续提交新任务，保证线程池始终满载。
• 即时结果推送：每个线程任务完成后立即调用this.push(result)，结果直接进入下游写入流，无需等待其他任务，完全利用流式处理的优势。
• 背压控制：当下游写入流处理速度跟不上时，lineStream.push()会返回false，此时暂停readline读取，直到流触发drain事件再恢复，防止内存溢出。
• 无需MessageChannel：Piscina已经封装了线程间的通信逻辑，直接通过piscina.run()的Promise接收结果即可，无需手动管理MessageChannel。

内容的提问来源于stack exchange，提问作者Dusty211