使用nestjs封装线程/进程池服务

简单使用，大佬绕路。

主要使用nodejs的worker_threads模块，worker_threads 是用于创建多线程而非多进程，但是也可以用来创建多进程，思路上是一样的，我这里主要用于多进程。
worker_threads 是 Node.js 的核心模块，不用额外引入。

创建worker.js线程

这里只做模拟，可以改成自己的逻辑

const { parentPort } = require('worker_threads');

// 模拟耗时操作
// 假设这是你的计算密集型任务 
function doWork(data) {
    return new Promise((resolve) => {
        setTimeout(() => {
            resolve(`Processed ${data}`);
        }, Math.random() * 3000);
    });
}

// parentPort.on用来接收从服务发送过来的参数
parentPort.on('message', async (data) => {
    try {
        const result = await doWork(data);
        parentPort.postMessage(result);
    } catch (error) {
        parentPort.postMessage({ error: error.message });
    }
});

封装线程池服务pool.service.ts

命令：nest g s pool
代码如下，有注释就不做过多解释了(代码写的不够优雅，可自行修改):

import { Injectable } from '@nestjs/common';
import { Worker, isMainThread, parentPort, workerData } from 'worker_threads';

@Injectable()
export class PoolService {
    private workers: Worker[] = [];//队列
    private poolSize = 4; // 可以根据 CPU 核心数动态调整  

    private workerUrl = 'your_url'
    private waitTask = [];//等待队列


    constructor() {
        if (!isMainThread) {
            // 如果不是在主线程，则运行 worker.js  
            require(this.workerUrl);
            return;
        }

        this.initPool();

        // 模拟使用
        setTimeout(async () => {
            for (let i = 0; i < 10; i++) {
                (async () => {
                    let res = await this.myPerformWork('hello' + i)
                    console.log(res)
                })()
            }
        }, 1000);
    }


    // 初始化
    private initPool() {
        for (let i = 0; i < this.poolSize; i++) {
            const worker = new Worker(this.workerUrl);
            worker.on('message', (message) => {
                // console.log(message)
            });
            worker.on('error', (err) => {
                console.error('Worker error:', err);
            });
            worker.on('exit', (code) => {
                if (code !== 0)
                    console.error(new Error(`Worker stopped with exit code ${code}`));
            });
            this.workers.push(worker);
        }
    }

    // 嵌套一层，为了将参数处理成统一格式
    public async myPerformWork(data: string): Promise<any> {
        return new Promise((resolve, reject) => {
            let newData = { data, resolve, reject };
            this.performWork(newData);
        })
    }

    private async performWork(data: any): Promise<any> {
        return new Promise((resolve, reject) => {
            // 这里简单使用第一个空闲 worker，实际中可能需要更复杂的调度策略  
            if (this.workers.length > 0) {
                const worker = this.workers.shift()!;
                worker.postMessage(data.data);

                worker.once('message', (message) => {
                    resolve(message);
                    data.resolve(message)
                    this.workers.push(worker); // 将 worker 放回池中
                    if (this.waitTask.length) {
                        this.performWork(this.waitTask.shift())
                    }
                });
                worker.once('error', (err) => {
                    reject(err);
                    data.reject();
                    // 可能需要移除或重启出错的 worker  
                });
            } else {
                this.waitTask.push(data)
            }
        });
    }
}

使用

就按照依赖注入的方式使用就行

constructor(private readonly poolService:PoolService) {}
async performWork(): Promise<string> {  
    return this.workerPoolService.myPerformWork('Hello');  
  }  

结束

版本

1. nodejs:20.9.0
2. npm:10.1.0
3. nestjs:10.0.0