shuzhiren-comfyui/任务队列后端/utils/queueRecovery.js

import redis from '../redis/index.js';
import initQueue from '../redis/initQueue.js';

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

async function recoverCallbackQueue() {
  logger.info('开始恢复回调队列...');
  
  try {
    const actualLength = await redis.lLen(initQueue.callback);
    const storedCount = await initQueue.getCQtasksALL();
    
    logger.info(`回调队列状态: 实际长度=${actualLength}, 存储计数=${storedCount}`);
    
    if (actualLength !== storedCount) {
      logger.warn(`检测到不一致，正在修复...`);
      await redis.json.set(initQueue.initInfoKey, '$.CQtasksALL', actualLength);
      logger.info(`已修复回调队列计数: CQtasksALL=${actualLength}`);
    }
    
    const taskIds = await redis.lRange(initQueue.callback, 0, -1);
    const orphanTaskIds = [];
    
    for (const taskId of taskIds) {
      const taskExists = await redis.exists(`${initQueue.prefix}:task:${taskId}`);
      if (!taskExists) {
        orphanTaskIds.push(taskId);
      }
    }
    
    if (orphanTaskIds.length > 0) {
      logger.warn(`发现${orphanTaskIds.length}个孤立任务，正在移除...`);
      const multi = redis.multi();
      for (const orphanId of orphanTaskIds) {
        multi.lRem(initQueue.callback, 0, orphanId);
      }
      await multi.exec();
      logger.info(`已移除${orphanTaskIds.length}个孤立任务`);
    }
    
    logger.info('回调队列恢复完成');
    return { actualLength, storedCount, orphanCount: orphanTaskIds.length };
  } catch (error) {
    logger.error('恢复回调队列失败:', error);
    throw error;
  }
}

async function recoverCallbackPendingQueue() {
  logger.info('开始恢复回调等待队列...');
  
  try {
    const pendingTasks = await initQueue.getCallbackPendingTasks();
    const now = Date.now();
    const staleRemoteTaskIds = [];
    
    for (const [remoteTaskId, taskInfoStr] of Object.entries(pendingTasks)) {
      try {
        const taskInfo = JSON.parse(taskInfoStr);
        
        const mappingExists = await redis.exists(`${initQueue.callback}:${remoteTaskId}`);
        const taskExists = await redis.exists(`${initQueue.prefix}:task:${taskInfo.taskId}`);
        
        if (!mappingExists || !taskExists) {
          staleRemoteTaskIds.push(remoteTaskId);
          logger.warn(`发现孤立的回调等待任务: remoteTaskId=${remoteTaskId}, mappingExists=${mappingExists}, taskExists=${taskExists}`);
        }
      } catch (parseError) {
        logger.error(`解析任务信息失败: ${remoteTaskId}`, parseError);
        staleRemoteTaskIds.push(remoteTaskId);
      }
    }
    
    if (staleRemoteTaskIds.length > 0) {
      logger.warn(`发现${staleRemoteTaskIds.length}个孤立的回调等待任务，正在移除...`);
      for (const remoteTaskId of staleRemoteTaskIds) {
        await initQueue.removeCallbackPendingTask(remoteTaskId);
      }
      logger.info(`已移除${staleRemoteTaskIds.length}个孤立的回调等待任务`);
    }
    
    logger.info('回调等待队列恢复完成');
    return { totalCount: Object.keys(pendingTasks).length, staleCount: staleRemoteTaskIds.length };
  } catch (error) {
    logger.error('恢复回调等待队列失败:', error);
    throw error;
  }
}

async function recoverPlatformCounts() {
  logger.info('开始恢复平台计数...');
  
  try {
    const platforms = await initQueue.getPlatforms();
    
    for (const [platformKey, platformInfo] of Object.entries(platforms)) {
      const waitQueueLength = await redis.lLen(platformInfo.waitQueue);
      
      const pollingKeys = await redis.keys(`${initQueue.prefix}:processPolling:${platformInfo.AIGC}:${platformInfo.platformName}`);
      let pollingCount = 0;
      for (const key of pollingKeys) {
        pollingCount += await redis.hLen(key);
      }
      
      const pendingCount = await initQueue.getCallbackPendingCount();
      
      logger.info(`平台 ${platformKey}: 等待队列=${waitQueueLength}, 轮询队列=${pollingCount}`);
    }
    
    logger.info('平台计数恢复完成');
  } catch (error) {
    logger.error('恢复平台计数失败:', error);
    throw error;
  }
}

async function cleanupExpiredMappings() {
  logger.info('开始清理过期的回调映射...');
  
  try {
    const keys = await redis.keys(`${initQueue.callback}:*`);
    let cleanedCount = 0;
    
    for (const key of keys) {
      const remoteTaskId = key.replace(`${initQueue.callback}:`, '');
      const taskId = await redis.get(key);
      
      if (taskId) {
        const taskExists = await redis.exists(`${initQueue.prefix}:task:${taskId}`);
        
        if (!taskExists) {
          await redis.del(key);
          await initQueue.removeCallbackPendingTask(remoteTaskId);
          cleanedCount++;
          logger.debug(`清理孤立映射: ${key} -> ${taskId}`);
        }
      }
    }
    
    logger.info(`清理完成，共清理${cleanedCount}个过期映射`);
    return cleanedCount;
  } catch (error) {
    logger.error('清理过期映射失败:', error);
    throw error;
  }
}

async function fullRecovery() {
  logger.info('========== 开始完整队列恢复 ==========');
  
  const results = {
    callbackQueue: null,
    callbackPending: null,
    platformCounts: null,
    expiredMappings: null
  };
  
  try {
    results.callbackQueue = await recoverCallbackQueue();
    results.callbackPending = await recoverCallbackPendingQueue();
    results.platformCounts = await recoverPlatformCounts();
    results.expiredMappings = await cleanupExpiredMappings();
    
    logger.info('========== 完整队列恢复完成 ==========');
    logger.info('恢复结果:', JSON.stringify(results, null, 2));
    
    return results;
  } catch (error) {
    logger.error('完整队列恢复失败:', error);
    throw error;
  }
}

const args = process.argv.slice(2);
const command = args[0] || 'full';

(async () => {
  try {
    if (!redis.isOpen) {
      await redis.connect();
      logger.info('Redis 连接成功');
    }
    
    switch (command) {
      case 'callback':
        await recoverCallbackQueue();
        break;
      case 'pending':
        await recoverCallbackPendingQueue();
        break;
      case 'platforms':
        await recoverPlatformCounts();
        break;
      case 'cleanup':
        await cleanupExpiredMappings();
        break;
      case 'full':
      default:
        await fullRecovery();
        break;
    }
    
    process.exit(0);
  } catch (error) {
    logger.error('恢复操作失败:', error);
    process.exit(1);
  }
})();