183a57e7f1
- 将全部路由文件（12个）的 jsonify 响应迁移至 success_response/error_response 统一格式 - 修复 sync_routes.py error_response 参数错误（P0） - 新增异步任务系统：task_registry + task_routes - 新增状态码：TASK_NOT_FOUND(4014)、TASK_CONFLICT(4015)、REINDEX_ERROR(5040) - 修正 task_routes/exam_pkg 中语义不匹配的状态码 - 更新 curl 测试手册、后端对接规范文档 - 添加缓存性能报告和 Redis 迁移计划
312 lines
11 KiB
Python
312 lines
11 KiB
Python
"""
|
|
异步任务注册表
|
|
|
|
为长时间运行的操作（同步、重建索引、上传向量化、出题、批阅）提供统一的
|
|
任务状态跟踪机制。支持 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
|