"""通用任务调度器 —— CAS 锁、并发控制、超时检测、后台执行 重构要点（v2）: - 不再继承 TaskHandler（scheduler 不是 handler，是编排层） - 接收 TaskRegistry 实例，消除全局 _TASK_HANDLER_REGISTRY - 使用结构化 LogRecord 替代自由 dict """ import asyncio import json import time from datetime import datetime, timedelta from typing import Any, Callable, Dict, List, Optional from supply.domain.task.config import TaskConfig, TaskStatus, get_default_task_config from supply.domain.task.base import TaskRegistry from supply.domain.task.exceptions import ( TaskConcurrencyError, TaskValidationError, ) from supply.infra import TaskScheduleResponse AlertCallback = Callable[[str, dict], Any] def _validate_table_name(table_name: str) -> str: cleaned = table_name.strip().replace(" ", "_") if not cleaned or "--" in cleaned or ";" in cleaned.lower(): raise TaskValidationError(f"非法的表名: {table_name}") return cleaned def _validate_task_name(task_name: Any) -> str: task_name = str(task_name).strip() if not task_name: raise TaskValidationError("task_name 不能为空") return task_name def _format_error(e: Exception) -> str: import traceback return "".join( traceback.format_exception(type(e), e, e.__traceback__) )[-1000:] class TaskScheduler: """通用任务调度器 —— CAS 锁、并发控制、超时检测、后台执行 使用方式: registry = TaskRegistry() # ... 注册 handler ... scheduler = TaskScheduler( data={"task_name": "my_crawl", "date_string": "2026-06-30"}, registry=registry, log_service=log, db_client=db, trace_id="trace-123", task_configs={"my_crawl": TaskConfig(timeout=600, max_concurrent=3)}, table_name="my_task_table", alert_callback=my_alert_func, ) result = await scheduler.deal() """ def __init__( self, data: dict, registry: TaskRegistry, *, log_service: Any = None, db_client: Any = None, trace_id: str = "", task_configs: Optional[Dict[str, TaskConfig]] = None, table_name: str = "task_manager", alert_callback: Optional[AlertCallback] = None, ): self.data = data self.registry = registry self.log_service = log_service self.db_client = db_client self.trace_id = trace_id self.table = _validate_table_name(table_name) self.task_configs = task_configs or {} self._alert_callback = alert_callback # ==================== 日志 ==================== async def _log(self, event: str, **kwargs) -> None: """结构化日志推送""" if self.log_service is None: return try: from supply.infra.observability.log_schema import ( LogRecord, LogLevel, LogCategory, ) record = LogRecord( level=LogLevel.INFO, category=LogCategory.TASK, event=event, trace_id=self.trace_id, task_name=self.data.get("task_name", ""), extra=kwargs, ) await self.log_service.log(record) except Exception: pass async def _alert(self, title: str, detail: dict) -> None: if self._alert_callback: try: await self._alert_callback(title, detail) except Exception: pass def get_task_config(self, task_name: str) -> TaskConfig: return self.task_configs.get(task_name, get_default_task_config()) # ==================== 数据库操作 ==================== async def _insert_or_ignore_task(self, task_name: str, date_str: str) -> None: query = f""" INSERT IGNORE INTO {self.table} (date_string, task_name, start_timestamp, task_status, trace_id, data) VALUES (%s, %s, %s, %s, %s, %s) """ await self.db_client.async_save( query=query, params=( date_str, task_name, int(time.time()), TaskStatus.INIT, self.trace_id, json.dumps(self.data, ensure_ascii=False), ), ) async def _try_lock_task(self) -> bool: query = f""" UPDATE {self.table} SET task_status = %s WHERE trace_id = %s AND task_status = %s """ result = await self.db_client.async_save( query=query, params=(TaskStatus.PROCESSING, self.trace_id, TaskStatus.INIT), ) return bool(result) async def _release_task(self, status: int) -> None: query = f""" UPDATE {self.table} SET task_status = %s, finish_timestamp = %s WHERE trace_id = %s AND task_status = %s """ await self.db_client.async_save( query=query, params=(status, int(time.time()), self.trace_id, TaskStatus.PROCESSING), ) async def _get_processing_tasks(self, task_name: str) -> List[Dict[str, Any]]: query = f""" SELECT trace_id, start_timestamp, data FROM {self.table} WHERE task_status = %s AND task_name = %s """ rows = await self.db_client.async_fetch( query=query, params=(TaskStatus.PROCESSING, task_name), ) return rows or [] # ==================== 并发/超时检查 ==================== async def _check_task_concurrency_and_timeout(self, task_name: str) -> None: processing_tasks = await self._get_processing_tasks(task_name) if not processing_tasks: return config = self.get_task_config(task_name) current_time = int(time.time()) # 超时检测：自动释放超时任务 timeout_tasks = [ t for t in processing_tasks if current_time - t["start_timestamp"] > config.timeout ] if timeout_tasks: await self._log( "task_timeout_detected", timeout_count=len(timeout_tasks), timeout_traces=[t["trace_id"] for t in timeout_tasks], ) await self._alert( title=f"Task Timeout Alert: {task_name}", detail={ "task_name": task_name, "timeout_count": len(timeout_tasks), "timeout_threshold": config.timeout, "timeout_tasks": [ { "trace_id": t["trace_id"], "running_time": current_time - t["start_timestamp"], } for t in timeout_tasks ], }, ) for t in timeout_tasks: await self._force_release_task(t["trace_id"], TaskStatus.FAILED) # 并发限制 active_tasks = [ t for t in processing_tasks if current_time - t["start_timestamp"] <= config.timeout ] if len(active_tasks) >= config.max_concurrent: await self._log( "task_concurrency_limit", current_count=len(active_tasks), max_concurrent=config.max_concurrent, ) await self._alert( title=f"Task Concurrency Limit: {task_name}", detail={ "task_name": task_name, "current_count": len(active_tasks), "max_concurrent": config.max_concurrent, "active_tasks": [t["trace_id"] for t in active_tasks], }, ) raise TaskConcurrencyError( f"Task {task_name} has reached max concurrency limit " f"({len(active_tasks)}/{config.max_concurrent})", task_name=task_name, ) # ==================== 任务执行 ==================== async def _run_with_guard( self, task_name: str, date_str: str, handler_callable: Callable, ) -> dict: """带保护的任务执行: 并发检查 → CAS 锁 → 后台执行 → 释放""" # 1. 并发/超时检查 try: await self._check_task_concurrency_and_timeout(task_name) except TaskConcurrencyError as e: return TaskScheduleResponse.fail("5005", str(e)) # 2. 创建任务记录并 CAS 获取锁 await self._insert_or_ignore_task(task_name, date_str) if not await self._try_lock_task(): return TaskScheduleResponse.fail("5001", "Task is already processing") # 3. 后台执行 async def _task_wrapper(): status = TaskStatus.FAILED config = self.get_task_config(task_name) start_time = time.time() try: await self._log("task_started") result = await handler_callable() status = result if isinstance(result, int) else TaskStatus.SUCCESS duration = time.time() - start_time await self._log( "task_completed", status=status, duration_ms=int(duration * 1000), ) except Exception as e: duration = time.time() - start_time error_detail = _format_error(e) await self._log( "task_failed", error=error_detail, duration_ms=int(duration * 1000), ) if config.alert_on_failure: await self._alert( title=f"Task Failed: {task_name}", detail={ "task_name": task_name, "trace_id": self.trace_id, "error": error_detail, "duration": duration, }, ) finally: await self._release_task(status) asyncio.create_task(_task_wrapper(), name=f"{task_name}_{self.trace_id}") return TaskScheduleResponse.success( task_name=task_name, data={ "code": 0, "message": "Task started successfully", "trace_id": self.trace_id, }, ) # ==================== 任务管理接口 ==================== async def get_task_status(self, trace_id: Optional[str] = None) -> Optional[Dict[str, Any]]: tid = trace_id or self.trace_id query = f"SELECT * FROM {self.table} WHERE trace_id = %s" return await self.db_client.async_fetch_one(query, params=(tid,)) async def cancel_task(self, trace_id: Optional[str] = None) -> bool: tid = trace_id or self.trace_id query = f""" UPDATE {self.table} SET task_status = %s, finish_timestamp = %s WHERE trace_id = %s AND task_status IN (%s, %s) """ result = await self.db_client.async_save( query, params=( TaskStatus.FAILED, int(time.time()), tid, TaskStatus.INIT, TaskStatus.PROCESSING, ), ) if result: await self._log("task_cancelled", trace_id=tid) return bool(result) async def retry_task(self, trace_id: Optional[str] = None) -> bool: tid = trace_id or self.trace_id query = f""" UPDATE {self.table} SET task_status = %s, start_timestamp = %s, finish_timestamp = NULL WHERE trace_id = %s """ result = await self.db_client.async_save( query, params=(TaskStatus.INIT, int(time.time()), tid) ) if result: await self._log("task_retried", trace_id=tid) return bool(result) async def _force_release_task(self, trace_id: str, status: int) -> None: query = f""" UPDATE {self.table} SET task_status = %s, finish_timestamp = %s WHERE trace_id = %s """ await self.db_client.async_save( query, params=(status, int(time.time()), trace_id) ) await self._log("task_force_released", trace_id=trace_id, status=status) # ==================== 主入口 ==================== async def deal(self) -> dict: """任务调度主入口""" task_name = self.data.get("task_name") if not task_name: return TaskScheduleResponse.fail("4003", "task_name is required") try: task_name = _validate_task_name(task_name) except TaskValidationError as e: return TaskScheduleResponse.fail("4003", str(e)) date_str = self.data.get("date_string") or ( datetime.utcnow() + timedelta(hours=8) ).strftime("%Y-%m-%d") # 从 TaskRegistry 查找 handler handler_cls = self.registry.get(task_name) if handler_cls is None: return TaskScheduleResponse.fail( "4001", f"Unknown task: {task_name}. " f"Available tasks: {', '.join(self.registry.list_names())}", ) # 实例化 handler handler = handler_cls(log_service=self.log_service, db_client=self.db_client) return await self._run_with_guard( task_name, date_str, lambda: handler.execute( type("Task", (), {"task_name": task_name, "params": self.data, "trace_id": self.trace_id})() ), )