import asyncio from sqlalchemy.ext.asyncio import AsyncSession from app.core.logging import get_logger from app.module.cleaning.repository import CleaningTaskRepository from app.module.cleaning.runtime_client import RuntimeClient logger = get_logger(__name__) # 模块级单例 _scheduler: "CleaningTaskScheduler | None" = None def get_scheduler() -> "CleaningTaskScheduler": global _scheduler if _scheduler is None: _scheduler = CleaningTaskScheduler( CleaningTaskRepository(None), RuntimeClient() ) return _scheduler class CleaningTaskScheduler: def __init__(self, task_repo: CleaningTaskRepository, runtime_client: RuntimeClient): self.task_repo = task_repo self.runtime_client = runtime_client self._polling_task_ids: set[str] = set() self._polling_started: bool = False self._poll_failure_count: dict[str, int] = {} self._not_found_counts: dict[str, int] = {} # 🔴 新增：记录 NOT_FOUND 次数 self._MAX_POLL_FAILURES = 1 async def execute_task(self, db: AsyncSession, task_id: str, retry_count: int) -> bool: from app.module.cleaning.schema import CleaningTaskDto, CleaningTaskStatus from datetime import datetime logger.info(f"🚀 [Scheduler] 任务 {task_id} 开始执行...") task = CleaningTaskDto() task.id = task_id task.status = CleaningTaskStatus.RUNNING task.started_at = datetime.now() task.retry_count = retry_count submitted = await self.runtime_client.submit_task(task_id, retry_count) if submitted: await self.task_repo.update_task(db, task) self._polling_task_ids.add(task_id) return submitted task.status = CleaningTaskStatus.FAILED task.finished_at = datetime.now() await self.task_repo.update_task(db, task) return submitted async def stop_task(self, db: AsyncSession, task_id: str) -> bool: from app.module.cleaning.schema import CleaningTaskDto, CleaningTaskStatus self._polling_task_ids.discard(task_id) self._poll_failure_count.pop(task_id, None) self._not_found_counts.pop(task_id, None) # 🔴 清理计数器 await self.runtime_client.stop_task(task_id) task = CleaningTaskDto() task.id = task_id task.status = CleaningTaskStatus.STOPPED await self.task_repo.update_task(db, task) return True async def startup(self): from app.module.cleaning.schema import CleaningTaskStatus from app.db.session import AsyncSessionLocal try: async with AsyncSessionLocal() as db: tasks = await self.task_repo.find_tasks(db, status=CleaningTaskStatus.RUNNING) for task in tasks: logger.info(f"[Polling] Recovered RUNNING task: {task.id}") self._polling_task_ids.add(task.id) except Exception as e: logger.error(f"[Polling] Failed to recover tasks: {e}") if not self._polling_started: self._polling_started = True asyncio.create_task(self._poll_all_tasks()) async def _poll_all_tasks(self): from app.module.cleaning.schema import CleaningTaskDto, CleaningTaskStatus from app.db.session import AsyncSessionLocal from datetime import datetime logger.info("[Polling] Global polling loop started") terminal_statuses = {"completed", "failed", "cancelled", "stopped"} while True: task_ids = list(self._polling_task_ids) if not task_ids: await asyncio.sleep(2) continue for task_id in task_ids: try: status_data = await self.runtime_client.get_task_status(task_id) if status_data is None: count = self._poll_failure_count.get(task_id, 0) + 1 self._poll_failure_count[task_id] = count if count >= self._MAX_POLL_FAILURES: logger.error(f"[Polling] Task {task_id} FAILED (runtime unreachable)") async with AsyncSessionLocal() as db: t = CleaningTaskDto() t.id = task_id t.status = CleaningTaskStatus.FAILED t.finished_at = datetime.now() await self.task_repo.update_task(db, t) await db.commit() self._polling_task_ids.discard(task_id) self._poll_failure_count.pop(task_id, None) continue self._poll_failure_count.pop(task_id, None) current_status = (status_data.get("status", "") or "").lower() # 🔴 核心修复：处理 NOT_FOUND if current_status == 'not_found': cnt = self._not_found_counts.get(task_id, 0) + 1 self._not_found_counts[task_id] = cnt # 阶段 1：启动容忍期 (10秒内认为是在初始化) if cnt <= 5: continue # 继续轮询，不报错 # 阶段 2：检查数据库状态 (判断是真丢失还是已完成) try: async with AsyncSessionLocal() as check_db: db_task = await self.task_repo.find_task_by_id(check_db, task_id) db_status = db_task.status if db_task else "UNKNOWN" # ✅ 成功收尾：如果数据库已经是终态，说明任务其实跑完了 # 此时 Runtime 清理了内存状态返回 NOT_FOUND，我们应停止轮询 if db_status in ["COMPLETED", "FAILED", "STOPPED"]: logger.info(f"[Scheduler] 任务 {task_id} 已结束 (DB: {db_status})，停止轮询。") self._polling_task_ids.discard(task_id) self._not_found_counts.pop(task_id, None) continue # ✅ 失败收尾：DB 还是 RUNNING，但 Runtime 彻底找不到它了 (>60秒) if cnt > 30: logger.error(f"[Scheduler] 任务 {task_id} 在 Runtime 丢失，标记 FAILED。") t = CleaningTaskDto() t.id = task_id t.status = CleaningTaskStatus.FAILED t.finished_at = datetime.now() await self.task_repo.update_task(check_db, t) await check_db.commit() self._polling_task_ids.discard(task_id) self._not_found_counts.pop(task_id, None) except Exception: pass continue # 继续等待 # 如果找到了，重置计数器 self._not_found_counts.pop(task_id, None) # ========================================== if current_status in terminal_statuses: logger.info(f"[Scheduler] 任务 {task_id} 结束: {current_status}") async with AsyncSessionLocal() as db: t = CleaningTaskDto() t.id = task_id t.status = CleaningTaskStatus.COMPLETED if current_status == "completed" else CleaningTaskStatus.FAILED t.finished_at = datetime.now() await self.task_repo.update_task(db, t) await db.commit() self._polling_task_ids.discard(task_id) self._poll_failure_count.pop(task_id, None) except Exception as e: logger.error(f"[Polling] Error polling task {task_id}: {e}") await asyncio.sleep(2)