vent_agent_2.0/scheduler/scheduler.py

"""APScheduler 装配。

用 ``AsyncIOScheduler`` 与主事件循环共用同一个 loop——非 async 的调度器
会在自己的线程里跑任务，那样 ``agent.astream`` / aiosqlite 的 loop 绑定
全部失效。这是 1.5 明确记录过的约束。

启动时把 DB 里 enabled 的任务全量装载进调度器：DB 是任务定义的**事实源**，
调度器内存里的 job 只是它的一份投影。所以改任务必须「先写 DB，再同步
job」，反之会在重启后丢失。
"""

from __future__ import annotations

import asyncio
from typing import Any

from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.triggers.cron import CronTrigger

from db import schedule_store
from obs import get_logger
from scheduler.runner import run_schedule

logger = get_logger(__name__)

_scheduler: AsyncIOScheduler | None = None
_active: set[asyncio.Task] = set()


def _job_id(schedule_id: str) -> str:
    return f"schedule:{schedule_id}"


def _parse_cron(expr: str) -> CronTrigger:
    """解析 5 段 cron。非法表达式抛 ValueError，由调用方转成 400。"""
    return CronTrigger.from_crontab((expr or "").strip(), timezone="Asia/Shanghai")


async def _invoke(schedule_id: str) -> None:
    """job 回调。吞掉异常——见 scheduler/runner.run_schedule 的说明。"""
    task = asyncio.current_task()
    _active.add(task)
    try:
        await run_schedule(schedule_id, trigger="cron")
    except Exception:  # runner 已记录带 schedule_id / trace_id 的错误，避免重复堆栈
        pass
    finally:
        _active.discard(task)


async def _invoke_batch(callback):
    """分日计划可能等待网关，关闭时与普通定时任务一起取消并等待。"""
    task = asyncio.current_task()
    _active.add(task)
    try:
        return await callback()
    finally:
        _active.discard(task)


def get_scheduler() -> AsyncIOScheduler | None:
    """取调度器实例；未启动返回 None。"""
    return _scheduler


def is_running() -> bool:
    return _scheduler is not None and _scheduler.running


async def start() -> int:
    """启动调度器并装载已启用的任务。返回装载条数。"""
    global _scheduler
    if _scheduler is not None and _scheduler.running:
        return len(_scheduler.get_jobs())

    _scheduler = AsyncIOScheduler(timezone="Asia/Shanghai")
    _scheduler.start()

    schedules = await asyncio.to_thread(schedule_store.list_schedules, None, True, 500)
    loaded = 0
    for schedule in schedules:
        try:
            add_job(schedule)
            loaded += 1
        except Exception as exc:
            # 单条 cron 非法不该让整个调度器起不来，但必须留下痕迹
            logger.error(
                "装载定时任务失败 id=%s cron=%r：%s: %s",
                schedule.get("id"), schedule.get("cron"), type(exc).__name__, exc,
            )
    import config
    expression = config.get_str('MAIN_FAN_BATCH_CRON')
    if config.get_bool('MAIN_FAN_LOCAL_BATCH_ENABLED', False) and expression:
        from extensions.main_fan.batch_service import daily_callback, monthly_callback
        from extensions.main_fan.monthly_schedule import configured_spread_days
        _scheduler.add_job(_invoke_batch, args=[monthly_callback], trigger=_parse_cron(expression),
            id='main-fan:monthly-local', name='主通风机月度本地批次', replace_existing=True,
            max_instances=1, misfire_grace_time=60, coalesce=True)
        loaded += 1
        if configured_spread_days() > 1:
            daily_expression = ' '.join(expression.split()[:2] + ['*', '*', '*'])
            _scheduler.add_job(_invoke_batch, args=[daily_callback], trigger=_parse_cron(daily_expression),
                id='main-fan:daily-local', name='主通风机月度分日批次', replace_existing=True,
                max_instances=1, misfire_grace_time=60, coalesce=True)
            loaded += 1
    logger.info("调度器已启动，装载 %d 个任务（扫描到 %d 条）", loaded, len(schedules))
    return loaded


async def shutdown() -> None:
    """停止调度器。"""
    global _scheduler
    if _scheduler is None:
        return
    try:
        _scheduler.pause()
        active = list(_active)
        for task in active:
            task.cancel()
        if active:
            await asyncio.gather(*active, return_exceptions=True)
        _scheduler.shutdown(wait=False)
    except Exception as exc:  # pragma: no cover
        logger.warning("关闭调度器异常：%s: %s", type(exc).__name__, exc)
    _scheduler = None
    logger.info("调度器已停止")


def add_job(schedule: dict[str, Any]) -> str:
    """把一个任务定义投影成调度器 job。"""
    if _scheduler is None:
        raise RuntimeError("调度器未启动")
    schedule_id = schedule["id"]
    trigger = _parse_cron(schedule["cron"])
    _scheduler.add_job(
        _invoke,
        trigger=trigger,
        args=[schedule_id],
        id=_job_id(schedule_id),
        name=schedule.get("title") or schedule_id,
        replace_existing=True,
        max_instances=1,
        # 错过触发时间（服务重启/长时间 GC）时不要补跑：定时任务多是
        # 「每日汇报」这类语义，补跑历史时点没有意义，还会突然烧一批额度
        misfire_grace_time=60,
        coalesce=True,
    )
    logger.info("已装载定时任务 id=%s cron=%s", schedule_id, schedule["cron"])
    return _job_id(schedule_id)


def remove_job(schedule_id: str) -> bool:
    """摘掉 job。任务定义被删或停用时调用。"""
    if _scheduler is None:
        return False
    job = _scheduler.get_job(_job_id(schedule_id))
    if job is None:
        return False
    _scheduler.remove_job(_job_id(schedule_id))
    logger.info("已摘除定时任务 id=%s", schedule_id)
    return True


def sync_job(schedule: dict[str, Any]) -> None:
    """按任务定义的当前状态同步 job：enabled 则装载，否则摘除。"""
    if schedule.get("enabled"):
        add_job(schedule)
    else:
        remove_job(schedule["id"])


def jobs() -> list[dict[str, Any]]:
    """当前调度器里的 job 快照，供 /health 展示。"""
    if _scheduler is None:
        return []
    return [
        {
            "id": job.id,
            "name": job.name,
            "next_run_time": job.next_run_time.isoformat() if job.next_run_time else None,
            "trigger": str(job.trigger),
        }
        for job in _scheduler.get_jobs()
    ]