shuzhiren-comfyui/任务队列后端/worker_threads/process/PollingThreadPool.js

88 lines
2.4 KiB
JavaScript
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import { Worker } from 'worker_threads';
// 创建固定大小的线程池6个线程
class QueryThreadPool {
constructor(size = 6) {
this.size = size; // 线程池大小默认为6
this.workers = []; // 存储worker对象的数组
this.taskQueue = []; // 任务队列,存储等待处理的任务
this.initWorkers(); // 初始化worker线程
}
initWorkers() {
for (let i = 0; i < this.size; i++) {
const worker = new Worker(new URL('./pollingTask.js', import.meta.url)); // 创建指向generatTask.js的worker
this.workers.push({
worker, // worker对象
busy: false // 是否正在处理任务的标志
});
}
}
async executeTask(taskBatch) {
return new Promise((resolve, reject) => {
// 寻找空闲的worker
const availableWorker = this.workers.find(w => !w.busy);
if (availableWorker) {
// 有空闲worker直接执行任务
this.runTaskOnWorker(availableWorker, taskBatch, resolve, reject);
} else {
// 没有空闲worker将任务加入队列
this.taskQueue.push({ taskBatch, resolve, reject });
}
});
}
runTaskOnWorker(workerObj, taskBatch, resolve, reject) {
workerObj.busy = true;
workerObj.worker.once('message', (result) => {
workerObj.busy = false;
resolve(result);
// 检查是否有排队的任务
this.processQueuedTasks();
});
workerObj.worker.once('error', (error) => {
workerObj.busy = false;
reject(error);
// 检查是否有排队的任务
this.processQueuedTasks();
});
workerObj.worker.postMessage(taskBatch);
}
processQueuedTasks() {
if (this.taskQueue.length > 0) {
const availableWorker = this.workers.find(w => !w.busy);
if (availableWorker) {
const queuedTask = this.taskQueue.shift();
this.runTaskOnWorker(
availableWorker,
queuedTask.taskBatch,
queuedTask.resolve,
queuedTask.reject
);
}
}
}
async executeAllTasks(taskBatches) {
const promises = taskBatches.map(batch => this.executeTask(batch));
return Promise.all(promises);
}
// 清理资源
terminate() {
this.workers.forEach(({ worker }) => {
worker.terminate();
});
}
}
export default new QueryThreadPool();