120 lines
3.2 KiB
Python
120 lines
3.2 KiB
Python
|
"""后台任务系统 - Phase 10.4"""
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import uuid
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from datetime import datetime
|
||
|
|
from typing import Any
|
||
|
|
from enum import Enum
|
||
|
|
|
||
|
|
|
||
|
|
class BackgroundTaskStatus(Enum):
|
||
|
|
PENDING = "pending"
|
||
|
|
RUNNING = "running"
|
||
|
|
COMPLETED = "completed"
|
||
|
|
FAILED = "failed"
|
||
|
|
CANCELLED = "cancelled"
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class BackgroundTask:
|
||
|
|
"""后台任务"""
|
||
|
|
|
||
|
|
id: str
|
||
|
|
name: str
|
||
|
|
status: BackgroundTaskStatus
|
||
|
|
created_at: datetime
|
||
|
|
started_at: datetime | None = None
|
||
|
|
completed_at: datetime | None = None
|
||
|
|
result: Any = None
|
||
|
|
error: str | None = None
|
||
|
|
|
||
|
|
|
||
|
|
class BackgroundTaskManager:
|
||
|
|
"""后台任务管理器"""
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
self._tasks: dict[str, BackgroundTask] = {}
|
||
|
|
self._.coroutines: dict[str, asyncio.Task] = {}
|
||
|
|
|
||
|
|
def submit_task(self, name: str, coro: Any, *args, **kwargs) -> str:
|
||
|
|
"""提交后台任务
|
||
|
|
|
||
|
|
Args:
|
||
|
|
name: 任务名称
|
||
|
|
coro: 协程函数
|
||
|
|
*args: 位置参数
|
||
|
|
**kwargs: 关键字参数
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
任务 ID
|
||
|
|
"""
|
||
|
|
task_id = str(uuid.uuid4())[:8]
|
||
|
|
|
||
|
|
# 创建任务记录
|
||
|
|
self._tasks[task_id] = BackgroundTask(
|
||
|
|
id=task_id,
|
||
|
|
name=name,
|
||
|
|
status=BackgroundTaskStatus.PENDING,
|
||
|
|
created_at=datetime.now(),
|
||
|
|
)
|
||
|
|
|
||
|
|
# 创建 asyncio task
|
||
|
|
async def run_task():
|
||
|
|
self._tasks[task_id].status = BackgroundTaskStatus.RUNNING
|
||
|
|
self._tasks[task_id].started_at = datetime.now()
|
||
|
|
try:
|
||
|
|
result = await coro(*args, **kwargs)
|
||
|
|
self._tasks[task_id].status = BackgroundTaskStatus.COMPLETED
|
||
|
|
self._tasks[task_id].result = result
|
||
|
|
except Exception as e:
|
||
|
|
self._tasks[task_id].status = BackgroundTaskStatus.FAILED
|
||
|
|
self._tasks[task_id].error = str(e)
|
||
|
|
finally:
|
||
|
|
self._tasks[task_id].completed_at = datetime.now()
|
||
|
|
if task_id in self._coroutines:
|
||
|
|
del self._coroutines[task_id]
|
||
|
|
|
||
|
|
self._coroutines[task_id] = asyncio.create_task(run_task())
|
||
|
|
return task_id
|
||
|
|
|
||
|
|
def cancel_task(self, task_id: str) -> bool:
|
||
|
|
"""取消任务
|
||
|
|
|
||
|
|
Args:
|
||
|
|
task_id: 任务 ID
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
是否成功取消
|
||
|
|
"""
|
||
|
|
if task_id not in self._tasks:
|
||
|
|
return False
|
||
|
|
|
||
|
|
if task_id in self._coroutines:
|
||
|
|
self._coroutines[task_id].cancel()
|
||
|
|
del self._coroutines[task_id]
|
||
|
|
|
||
|
|
self._tasks[task_id].status = BackgroundTaskStatus.CANCELLED
|
||
|
|
self._tasks[task_id].completed_at = datetime.now()
|
||
|
|
return True
|
||
|
|
|
||
|
|
def get_task_status(self, task_id: str) -> BackgroundTask | None:
|
||
|
|
"""获取任务状态"""
|
||
|
|
return self._tasks.get(task_id)
|
||
|
|
|
||
|
|
def list_tasks(self) -> list[BackgroundTask]:
|
||
|
|
"""列出所有任务"""
|
||
|
|
return list(self._tasks.values())
|
||
|
|
|
||
|
|
|
||
|
|
# 全局单例
|
||
|
|
_manager: BackgroundTaskManager | None = None
|
||
|
|
|
||
|
|
|
||
|
|
def get_background_task_manager() -> BackgroundTaskManager:
|
||
|
|
"""获取全局后台任务管理器"""
|
||
|
|
global _manager
|
||
|
|
if _manager is None:
|
||
|
|
_manager = BackgroundTaskManager()
|
||
|
|
return _manager
|