89 lines
2.1 KiB
Python
89 lines
2.1 KiB
Python
|
"""异步任务处理模块
|
|||
|
|
|
|||
|
|
TODO: 后续替换为 RabbitMQ 消息队列
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import asyncio
|
|||
|
|
import uuid
|
|||
|
|
from datetime import datetime
|
|||
|
|
from typing import Callable, Awaitable
|
|||
|
|
from app.utils.logger import logger
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 任务状态存储（后续替换为 Redis）
|
|||
|
|
_task_store: Dict[str, Dict] = {}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def generate_task_id() -> str:
|
|||
|
|
"""生成任务ID"""
|
|||
|
|
return f"task-{uuid.uuid4().hex[:12]}"
|
|||
|
|
|
|||
|
|
|
|||
|
|
class AsyncTaskManager:
|
|||
|
|
"""异步任务管理器"""
|
|||
|
|
|
|||
|
|
def __init__(self):
|
|||
|
|
self._running_tasks: dict[str, asyncio.Task] = {}
|
|||
|
|
|
|||
|
|
def create_task(
|
|||
|
|
self,
|
|||
|
|
task_id: str,
|
|||
|
|
task_func: Callable[[str], Awaitable[None]],
|
|||
|
|
) -> asyncio.Task:
|
|||
|
|
"""
|
|||
|
|
创建异步任务
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
task_id: 任务ID
|
|||
|
|
task_func: 任务执行函数
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
asyncio.Task
|
|||
|
|
"""
|
|||
|
|
task = asyncio.create_task(self._run_task(task_id, task_func))
|
|||
|
|
self._running_tasks[task_id] = task
|
|||
|
|
return task
|
|||
|
|
|
|||
|
|
async def _run_task(
|
|||
|
|
self,
|
|||
|
|
task_id: str,
|
|||
|
|
task_func: Callable[[str], Awaitable[None]],
|
|||
|
|
):
|
|||
|
|
"""运行任务并处理状态"""
|
|||
|
|
try:
|
|||
|
|
await task_func(task_id)
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"Task {task_id} failed: {e}")
|
|||
|
|
_task_store[task_id] = {
|
|||
|
|
"status": "failed",
|
|||
|
|
"error": str(e),
|
|||
|
|
"completed_at": datetime.now(),
|
|||
|
|
}
|
|||
|
|
finally:
|
|||
|
|
if task_id in self._running_tasks:
|
|||
|
|
del self._running_tasks[task_id]
|
|||
|
|
|
|||
|
|
def get_task_status(self, task_id: str) -> dict | None:
|
|||
|
|
"""获取任务状态"""
|
|||
|
|
return _task_store.get(task_id)
|
|||
|
|
|
|||
|
|
def cancel_task(self, task_id: str) -> bool:
|
|||
|
|
"""取消任务"""
|
|||
|
|
if task_id in self._running_tasks:
|
|||
|
|
self._running_tasks[task_id].cancel()
|
|||
|
|
return True
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 单例
|
|||
|
|
task_manager = AsyncTaskManager()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_task_status(task_id: str) -> dict | None:
|
|||
|
|
"""获取任务状态"""
|
|||
|
|
return _task_store.get(task_id)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def set_task_status(task_id: str, status: dict):
|
|||
|
|
"""设置任务状态"""
|
|||
|
|
_task_store[task_id] = status
|