321 lines
8.8 KiB
Python
321 lines
8.8 KiB
Python
|
"""
|
|||
|
|
用于后台任务处理的异步作业管理器。
|
|||
|
|
"""
|
|||
|
|
import asyncio
|
|||
|
|
import uuid
|
|||
|
|
from datetime import datetime
|
|||
|
|
from typing import Dict, Any, Callable, Optional
|
|||
|
|
from enum import Enum
|
|||
|
|
|
|||
|
|
from app.utils.logger import get_logger
|
|||
|
|
|
|||
|
|
logger = get_logger(__name__)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class JobStatus(str, Enum):
|
|||
|
|
"""作业状态枚举。"""
|
|||
|
|
PENDING = "pending"
|
|||
|
|
RUNNING = "running"
|
|||
|
|
COMPLETED = "completed"
|
|||
|
|
FAILED = "failed"
|
|||
|
|
|
|||
|
|
|
|||
|
|
class Job:
|
|||
|
|
"""作业表示。"""
|
|||
|
|
|
|||
|
|
def __init__(self, job_id: str, job_fn: Callable, *args, **kwargs):
|
|||
|
|
self.id = job_id
|
|||
|
|
self.job_fn = job_fn
|
|||
|
|
self.args = args
|
|||
|
|
self.kwargs = kwargs
|
|||
|
|
self.status = JobStatus.PENDING
|
|||
|
|
self.result = None
|
|||
|
|
self.error = None
|
|||
|
|
self.created_at = datetime.utcnow()
|
|||
|
|
self.started_at: Optional[datetime] = None
|
|||
|
|
self.completed_at: Optional[datetime] = None
|
|||
|
|
|
|||
|
|
def to_dict(self) -> Dict[str, Any]:
|
|||
|
|
"""将作业转换为字典。"""
|
|||
|
|
return {
|
|||
|
|
"id": self.id,
|
|||
|
|
"status": self.status.value,
|
|||
|
|
"created_at": self.created_at.isoformat(),
|
|||
|
|
"started_at": self.started_at.isoformat() if self.started_at else None,
|
|||
|
|
"completed_at": self.completed_at.isoformat() if self.completed_at else None,
|
|||
|
|
"error": str(self.error) if self.error else None,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
class AsyncJobManager:
|
|||
|
|
"""
|
|||
|
|
使用 asyncio 管理异步后台作业。
|
|||
|
|
|
|||
|
|
注意：这是用于教育目的的简化实现。
|
|||
|
|
在生产环境中，请考虑使用 Celery + Redis 或类似解决方案。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self):
|
|||
|
|
"""初始化作业管理器。"""
|
|||
|
|
self.jobs: Dict[str, Job] = {}
|
|||
|
|
self._lock = asyncio.Lock()
|
|||
|
|
|
|||
|
|
async def submit_job(
|
|||
|
|
self,
|
|||
|
|
job_fn: Callable,
|
|||
|
|
*args,
|
|||
|
|
**kwargs
|
|||
|
|
) -> str:
|
|||
|
|
"""
|
|||
|
|
提交作业进行异步执行。
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
job_fn: 要执行的可调用函数
|
|||
|
|
*args: 函数的位置参数
|
|||
|
|
**kwargs: 函数的关键字参数
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
str: 作业 ID
|
|||
|
|
"""
|
|||
|
|
job_id = str(uuid.uuid4())
|
|||
|
|
job = Job(job_id, job_fn, *args, **kwargs)
|
|||
|
|
|
|||
|
|
async with self._lock:
|
|||
|
|
self.jobs[job_id] = job
|
|||
|
|
|
|||
|
|
logger.info("job_submitted", job_id=job_id)
|
|||
|
|
|
|||
|
|
# 在后台启动作业执行
|
|||
|
|
asyncio.create_task(self._execute_job(job))
|
|||
|
|
|
|||
|
|
return job_id
|
|||
|
|
|
|||
|
|
async def _execute_job(self, job: Job) -> None:
|
|||
|
|
"""
|
|||
|
|
在后台执行作业。
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
job: 要执行的作业实例
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
job.status = JobStatus.RUNNING
|
|||
|
|
job.started_at = datetime.utcnow()
|
|||
|
|
|
|||
|
|
logger.info("job_started", job_id=job.id)
|
|||
|
|
|
|||
|
|
# 检查函数是异步还是同步
|
|||
|
|
if asyncio.iscoroutinefunction(job.job_fn):
|
|||
|
|
result = await job.job_fn(*job.args, **job.kwargs)
|
|||
|
|
else:
|
|||
|
|
# 在执行器中运行同步函数以避免阻塞
|
|||
|
|
loop = asyncio.get_event_loop()
|
|||
|
|
result = await loop.run_in_executor(
|
|||
|
|
None,
|
|||
|
|
lambda: job.job_fn(*job.args, **job.kwargs)
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
job.result = result
|
|||
|
|
job.status = JobStatus.COMPLETED
|
|||
|
|
job.completed_at = datetime.utcnow()
|
|||
|
|
|
|||
|
|
logger.info("job_completed", job_id=job.id)
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
job.error = e
|
|||
|
|
job.status = JobStatus.FAILED
|
|||
|
|
job.completed_at = datetime.utcnow()
|
|||
|
|
|
|||
|
|
logger.error("job_failed", job_id=job.id, error=str(e))
|
|||
|
|
|
|||
|
|
async def get_job_status(self, job_id: str) -> Dict[str, Any]:
|
|||
|
|
"""
|
|||
|
|
获取作业状态和元数据。
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
job_id: 作业 ID
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
Dict[str, Any]: 作业状态信息
|
|||
|
|
|
|||
|
|
Raises:
|
|||
|
|
KeyError: 如果找不到作业
|
|||
|
|
"""
|
|||
|
|
async with self._lock:
|
|||
|
|
if job_id not in self.jobs:
|
|||
|
|
raise KeyError(f"Job not found: {job_id}")
|
|||
|
|
|
|||
|
|
job = self.jobs[job_id]
|
|||
|
|
return job.to_dict()
|
|||
|
|
|
|||
|
|
async def get_job_result(self, job_id: str) -> Any:
|
|||
|
|
"""
|
|||
|
|
获取作业结果。
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
job_id: 作业 ID
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
Any: 作业结果
|
|||
|
|
|
|||
|
|
Raises:
|
|||
|
|
KeyError: 如果找不到作业
|
|||
|
|
RuntimeError: 如果作业未完成或失败
|
|||
|
|
"""
|
|||
|
|
async with self._lock:
|
|||
|
|
if job_id not in self.jobs:
|
|||
|
|
raise KeyError(f"Job not found: {job_id}")
|
|||
|
|
|
|||
|
|
job = self.jobs[job_id]
|
|||
|
|
|
|||
|
|
if job.status == JobStatus.FAILED:
|
|||
|
|
raise RuntimeError(f"Job failed: {job.error}")
|
|||
|
|
|
|||
|
|
if job.status != JobStatus.COMPLETED:
|
|||
|
|
raise RuntimeError(f"Job not completed yet. Status: {job.status}")
|
|||
|
|
|
|||
|
|
return job.result
|
|||
|
|
|
|||
|
|
async def wait_for_job(
|
|||
|
|
self,
|
|||
|
|
job_id: str,
|
|||
|
|
timeout: Optional[float] = None
|
|||
|
|
) -> Any:
|
|||
|
|
"""
|
|||
|
|
等待作业完成并返回结果。
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
job_id: 作业 ID
|
|||
|
|
timeout: 可选的超时时间（秒）
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
Any: 作业结果
|
|||
|
|
|
|||
|
|
Raises:
|
|||
|
|
KeyError: 如果找不到作业
|
|||
|
|
asyncio.TimeoutError: 如果超时
|
|||
|
|
RuntimeError: 如果作业失败
|
|||
|
|
"""
|
|||
|
|
start_time = asyncio.get_event_loop().time()
|
|||
|
|
|
|||
|
|
while True:
|
|||
|
|
async with self._lock:
|
|||
|
|
if job_id not in self.jobs:
|
|||
|
|
raise KeyError(f"Job not found: {job_id}")
|
|||
|
|
|
|||
|
|
job = self.jobs[job_id]
|
|||
|
|
|
|||
|
|
if job.status == JobStatus.COMPLETED:
|
|||
|
|
return job.result
|
|||
|
|
|
|||
|
|
if job.status == JobStatus.FAILED:
|
|||
|
|
raise RuntimeError(f"Job failed: {job.error}")
|
|||
|
|
|
|||
|
|
# 检查超时
|
|||
|
|
if timeout is not None:
|
|||
|
|
elapsed = asyncio.get_event_loop().time() - start_time
|
|||
|
|
if elapsed >= timeout:
|
|||
|
|
raise asyncio.TimeoutError(f"Job timeout after {timeout}s")
|
|||
|
|
|
|||
|
|
# 稍等片刻再检查
|
|||
|
|
await asyncio.sleep(0.1)
|
|||
|
|
|
|||
|
|
async def cancel_job(self, job_id: str) -> bool:
|
|||
|
|
"""
|
|||
|
|
取消待处理的作业（尽力而为）。
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
job_id: 作业 ID
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
bool: 如果已取消则为 True，如果已运行/完成则为 False
|
|||
|
|
"""
|
|||
|
|
async with self._lock:
|
|||
|
|
if job_id not in self.jobs:
|
|||
|
|
raise KeyError(f"Job not found: {job_id}")
|
|||
|
|
|
|||
|
|
job = self.jobs[job_id]
|
|||
|
|
|
|||
|
|
if job.status == JobStatus.PENDING:
|
|||
|
|
job.status = JobStatus.FAILED
|
|||
|
|
job.error = Exception("Job cancelled by user")
|
|||
|
|
job.completed_at = datetime.utcnow()
|
|||
|
|
logger.info("job_cancelled", job_id=job_id)
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
async def list_jobs(
|
|||
|
|
self,
|
|||
|
|
status: Optional[JobStatus] = None,
|
|||
|
|
limit: int = 100
|
|||
|
|
) -> list[Dict[str, Any]]:
|
|||
|
|
"""
|
|||
|
|
列出作业，可选状态过滤器。
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
status: 可选的状态过滤器
|
|||
|
|
limit: 返回的最大作业数
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
list[Dict[str, Any]]: 作业信息列表
|
|||
|
|
"""
|
|||
|
|
async with self._lock:
|
|||
|
|
jobs = list(self.jobs.values())
|
|||
|
|
|
|||
|
|
if status:
|
|||
|
|
jobs = [j for j in jobs if j.status == status]
|
|||
|
|
|
|||
|
|
# 按创建时间排序（最新的在前）
|
|||
|
|
jobs.sort(key=lambda j: j.created_at, reverse=True)
|
|||
|
|
|
|||
|
|
return [j.to_dict() for j in jobs[:limit]]
|
|||
|
|
|
|||
|
|
async def cleanup_old_jobs(self, max_age_hours: int = 24) -> int:
|
|||
|
|
"""
|
|||
|
|
清理旧的已完成/失败作业。
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
max_age_hours: 最大年龄（小时）
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
int: 清理的作业数量
|
|||
|
|
"""
|
|||
|
|
from datetime import timedelta
|
|||
|
|
|
|||
|
|
cutoff_time = datetime.utcnow() - timedelta(hours=max_age_hours)
|
|||
|
|
cleaned = 0
|
|||
|
|
|
|||
|
|
async with self._lock:
|
|||
|
|
job_ids_to_remove = []
|
|||
|
|
|
|||
|
|
for job_id, job in self.jobs.items():
|
|||
|
|
if job.status in [JobStatus.COMPLETED, JobStatus.FAILED]:
|
|||
|
|
if job.completed_at and job.completed_at < cutoff_time:
|
|||
|
|
job_ids_to_remove.append(job_id)
|
|||
|
|
|
|||
|
|
for job_id in job_ids_to_remove:
|
|||
|
|
del self.jobs[job_id]
|
|||
|
|
cleaned += 1
|
|||
|
|
|
|||
|
|
if cleaned > 0:
|
|||
|
|
logger.info("jobs_cleaned_up", count=cleaned)
|
|||
|
|
|
|||
|
|
return cleaned
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 全局作业管理器实例
|
|||
|
|
_job_manager: Optional[AsyncJobManager] = None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_job_manager() -> AsyncJobManager:
|
|||
|
|
"""
|
|||
|
|
获取或创建全局作业管理器实例。
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
AsyncJobManager: 作业管理器实例
|
|||
|
|
"""
|
|||
|
|
global _job_manager
|
|||
|
|
if _job_manager is None:
|
|||
|
|
_job_manager = AsyncJobManager()
|
|||
|
|
return _job_manager
|