|
|
@@ -1,402 +0,0 @@
|
|
|
-"""通用任务调度器 —— 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})()
|
|
|
- ),
|
|
|
- )
|