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

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

// 日志工具函数
const logger = {
  info: (message) => {
    const timestamp = new Date().toISOString();
    console.log(`[${timestamp}] INFO: ${message}`);
  },
  error: (message, error) => {
    const timestamp = new Date().toISOString();
    console.error(`[${timestamp}] ERROR: ${message}`, error || '');
  },
  debug: (message) => {
    const timestamp = new Date().toISOString();
    console.debug(`[${timestamp}] DEBUG: ${message}`);
  }
};

// 创建线程池实例
const threadPool = QueryThreadPool; // 6个线程的线程池

// 批量获取处理任务的信息 
async function getTasks(tasks) {
  // 每100个任务为一组
  const batchSize = 100;
  const batches = [];
  
  // 将任务数据按批次分割
  const taskEntries = Object.entries(tasks);
  for (let i = 0; i < taskEntries.length; i += batchSize) {
    const batchEntries = taskEntries.slice(i, i + batchSize);
    batches.push(Object.fromEntries(batchEntries));
  }
  
  try {
    // 等待所有子线程完成处理
    const workerResults = await threadPool.executeAllTasks(batches);
    return workerResults.flat();
  } catch (error) {
    logger.error('子线程处理任务时出错:', error);
    throw error;
  }
}

// 获取系统负载（简化版，实际可根据CPU/内存使用率调整）
async function getSystemLoad() {
  return 0.5; // 模拟系统负载，0-1之间
}

// 动态计算轮询间隔
function getDynamicInterval(taskCount) {
  if (taskCount > 100) return 5000; // 任务多，缩短间隔
  if (taskCount > 0) return 10000;  // 有任务，正常间隔
  return 30000;                     // 无任务，延长间隔
}

// 获取支持的平台和AIGC类型组合
async function getSupportedPlatforms() {
  // 从redis获取所有平台AIGC组合
  const platforms = new Set();
  try {
    // 直接从轮询队列中获取平台信息，这样更准确
    const pollingKeys = await redis.keys(`${initQueue.prefix}:processPolling:*`);
    logger.debug(`[getSupportedPlatforms] 找到轮询队列键: ${JSON.stringify(pollingKeys)}`);
    for (const key of pollingKeys) {
      // 从轮询键中提取AIGC和platform
      const match = key.match(/processPolling:(.+?):(.+)$/);
      if (match) {
        const [, aigc, platform] = match;
        platforms.add(`${aigc}:${platform}`);
        logger.debug(`[getSupportedPlatforms] 添加平台组合: ${aigc}:${platform}`);
      }
    }
    
    // 如果轮询队列中没有任务，从所有任务中获取平台信息作为备用
    if (platforms.size === 0) {
      logger.debug('[getSupportedPlatforms] 轮询队列为空，从任务数据中获取平台信息');
      const taskKeys = await redis.keys(`${initQueue.prefix}:task:*`);
      logger.debug(`[getSupportedPlatforms] 找到任务键: ${taskKeys.length} 个`);
      for (const key of taskKeys) {
        // 获取任务信息
        const infoStr = await redis.hGet(key, 'info');
        if (infoStr) {
          try {
            const info = JSON.parse(infoStr);
            platforms.add(`${info.AIGC}:${info.platform}`);
            logger.debug(`[getSupportedPlatforms] 从任务添加平台组合: ${info.AIGC}:${info.platform}`);
          } catch (parseError) {
            logger.error('解析任务信息失败:', parseError);
          }
        }
      }
    }
    
    logger.info(`[getSupportedPlatforms] 最终支持的平台组合: ${JSON.stringify(Array.from(platforms))}`);
  } catch (error) {
    logger.error('获取支持的平台组合失败:', error);
  }
  return Array.from(platforms);
}

// 轮询单个平台组合
async function pollPlatform(platformKey) {
  logger.info(`开始轮询平台组合: ${platformKey}`);
  
  while(true) {
    try {
      // 获取该平台的轮询任务
      const [aigc, platform] = platformKey.split(':');
      const pollingKey = `${initQueue.prefix}:processPolling:${aigc}:${platform}`;
      
      // 检查是否有任务
      const taskCount = await redis.hLen(pollingKey);
      logger.debug(`[pollPlatform] 检查轮询队列: ${pollingKey}, 任务数量: ${taskCount}`);
      
      if(taskCount > 0) {
        logger.info(`平台 ${platformKey} 有可处理任务，数量: ${taskCount}`);
        
        // 动态计算批量大小
        const systemLoad = await getSystemLoad();
        const batchSize = Math.max(50, Math.min(200, 100 - systemLoad * 50));
        
        // 获取所有任务
        const tasks = await redis.hGetAll(pollingKey);
        logger.debug(`批量获取平台 ${platformKey} 任务信息: ${JSON.stringify(Object.keys(tasks))}`);
        
        // 处理任务
        await getTasks(tasks);
        logger.info(`平台 ${platformKey} 处理完成`);
      }
      
      // 动态调整轮询间隔
      const interval = getDynamicInterval(taskCount);
      logger.debug(`平台 ${platformKey} 轮询间隔: ${interval}ms`);
      await new Promise(resolve => setTimeout(resolve, interval));
    } catch (error) {
      logger.error(`处理平台 ${platformKey} 任务时出错:`, error);
      // 出错后等待5秒再重试
      await new Promise(resolve => setTimeout(resolve, 5000));
    }
  }
}

// 持续执行批量处理
(async () => {
  logger.info('[process.js] 轮询线程启动');
  // 获取支持的平台组合
  const platforms = await getSupportedPlatforms();
  logger.info(`支持的平台组合: ${platforms.join(', ')}`);
  
  // 为每个平台组合启动独立的轮询线程
  for (const platform of platforms) {
    logger.info(`[process.js] 启动平台轮询线程: ${platform}`);
    pollPlatform(platform);
  }
  
  // 定期检查新增的平台组合
  setInterval(async () => {
    logger.debug('[process.js] 检查新增的平台组合');
    const currentPlatforms = await getSupportedPlatforms();
    for (const platform of currentPlatforms) {
      if (!platforms.includes(platform)) {
        platforms.push(platform);
        logger.info(`发现新平台组合: ${platform}`);
        pollPlatform(platform);
      }
    }
  }, 60000); // 每分钟检查一次
})()