import uuid import asyncio import hashlib from datetime import datetime, timedelta from typing import List, Optional, Dict, Any import json import sys from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select from sqlalchemy import delete, func, and_, text from ..config import read_int_env from ..models import SystemTaskORM, TaskLogORM from .. import database # --- Configuration --- # 任务超时时间 (分钟)。如果任务超过此时间未更新,视为僵尸任务。 # ENVI 任务单步可能超过 60 分钟,默认设为 120 分钟,可通过环境变量覆盖。 TASK_TIMEOUT_MINUTES = read_int_env("TASK_TIMEOUT_MINUTES", 120, minimum=10, maximum=1440) TASK_ACTIVE_DEFAULT_LIMIT = read_int_env( "TASK_ACTIVE_DEFAULT_LIMIT", 100, minimum=1, maximum=2000, ) TASK_ACTIVE_MAX_LIMIT = read_int_env( "TASK_ACTIVE_MAX_LIMIT", 500, minimum=1, maximum=10000, ) TASK_LOG_DEFAULT_LIMIT = read_int_env( "TASK_LOG_DEFAULT_LIMIT", 100, minimum=1, maximum=5000, ) TASK_LOG_MAX_LIMIT = read_int_env( "TASK_LOG_MAX_LIMIT", 1000, minimum=1, maximum=20000, ) TASK_QUERY_MAX_OFFSET = read_int_env( "TASK_QUERY_MAX_OFFSET", 500000, minimum=0, maximum=20000000, ) def _clamp_pagination(limit: int, offset: int, *, default_limit: int, max_limit: int) -> tuple[int, int]: try: parsed_limit = int(limit) except (TypeError, ValueError): parsed_limit = default_limit try: parsed_offset = int(offset) except (TypeError, ValueError): parsed_offset = 0 safe_limit = min(max_limit, max(1, parsed_limit)) safe_offset = min(TASK_QUERY_MAX_OFFSET, max(0, parsed_offset)) return safe_limit, safe_offset def _normalize_string_list(values: Optional[List[str]], *, uppercase: bool = False) -> List[str]: normalized: List[str] = [] for raw in values or []: text = str(raw or "").strip() if not text: continue if uppercase: text = text.upper() if text not in normalized: normalized.append(text) return normalized def _task_type_lock_key(task_type: str) -> int: normalized = (task_type or "").strip().lower().encode("utf-8") digest = hashlib.sha256(normalized).digest() return int.from_bytes(digest[:8], byteorder="big", signed=True) def _is_postgresql_session(db: AsyncSession) -> bool: try: bind = db.get_bind() dialect_name = getattr(getattr(bind, "dialect", None), "name", "") return (dialect_name or "").lower() == "postgresql" except Exception as exc: print(f"[WARN] _is_postgresql: {exc}") return False def get_db_session(): """ 获取数据库会话的辅助函数。 如果数据库尚未初始化,则尝试初始化。 """ if database.AsyncSessionLocal is None: print(">>> [TaskService] 检测到 AsyncSessionLocal 未初始化,正在尝试重新初始化数据库...") try: database.init_db() print(">>> [TaskService] 数据库初始化成功。") except Exception as e: print(f">>> [TaskService] 数据库初始化失败: {e}") raise if database.AsyncSessionLocal is None: raise RuntimeError("无法获取数据库会话,请检查数据库连接配置。") return database.AsyncSessionLocal() class TaskService: """ 持久化任务服务。 负责管理任务的生命周期、日志记录和状态持久化。 """ def __init__(self): pass async def _expire_zombie_tasks(self, db: AsyncSession) -> None: timeout_threshold = datetime.now() - timedelta(minutes=TASK_TIMEOUT_MINUTES) result = await db.execute( select(SystemTaskORM).where( and_( SystemTaskORM.status == "RUNNING", SystemTaskORM.updated_at < timeout_threshold ) ) ) zombie_tasks = result.scalars().all() for task in zombie_tasks: task.status = "FAILED" task.message = "系统检测超时: 任务被认为已失效 (心跳超时)" log = TaskLogORM(task_id=task.task_id, log_level="WARNING", message="任务因超时被自动标记为失败") db.add(log) if zombie_tasks: await db.commit() async def create_task( self, task_type: str, task_name: str, params: Optional[Dict[str, Any]] = None, db: Optional[AsyncSession] = None ) -> str: """ 创建新任务。 实现逻辑:检查是否存在同类型且运行中的任务。如果存在,则抛出异常。 """ gen_db = db is None if gen_db: db = get_db_session() try: # 1. 并发冲突检查:是否已有同类型的任务在运行? if _is_postgresql_session(db): lock_key = _task_type_lock_key(task_type) await db.execute( text("SELECT pg_advisory_xact_lock(:lock_key)"), {"lock_key": lock_key}, ) status_check = await db.execute( select(SystemTaskORM) .where( and_( SystemTaskORM.task_type == task_type, SystemTaskORM.status.in_(["PENDING", "RUNNING"]) ) ) ) existing_task = status_check.scalar_one_or_none() if existing_task: raise ValueError(f"任务冲突: 已存在一个正在运行的 {task_type} 任务 (ID: {existing_task.task_id})。请等待该任务完成。") # 2. 创建新任务记录 task_id = str(uuid.uuid4()) new_task = SystemTaskORM( task_id=task_id, task_type=task_type, task_name=task_name, status="PENDING", params=params ) db.add(new_task) # 3. 记录创建日志 log_entry = TaskLogORM( task_id=task_id, log_level="INFO", message=f"任务已创建: {task_name} ({task_type})" ) db.add(log_entry) await db.commit() return task_id except Exception as e: await db.rollback() raise e finally: if gen_db: await db.close() async def start_task(self, task_id: str, message: str = "任务开始执行", db: Optional[AsyncSession] = None): """将任务状态更新为 RUNNING""" gen_db = db is None if gen_db: db = get_db_session() try: result = await db.execute( select(SystemTaskORM).where(SystemTaskORM.task_id == task_id) ) task = result.scalar_one_or_none() if task: task.status = "RUNNING" task.message = message task.started_at = datetime.now() await self._add_log(task_id, "INFO", "任务已开始", db=db) await db.commit() else: raise ValueError(f"任务未找到: {task_id}") finally: if gen_db: await db.close() async def update_task( self, task_id: str, status: Optional[str] = None, progress: Optional[int] = None, message: Optional[str] = None, db: Optional[AsyncSession] = None ): """ 更新任务状态。 如果状态变为 COMPLETED 或 FAILED,标记结束时间。 """ gen_db = db is None if gen_db: db = get_db_session() try: result = await db.execute( select(SystemTaskORM).where(SystemTaskORM.task_id == task_id) ) task = result.scalar_one_or_none() if task: if status: task.status = status if progress is not None: task.progress = progress if message: task.message = message # 如果任务结束,更新结束时间 if status in ["COMPLETED", "FAILED", "CANCELLED"]: task.ended_at = datetime.now() await self._add_log(task_id, "INFO", f"任务已结束: {status}", db=db) else: # 显式更新心跳时间,防止 SQLAlchemy 因属性未变而跳过 UPDATE task.updated_at = datetime.now() await db.commit() else: print(f"警告: 尝试更新一个不存在的任务 {task_id}") finally: if gen_db: await db.close() async def get_active_tasks( self, limit: int = TASK_ACTIVE_DEFAULT_LIMIT, offset: int = 0, db: Optional[AsyncSession] = None, ) -> List[SystemTaskORM]: """ 获取所有非结束状态的任务。 注意:这里应该加上超时检测逻辑,如果 updated_at 太旧,应该标记为 FAILED。 """ gen_db = db is None if gen_db: db = get_db_session() try: safe_limit, safe_offset = _clamp_pagination( limit, offset, default_limit=TASK_ACTIVE_DEFAULT_LIMIT, max_limit=TASK_ACTIVE_MAX_LIMIT, ) await self._expire_zombie_tasks(db) # 2. 获取活跃任务 active_result = await db.execute( select(SystemTaskORM) .where(SystemTaskORM.status.in_(["PENDING", "RUNNING"])) .order_by(SystemTaskORM.created_at.desc()) .offset(safe_offset) .limit(safe_limit) ) return active_result.scalars().all() finally: if gen_db: await db.close() async def list_tasks( self, task_types: Optional[List[str]] = None, statuses: Optional[List[str]] = None, limit: int = TASK_ACTIVE_DEFAULT_LIMIT, offset: int = 0, db: Optional[AsyncSession] = None, ) -> List[SystemTaskORM]: gen_db = db is None if gen_db: db = get_db_session() try: safe_limit, safe_offset = _clamp_pagination( limit, offset, default_limit=TASK_ACTIVE_DEFAULT_LIMIT, max_limit=TASK_ACTIVE_MAX_LIMIT, ) await self._expire_zombie_tasks(db) normalized_task_types = _normalize_string_list(task_types) normalized_statuses = _normalize_string_list(statuses, uppercase=True) stmt = select(SystemTaskORM) if normalized_task_types: stmt = stmt.where(SystemTaskORM.task_type.in_(normalized_task_types)) if normalized_statuses: stmt = stmt.where(SystemTaskORM.status.in_(normalized_statuses)) stmt = ( stmt .order_by(SystemTaskORM.created_at.desc(), SystemTaskORM.id.desc()) .offset(safe_offset) .limit(safe_limit) ) result = await db.execute(stmt) return result.scalars().all() finally: if gen_db: await db.close() async def get_task(self, task_id: str, db: Optional[AsyncSession] = None) -> Optional[SystemTaskORM]: gen_db = db is None if gen_db: db = get_db_session() try: result = await db.execute( select(SystemTaskORM).where(SystemTaskORM.task_id == task_id) ) return result.scalar_one_or_none() finally: if gen_db: await db.close() async def get_logs( self, task_id: str, limit: int = TASK_LOG_DEFAULT_LIMIT, offset: int = 0, db: Optional[AsyncSession] = None, ) -> List[TaskLogORM]: """获取任务日志""" gen_db = db is None if gen_db: db = get_db_session() try: safe_limit, safe_offset = _clamp_pagination( limit, offset, default_limit=TASK_LOG_DEFAULT_LIMIT, max_limit=TASK_LOG_MAX_LIMIT, ) result = await db.execute( select(TaskLogORM) .where(TaskLogORM.task_id == task_id) .order_by(TaskLogORM.timestamp.desc()) .offset(safe_offset) .limit(safe_limit) ) # 返回正序列表以便阅读 logs = result.scalars().all() return list(reversed(logs)) finally: if gen_db: await db.close() async def _add_log(self, task_id: str, level: str, message: str, db: AsyncSession): """内部方法:添加日志""" log_entry = TaskLogORM(task_id=task_id, log_level=level, message=message) db.add(log_entry) # 不在这里 commit,由上层控制事务 async def add_log(self, task_id: str, level: str, message: str): """公开方法:添加日志 (带独立事务)""" db = get_db_session() try: await self._add_log(task_id, level, message, db=db) await db.commit() finally: await db.close() # 创建全局单例实例 async def delete_log( self, task_id: str, log_id: int, db: Optional[AsyncSession] = None, ) -> bool: gen_db = db is None if gen_db: db = get_db_session() try: result = await db.execute( select(TaskLogORM).where( TaskLogORM.id == int(log_id), TaskLogORM.task_id == task_id, ) ) log_entry = result.scalar_one_or_none() if log_entry is None: return False await db.delete(log_entry) await db.commit() return True except Exception: await db.rollback() raise finally: if gen_db: await db.close() async def clear_logs( self, task_id: str, db: Optional[AsyncSession] = None, ) -> int: gen_db = db is None if gen_db: db = get_db_session() try: result = await db.execute( delete(TaskLogORM).where(TaskLogORM.task_id == task_id) ) await db.commit() return int(result.rowcount or 0) except Exception: await db.rollback() raise finally: if gen_db: await db.close() task_service = TaskService()