shuzhiren-comfyui/任务队列后端/worker_threads/callback_result/result.js

import { parentPort } from 'worker_threads'
import redis from '../../redis/index.js'
import initQueue from '../../redis/initQueue.js'

const logger = {
  info: (message) => {
    const timestamp = new Date().toISOString();
    console.log(`[${timestamp}] [CallbackResult] INFO: ${message}`);
  },
  warn: (message) => {
    const timestamp = new Date().toISOString();
    console.warn(`[${timestamp}] [CallbackResult] WARN: ${message}`);
  },
  error: (message, error) => {
    const timestamp = new Date().toISOString();
    console.error(`[${timestamp}] [CallbackResult] ERROR: ${message}`, error || '');
  },
  debug: (message) => {
    const timestamp = new Date().toISOString();
    console.debug(`[${timestamp}] [CallbackResult] DEBUG: ${message}`);
  }
};

const MAX_RETRIES = 3;
const STUCK_TASK_THRESHOLD = 300000;

async function getTasks() {
  try {
    const taskIds = await redis.lRange(initQueue.callback, 0, 49);
    const taskCountMap = new Map();
    const orphanTaskIds = [];
    const processedTaskIds = [];
    const seenTaskIds = new Set();

    if (taskIds.length === 0) {
      await redis.json.set(initQueue.initInfoKey, '$.CQtasksALL', 0);
      logger.debug('回调结果队列为空，已重置任务数为0');
      return [];
    }
    
    // logger.debug(`回调结果队列任务ID: ${taskIds.slice(0, 5).join(', ')}${taskIds.length > 5 ? '...' : ''}`);
    
    const multi = redis.multi();
    for (const taskId of taskIds) {
      multi.hGetAll(`${initQueue.prefix}:task:${taskId}`);
    }
    
    const results = await multi.exec();

    for (let i = 0; i < taskIds.length; i++) {
      const taskId = taskIds[i];

      if (seenTaskIds.has(taskId)) {
        continue;
      }
      seenTaskIds.add(taskId);

      const taskInfo = results[i];
      
      if (!taskInfo || Object.keys(taskInfo).length === 0) {
        logger.warn(`任务数据不存在或为空: ${taskId}，标记为孤立任务`);
        orphanTaskIds.push(taskId);
        continue;
      }
      
      const backendId = taskInfo.backendId;
      const resultData = taskInfo.resultData;
      const aigc = taskInfo.AIGC || 'default';
      const platform = taskInfo.platform || 'default';
      
      if (!backendId) {
        logger.warn(`任务缺少backendId: ${taskId}，标记为孤立任务`);
        orphanTaskIds.push(taskId);
        continue;
      }
      
      try {
        const resultWithTaskId = {
          taskId: taskId,
          result: resultData
        };

        parentPort.postMessage({
          type: 'success',
          backendId: backendId,
          message: JSON.stringify(resultWithTaskId)
        });

        logger.debug(`成功发送结果给客户端，taskId: ${taskId}`)

        const key = `${aigc}:${platform}`;
        if(taskCountMap.has(key)){
          taskCountMap.set(key, taskCountMap.get(key) + 1);
        } else {
          taskCountMap.set(key, 1);
        }
        processedTaskIds.push(taskId);
      } catch (sendError) {
        logger.error(`发送结果给客户端失败: ${taskId}`, sendError);
        processedTaskIds.push(taskId);
      }
    }
    
    if (orphanTaskIds.length > 0) {
      logger.warn(`发现${orphanTaskIds.length}个孤立任务，将从队列中移除`);
      const multiOrphan = redis.multi();
      for (const orphanId of orphanTaskIds) {
        multiOrphan.lRem(initQueue.callback, 0, orphanId);
      }
      await multiOrphan.exec();
    }
    
    if (processedTaskIds.length > 0 || orphanTaskIds.length > 0) {
      const multiFinal = redis.multi();
      
      if (processedTaskIds.length > 0) {
        multiFinal.lTrim(initQueue.callback, processedTaskIds.length, -1);
      }
      
      await multiFinal.exec();
      
      if (taskCountMap.size > 0) {
        await initQueue.reducePlatformsProcess(taskCountMap);
      }
      
      const totalProcessed = processedTaskIds.length + orphanTaskIds.length;
      await initQueue.reduceCQtasksALL(totalProcessed);
      
      logger.info(`已处理${processedTaskIds.length}个回调结果任务，移除${orphanTaskIds.length}个孤立任务`);
    }
    
    return taskIds; 
  } catch (error) {
    logger.error('处理回调结果任务失败:', error);
    await redis.json.set(initQueue.initInfoKey, '$.CQtasksALL', 0);
    return [];
  }
}

async function syncQueueState() {
  try {
    const actualLength = await redis.lLen(initQueue.callback);
    const storedCount = await initQueue.getCQtasksALL();
    
    if (actualLength !== storedCount) {
      logger.warn(`检测到队列状态不一致: 实际长度=${actualLength}, 存储计数=${storedCount}`);
      await redis.json.set(initQueue.initInfoKey, '$.CQtasksALL', actualLength);
      logger.info(`已同步队列状态: CQtasksALL=${actualLength}`);
    }
  } catch (error) {
    logger.error('同步队列状态失败:', error);
  }
}

(async () => {
  logger.info('回调结果处理线程启动');
  
  let checkCount = 0;
  
  while (true) {
    try {
      checkCount++;
      
      const rqTasksAll = await initQueue.getCQtasksALL();

      if (rqTasksAll !== 0) {
        logger.info(`回调结果队列有任务可处理，数量: ${rqTasksAll}`);
        await getTasks();
        
        await new Promise(resolve => setTimeout(resolve, 1000));
      } else { 
        if (checkCount % 10 === 0) {
          await syncQueueState();
        }
        await new Promise(resolve => setTimeout(resolve, 10000));
      }
      
    } catch (error) {
      logger.error('持续处理回调结果任务失败:', error);
      await new Promise(resolve => setTimeout(resolve, 5000));
    }
  }
})()