""" 异步任务注册表 为长时间运行的操作（同步、重建索引、上传向量化、出题、批阅）提供统一的 任务状态跟踪机制。支持 JSON 轮询和 SSE 流式两种进度查询方式。 核心设计： - 进程内字典存储，线程安全 - 后台线程执行长操作，立即返回 task_id - 任务完成后保留 1 小时自动清理 - 单 Worker 环境下足够可靠 使用示例： from core.task_registry import get_registry registry = get_registry() task_id = registry.start_task('sync', '文档同步', total=10) # 在后台线程中更新进度 registry.update_progress(task_id, current=5, message='处理中...') registry.complete_task(task_id, result={...}) """ import uuid import time import threading import logging from dataclasses import dataclass, field from typing import Any, Callable, Dict, List, Optional from concurrent.futures import ThreadPoolExecutor from datetime import datetime logger = logging.getLogger(__name__) # ==================== 数据模型 ==================== @dataclass class TaskInfo: """任务信息""" id: str # 任务唯一标识 type: str # 任务类型: sync / reindex / upload / batch_upload / exam_generate / exam_grade description: str # 人类可读描述 status: str = 'pending' # pending / running / completed / failed progress: float = 0.0 # 0-100 百分比 current: int = 0 # 当前处理项数 total: int = 0 # 总项数 stage: str = '' # 当前阶段描述 message: str = '' # 当前步骤消息 result: Any = None # 完成后的结果数据 error: Optional[str] = None # 失败时的错误信息 created_at: float = 0.0 # 创建时间戳 started_at: Optional[float] = None completed_at: Optional[float] = None _events: List[dict] = field(default_factory=list, repr=False) # SSE 事件缓冲 _lock: threading.Lock = field(default_factory=threading.Lock, repr=False) def add_event(self, event: dict): """添加 SSE 事件到缓冲区""" with self._lock: self._events.append(event) def drain_events(self) -> List[dict]: """取出并清空事件缓冲区""" with self._lock: events = list(self._events) self._events.clear() return events def to_dict(self) -> dict: """转为可序列化的字典""" d = { 'task_id': self.id, 'type': self.type, 'description': self.description, 'status': self.status, 'progress': round(self.progress, 1), 'current': self.current, 'total': self.total, 'stage': self.stage, 'message': self.message, 'created_at': datetime.fromtimestamp(self.created_at).isoformat(), } if self.started_at: d['started_at'] = datetime.fromtimestamp(self.started_at).isoformat() if self.completed_at: d['completed_at'] = datetime.fromtimestamp(self.completed_at).isoformat() d['duration_ms'] = int((self.completed_at - self.started_at) * 1000) if self.result is not None: d['result'] = self.result if self.error is not None: d['error'] = self.error return d # ==================== 任务注册表 ==================== class TaskRegistry: """ 全局任务注册表（单例） 管理所有异步任务的创建、执行、状态更新和查询。 """ def __init__(self, max_workers: int = 4, task_ttl: int = 3600): self._tasks: Dict[str, TaskInfo] = {} self._lock = threading.Lock() self._executor = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix='task') self._task_ttl = task_ttl # 已完成任务的保留时间（秒） self._cleanup_interval = 300 # 清理间隔（秒） self._last_cleanup = time.time() def create_task(self, task_type: str, description: str, total: int = 0) -> TaskInfo: """ 创建新任务 Args: task_type: 任务类型标识 description: 人类可读描述 total: 预计处理的总项数（用于计算百分比） Returns: TaskInfo 实例 """ task_id = uuid.uuid4().hex[:12] task = TaskInfo( id=task_id, type=task_type, description=description, status='pending', total=total, created_at=time.time(), ) with self._lock: self._tasks[task_id] = task logger.info(f"[任务] 创建 {task_id}: {description}") return task def start_task(self, task_id: str, fn: Callable, *args, **kwargs) -> str: """ 在后台线程中执行任务 Args: task_id: 任务 ID fn: 要执行的函数，签名为 fn(task, *args, **kwargs) *args, **kwargs: 传给 fn 的参数 Returns: task_id """ task = self._tasks.get(task_id) if not task: raise ValueError(f"任务不存在: {task_id}") self._submit(fn, task, args, kwargs) return task_id def _submit(self, fn, task, args, kwargs): """提交任务到线程池""" def _wrapper(): task.status = 'running' task.started_at = time.time() task.add_event({'type': 'start', 'data': {'stage': task.stage or '初始化'}}) try: result = fn(task, *args, **kwargs) task.status = 'completed' task.progress = 100.0 task.completed_at = time.time() task.result = result task.add_event({'type': 'complete', 'data': task.to_dict()}) logger.info(f"[任务] 完成 {task.id}: {task.description}") except Exception as e: task.status = 'failed' task.error = str(e) task.completed_at = time.time() task.add_event({'type': 'error', 'data': {'message': str(e)}}) logger.error(f"[任务] 失败 {task.id}: {e}", exc_info=True) self._executor.submit(_wrapper) def get_task(self, task_id: str) -> Optional[TaskInfo]: """获取任务信息""" with self._lock: return self._tasks.get(task_id) def list_tasks(self, status: Optional[str] = None, task_type: Optional[str] = None, limit: int = 50) -> List[TaskInfo]: """列出任务（按创建时间倒序）""" with self._lock: tasks = list(self._tasks.values()) if status: tasks = [t for t in tasks if t.status == status] if task_type: tasks = [t for t in tasks if t.type == task_type] tasks.sort(key=lambda t: t.created_at, reverse=True) return tasks[:limit] def update_progress(self, task_id: str, *, current: Optional[int] = None, total: Optional[int] = None, stage: Optional[str] = None, message: Optional[str] = None): """ 更新任务进度（线程安全） Args: task_id: 任务 ID current: 当前处理项数 total: 更新总项数 stage: 当前阶段 message: 当前步骤描述 """ task = self._tasks.get(task_id) if not task: return with task._lock: if current is not None: task.current = current if total is not None: task.total = total if stage is not None: task.stage = stage if message is not None: task.message = message # 计算百分比 if task.total > 0: task.progress = min(99.0, (task.current / task.total) * 100) # 推送进度事件 task.add_event({ 'type': 'progress', 'data': { 'progress': round(task.progress, 1), 'current': task.current, 'total': task.total, 'stage': task.stage, 'message': task.message, } }) def complete_task(self, task_id: str, result: Any = None): """手动标记任务完成""" task = self._tasks.get(task_id) if not task: return task.status = 'completed' task.progress = 100.0 task.completed_at = time.time() if result is not None: task.result = result task.add_event({'type': 'complete', 'data': task.to_dict()}) def fail_task(self, task_id: str, error: str): """手动标记任务失败""" task = self._tasks.get(task_id) if not task: return task.status = 'failed' task.error = error task.completed_at = time.time() task.add_event({'type': 'error', 'data': {'message': error}}) def cleanup(self): """清理过期的已完成任务""" now = time.time() with self._lock: expired = [ tid for tid, t in self._tasks.items() if t.status in ('completed', 'failed') and t.completed_at and (now - t.completed_at) > self._task_ttl ] for tid in expired: del self._tasks[tid] if expired: logger.debug(f"[任务] 清理 {len(expired)} 个过期任务") def maybe_cleanup(self): """定期清理（避免频繁操作）""" if time.time() - self._last_cleanup > self._cleanup_interval: self.cleanup() self._last_cleanup = time.time() def get_stats(self) -> dict: """获取任务统计""" with self._lock: tasks = list(self._tasks.values()) stats = {'total': len(tasks), 'by_status': {}, 'by_type': {}} for t in tasks: stats['by_status'][t.status] = stats['by_status'].get(t.status, 0) + 1 stats['by_type'][t.type] = stats['by_type'].get(t.type, 0) + 1 return stats # ==================== 全局单例 ==================== _registry: Optional[TaskRegistry] = None _registry_lock = threading.Lock() def get_registry() -> TaskRegistry: """获取全局任务注册表单例""" global _registry if _registry is None: with _registry_lock: if _registry is None: _registry = TaskRegistry() return _registry