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

import dotenv from 'dotenv';
import { createClient } from 'redis';
import initQueue from './redis/initQueue.js';

dotenv.config();

const prefix = process.env.PROJECT_PREFIX || 'default';
const initInfoKey = `${prefix}:InitInfo`;

const redis = createClient({
  url: process.env.REDIS_URL || 'redis://localhost:6379'
});

async function resetCounters() {
  try {
    console.log('正在重置计数器...');
    
    const initInfoResult = await redis.json.get(initInfoKey, { path: '$' });
    
    if (!initInfoResult || !initInfoResult[0]) {
      console.log('未找到初始化信息，跳过计数器重置');
      return;
    }
    
    const initInfo = initInfoResult[0];
    
    console.log('\n当前计数器状态:');
    console.log(`  PQtasksALL (总处理任务数): ${initInfo.PQtasksALL || 0}`);
    console.log(`  RQtasksALL (总结果任务数): ${initInfo.RQtasksALL || 0}`);
    console.log(`  CQtasksALL (总回调任务数): ${initInfo.CQtasksALL || 0}`);
    console.log(`  EQtaskALL (总错误任务数): ${initInfo.EQtaskALL || 0}`);
    
    if (initInfo.platforms) {
      console.log('\n各平台计数器状态:');
      for (const [key, platform] of Object.entries(initInfo.platforms)) {
        console.log(`  ${key}: WQtasks=${platform.WQtasks || 0}, PQtasks=${platform.PQtasks || 0}`);
      }
    }
    
    const multi = redis.multi();
    
    multi.json.set(initInfoKey, '$.PQtasksALL', 0);
    multi.json.set(initInfoKey, '$.RQtasksALL', 0);
    multi.json.set(initInfoKey, '$.CQtasksALL', 0);
    multi.json.set(initInfoKey, '$.EQtaskALL', 0);
    
    if (initInfo.platforms) {
      for (const key of Object.keys(initInfo.platforms)) {
        multi.json.set(initInfoKey, `$.platforms.${key}.WQtasks`, 0);
        multi.json.set(initInfoKey, `$.platforms.${key}.PQtasks`, 0);
      }
    }
    
    await multi.exec();
    
    console.log('\n✓ 所有计数器已重置为0');
    
  } catch (error) {
    console.error('重置计数器失败:', error.message);
    throw error;
  }
}

async function clearCallbackPendingQueue() {
  try {
    const callbackPendingKey = `${prefix}:callback:pending`;
    const exists = await redis.exists(callbackPendingKey);
    
    if (exists) {
      const count = await redis.hLen(callbackPendingKey);
      if (count > 0) {
        await redis.del(callbackPendingKey);
        console.log(`\n✓ 已清除回调等待队列，共 ${count} 个任务`);
      }
    }
  } catch (error) {
    console.error('清除回调等待队列失败:', error.message);
  }
}

async function clearAllProjectData() {
  try {
    console.log('正在连接Redis...');
    console.log('REDIS_URL:', process.env.REDIS_URL);
    
    await redis.connect();
    console.log('Redis连接成功');
    
    console.log(`\n开始清除项目 "${prefix}" 的所有Redis数据...`);
    
    try {
      const initInfoResult = await redis.json.get(initInfoKey, { path: '$' });
      if (initInfoResult) {
        console.log('\n当前初始化信息:', JSON.stringify(initInfoResult, null, 2));
      } else {
        console.log('\n未找到初始化信息');
      }
    } catch (error) {
      console.log('获取初始化信息失败（可能不存在）:', error.message);
    }
    
    const keysToDelete = [];
    
    const patterns = [
      `${prefix}:*`,
      `${initQueue?.prefix || prefix}:*`,
    ];
    
    for (const pattern of patterns) {
      console.log(`\n正在搜索模式: ${pattern}`);
      let cursor = '0';
      let totalDeleted = 0;
      
      do {
        try {
          const result = await redis.scan(cursor, {
            MATCH: pattern,
            COUNT: 100
          });
          
          const newCursor = result.cursor;
          const keys = result.keys || [];
          
          console.log(`  游标: ${newCursor}, 找到键数量: ${keys.length}`);
          
          if (keys.length > 0) {
            const validKeys = keys.filter(key => {
              return typeof key === 'string' && key.trim() !== '';
            });
            
            if (validKeys.length > 0) {
              await redis.del(...validKeys);
              totalDeleted += validKeys.length;
              console.log(`  已删除 ${validKeys.length} 个键`);
            }
          }
          
          cursor = newCursor;
        } catch (error) {
          console.error(`  搜索过程中出错:`, error.message);
          break;
        }
      } while (cursor !== '0');
      
      console.log(`模式 "${pattern}" 共删除 ${totalDeleted} 个键`);
    }
    
    try {
      const platforms = await redis.json.get(initInfoKey, { path: '$.platforms' });
      if (platforms && platforms[0]) {
        for (const [key, platform] of Object.entries(platforms[0])) {
          if (platform.waitQueue) {
            console.log(`\n删除等待队列: ${platform.waitQueue}`);
            await redis.del(platform.waitQueue);
          }
        }
      }
    } catch (error) {
      console.log('删除等待队列失败:', error.message);
    }
    
    console.log('\n========================================');
    console.log('项目所有Redis数据已清除完成！');
    console.log('========================================');
    
  } catch (error) {
    console.error('清除Redis数据失败:', error);
  } finally {
    if (redis.isOpen) {
      await redis.disconnect();
      console.log('\nRedis连接已关闭');
    }
    process.exit();
  }
}

async function resetQueueCounters() {
  try {
    console.log('正在连接Redis...');
    console.log('REDIS_URL:', process.env.REDIS_URL);
    
    await redis.connect();
    console.log('Redis连接成功');
    
    console.log(`\n========== 重置队列计数器 ==========`);
    console.log(`项目前缀: ${prefix}`);
    
    await resetCounters();
    await clearCallbackPendingQueue();
    
    console.log('\n========================================');
    console.log('队列计数器重置完成！');
    console.log('现在可以正常处理新任务了');
    console.log('========================================');
    
  } catch (error) {
    console.error('重置队列计数器失败:', error);
  } finally {
    if (redis.isOpen) {
      await redis.disconnect();
      console.log('\nRedis连接已关闭');
    }
    process.exit();
  }
}

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

if (mode === 'clear') {
  clearAllProjectData();
} else {
  resetQueueCounters();
}