115 lines
4.4 KiB
JavaScript
115 lines
4.4 KiB
JavaScript
import { parentPort } from 'worker_threads';
|
||
import redis from '../../redis/index.js';
|
||
import initQueue from '../../redis/initQueue.js';
|
||
import { externalGetRequest } from '../../outside/polling.js';
|
||
|
||
async function getTask(tasks) {
|
||
console.log(`[pollingTask] 开始处理 ${Object.keys(tasks).length} 个轮询任务`);
|
||
console.log(`[pollingTask] 轮询任务数据: ${JSON.stringify(tasks)}`);
|
||
const queryTasks = []
|
||
for (const [remoteTaskId, value] of Object.entries(tasks)) {
|
||
console.log(`[pollingTask] 准备查询任务: remoteTaskId=${remoteTaskId}, value=${value}`);
|
||
// 查询外部平台任务结果
|
||
const queryTaskPromise = externalGetRequest(remoteTaskId, value) // { platform, taskid, AIGC }
|
||
queryTasks.push(queryTaskPromise)
|
||
}
|
||
|
||
try {
|
||
const responseTasks = await Promise.all(queryTasks)
|
||
console.log(`[pollingTask] 轮询查询完成，收到 ${responseTasks.length} 个响应`);
|
||
console.log(`[pollingTask] 轮询响应详情: ${JSON.stringify(responseTasks)}`);
|
||
return responseTasks
|
||
} catch (error) {
|
||
console.error('[pollingTask] 轮询查询出错:', error);
|
||
}
|
||
}
|
||
|
||
// 批量将完成的任务移动到结果队列中
|
||
async function storeSuccessTasks(SuccessTasks) {
|
||
console.log(`[pollingTask] 开始存储 ${SuccessTasks.length} 个成功完成的任务到结果队列`);
|
||
const taskCountMap = new Map();
|
||
const taskErrorCountMap = new Map();
|
||
|
||
// 准备批量操作
|
||
const multi = redis.multi();
|
||
|
||
for (const task of SuccessTasks) {
|
||
const taskId = task.taskid || task.taskId;
|
||
const remoteTaskId = task.remoteTaskId;
|
||
const aigc = task.aigc;
|
||
const platform = task.platform;
|
||
|
||
console.log(`[pollingTask] 处理任务: taskId=${taskId}, remoteTaskId=${remoteTaskId}, status=${task.status}`);
|
||
|
||
// 存储结果到 task 的 resultData 里
|
||
multi.hSet(`${initQueue.prefix}:task:${taskId}`, 'resultData', task.result);
|
||
multi.hSet(`${initQueue.prefix}:task:${taskId}`, 'status', task.status);
|
||
|
||
// 判断是否为错误任务
|
||
if (task.status === 'failed') {
|
||
// 推送任务 ID 到错误列表
|
||
multi.lPush(initQueue.errorList, taskId);
|
||
|
||
// 计算错误队列任务数
|
||
const key = `${aigc}:${platform}`;
|
||
if(taskErrorCountMap.has(key)){
|
||
taskErrorCountMap.set(key, taskErrorCountMap.get(key) + 1);
|
||
} else {
|
||
taskErrorCountMap.set(key, 1);
|
||
}
|
||
} else {
|
||
// 推送任务 ID 到结果列表
|
||
multi.lPush(initQueue.resultList, taskId);
|
||
|
||
// 计算各平台处理队列的已完成待释放任务数
|
||
const key = `${aigc}:${platform}`;
|
||
if(taskCountMap.has(key)){
|
||
taskCountMap.set(key, taskCountMap.get(key) + 1);
|
||
} else {
|
||
taskCountMap.set(key, 1);
|
||
}
|
||
}
|
||
|
||
// 按平台+AIGC类型删除轮询任务
|
||
const platformKey = `${initQueue.prefix}:processPolling:${aigc}:${platform}`;
|
||
multi.hDel(platformKey, remoteTaskId);
|
||
}
|
||
|
||
// 执行所有Redis操作
|
||
await multi.exec();
|
||
console.log(`[pollingTask] 已完成Redis批量操作，删除了轮询任务并存储到结果队列`);
|
||
|
||
// 更新平台计数（使用原子操作）
|
||
if (taskCountMap.size > 0) {
|
||
await initQueue.reducePlatformsProcess(taskCountMap);
|
||
console.log(`[pollingTask] 已更新平台计数: ${JSON.stringify(Array.from(taskCountMap.entries()))}`);
|
||
}
|
||
|
||
// 更新错误队列计数
|
||
if (taskErrorCountMap.size > 0) {
|
||
const totalErrorCount = Object.values(taskErrorCountMap).reduce((a, b) => a + b, 0);
|
||
initQueue.addEQtaskALL(totalErrorCount);
|
||
console.log(`[pollingTask] 已更新错误队列计数: ${JSON.stringify(Array.from(taskErrorCountMap.entries()))}`);
|
||
}
|
||
|
||
return true
|
||
}
|
||
|
||
parentPort.on('message', async (tasks) => {
|
||
console.log(`[pollingTask] 收到主线程消息，任务数量: ${Object.keys(tasks).length}`);
|
||
try {
|
||
const responseTasks = await getTask(tasks);
|
||
// 过滤掉 undefined 元素
|
||
const successTasks = responseTasks.filter(task => task !== undefined);
|
||
console.log(`[pollingTask] 过滤后的成功任务数量: ${successTasks.length}`);
|
||
if (successTasks.length > 0) {
|
||
await storeSuccessTasks(successTasks);
|
||
}
|
||
parentPort.postMessage({ status: 'completed', processed: successTasks.length });
|
||
console.log(`[pollingTask] 任务处理完成，已通知主线程`);
|
||
} catch (error) {
|
||
console.error('[pollingTask] 处理轮询任务时出错:', error);
|
||
parentPort.postMessage({ status: 'error', error: error.message });
|
||
}
|
||
});
|