jk-rag-python/rag/upload_queue_manager.py

import asyncio
from queue import Queue
from typing import Dict, Callable, Any
from utils.get_logger import setup_logger
import time

logger = setup_logger(__name__)


class UploadQueueManager:
    """
    文档上传队列管理器
    
    功能：
    1. 控制最多10个并发任务
    2. 前一个任务完成后，下一个任务自动开始
    3. 提供队列状态监控
    """
    
    def __init__(self, max_concurrent: int = 10):
        """
        初始化队列管理器
        
        参数:
            max_concurrent: 最大并发数，默认10
        """
        self.max_concurrent = max_concurrent
        self.task_queue = asyncio.Queue()  # 任务队列
        self.semaphore = asyncio.Semaphore(max_concurrent)  # 信号量控制并发
        self.active_tasks: Dict[str, Dict[str, Any]] = {}  # 活跃任务字典 {task_id: task_info}
        self.completed_tasks: Dict[str, Dict[str, Any]] = {}  # 已完成任务字典
        self.failed_tasks: Dict[str, Dict[str, Any]] = {}  # 失败任务字典
        self.workers = []  # 工作协程列表
        self.is_running = False
        self._lock = asyncio.Lock()  # 用于保护共享状态
        
        logger.info(f"队列管理器初始化完成，最大并发数: {max_concurrent}")
    
    async def start(self):
        """启动队列管理器的工作协程"""
        if self.is_running:
            logger.warning("队列管理器已经在运行中")
            return
        
        self.is_running = True
        # 启动多个工作协程
        for i in range(self.max_concurrent):
            worker = asyncio.create_task(self._worker(i))
            self.workers.append(worker)
        
        logger.info(f"队列管理器已启动，工作协程数: {self.max_concurrent}")
    
    async def stop(self):
        """停止队列管理器"""
        self.is_running = False
        
        # 等待所有工作协程完成
        if self.workers:
            await asyncio.gather(*self.workers, return_exceptions=True)
            self.workers.clear()
        
        logger.info("队列管理器已停止")
    
    async def _worker(self, worker_id: int):
        """
        工作协程，从队列中取任务并执行
        
        参数:
            worker_id: 工作协程ID
        """
        logger.info(f"工作协程 {worker_id} 已启动")
        
        while self.is_running:
            try:
                # 从队列中获取任务（带超时，避免永久阻塞）
                try:
                    task_info = await asyncio.wait_for(
                        self.task_queue.get(),
                        timeout=1.0
                    )
                except asyncio.TimeoutError:
                    continue
                
                task_id = task_info["task_id"]
                task_func = task_info["task_func"]
                task_args = task_info["task_args"]
                task_kwargs = task_info["task_kwargs"]
                
                logger.info(f"工作协程 {worker_id} 开始处理任务: {task_id}")
                
                # 使用信号量控制并发
                async with self.semaphore:
                    # 记录任务开始
                    async with self._lock:
                        self.active_tasks[task_id] = {
                            "task_id": task_id,
                            "worker_id": worker_id,
                            "start_time": time.time(),
                            "status": "running"
                        }
                    
                    try:
                        # 执行任务
                        result = await task_func(*task_args, **task_kwargs)
                        
                        # 记录任务完成
                        async with self._lock:
                            if task_id in self.active_tasks:
                                task_info = self.active_tasks.pop(task_id)
                                task_info["end_time"] = time.time()
                                task_info["duration"] = task_info["end_time"] - task_info["start_time"]
                                task_info["status"] = "completed"
                                task_info["result"] = result
                                self.completed_tasks[task_id] = task_info
                        
                        logger.info(f"工作协程 {worker_id} 完成任务: {task_id}, 耗时: {task_info['duration']:.2f}秒")
                    
                    except Exception as e:
                        # 记录任务失败
                        logger.error(f"工作协程 {worker_id} 执行任务 {task_id} 失败: {e}", exc_info=True)
                        
                        async with self._lock:
                            if task_id in self.active_tasks:
                                task_info = self.active_tasks.pop(task_id)
                                task_info["end_time"] = time.time()
                                task_info["duration"] = task_info["end_time"] - task_info["start_time"]
                                task_info["status"] = "failed"
                                task_info["error"] = str(e)
                                self.failed_tasks[task_id] = task_info
                
                # 标记任务完成
                self.task_queue.task_done()
            
            except Exception as e:
                logger.error(f"工作协程 {worker_id} 发生异常: {e}", exc_info=True)
        
        logger.info(f"工作协程 {worker_id} 已停止")
    
    async def submit_task(self, task_id: str, task_func: Callable, *args, **kwargs) -> None:
        """
        提交任务到队列
        
        参数:
            task_id: 任务ID（通常是 document_id）
            task_func: 任务函数（异步函数）
            *args: 任务函数的位置参数
            **kwargs: 任务函数的关键字参数
        """
        task_info = {
            "task_id": task_id,
            "task_func": task_func,
            "task_args": args,
            "task_kwargs": kwargs,
            "submit_time": time.time()
        }
        
        await self.task_queue.put(task_info)
        
        queue_size = self.task_queue.qsize()
        active_count = len(self.active_tasks)
        
        logger.info(f"任务已提交到队列: {task_id}, 队列长度: {queue_size}, 活跃任务数: {active_count}")
    
    def get_queue_status(self) -> Dict[str, Any]:
        """
        获取队列状态
        
        返回:
            包含队列状态信息的字典
        """
        return {
            "is_running": self.is_running,
            "max_concurrent": self.max_concurrent,
            "queue_size": self.task_queue.qsize(),
            "active_tasks_count": len(self.active_tasks),
            "completed_tasks_count": len(self.completed_tasks),
            "failed_tasks_count": len(self.failed_tasks),
            "active_tasks": list(self.active_tasks.keys()),
            "worker_count": len(self.workers)
        }
    
    def get_task_status(self, task_id: str) -> Dict[str, Any]:
        """
        获取指定任务的状态
        
        参数:
            task_id: 任务ID
        
        返回:
            任务状态信息，如果任务不存在返回 None
        """
        if task_id in self.active_tasks:
            return {**self.active_tasks[task_id], "status": "running"}
        elif task_id in self.completed_tasks:
            return {**self.completed_tasks[task_id], "status": "completed"}
        elif task_id in self.failed_tasks:
            return {**self.failed_tasks[task_id], "status": "failed"}
        else:
            return {"status": "not_found"}
    
    async def wait_for_completion(self):
        """等待队列中的所有任务完成"""
        await self.task_queue.join()
        logger.info("队列中的所有任务已完成")
    
    async def remove_from_queue(self, task_id: str) -> bool:
        """
        从等待队列中删除指定任务
        
        参数:
            task_id: 任务ID
        
        返回:
            bool: 是否成功删除
        """
        async with self._lock:
            # 获取队列中的所有任务
            temp_tasks = []
            found = False
            
            # 从队列中取出所有任务
            while not self.task_queue.empty():
                try:
                    task_info = self.task_queue.get_nowait()
                    if task_info["task_id"] == task_id:
                        found = True
                        logger.info(f"从队列中找到并删除任务: {task_id}")
                        self.task_queue.task_done()
                    else:
                        temp_tasks.append(task_info)
                except:
                    break
            
            # 将其他任务放回队列
            for task_info in temp_tasks:
                await self.task_queue.put(task_info)
            
            return found
    
    def clear_history(self):
        """清空已完成和失败任务的历史记录"""
        cleared_completed = len(self.completed_tasks)
        cleared_failed = len(self.failed_tasks)
        
        self.completed_tasks.clear()
        self.failed_tasks.clear()
        
        logger.info(f"已清空历史记录: 完成任务 {cleared_completed} 个, 失败任务 {cleared_failed} 个")


# 全局队列管理器实例（单例模式）
_global_queue_manager: UploadQueueManager = None


def get_queue_manager(max_concurrent: int = 10) -> UploadQueueManager:
    """
    获取全局队列管理器实例（单例模式）
    
    参数:
        max_concurrent: 最大并发数，默认10
    
    返回:
        UploadQueueManager 实例
    """
    global _global_queue_manager
    
    if _global_queue_manager is None:
        _global_queue_manager = UploadQueueManager(max_concurrent)
        logger.info("创建全局队列管理器实例")
    
    return _global_queue_manager