"""三层定时器循环。 在 FastAPI startup 事件中用 asyncio.create_task() 启动。 定时器在后台跑，每次迭代独立 try/except + timeout，不会卡死主 event loop。 """ from __future__ import annotations import asyncio import functools import logging import os import subprocess import time from datetime import datetime, timezone from typing import Any, Dict, List, Optional from app.config import get_settings from app.db import SessionLocal from app.diagnostic.config import DiagConfig, get_diag_config from app.diagnostic.engine import DiagnosticEngine, RuleResult from app.services.scanner import Scanner logger = logging.getLogger(__name__) # 诊断套件全局运行时状态（单例 dict） _DIAG_STATE: Dict[str, Any] = { "enabled": False, "docker_status": "none", "docker_info": {}, "started_at": None, "rules_cache": {}, # {rule_id: {result, ts, run_by}} "latest": {"rules": [], "generated_at": None, "docker_status": "none"}, } # 后台引用（防止 GC 收走） _TIMER_TASKS: List[asyncio.Task] = [] # ====================================================================== # 外部访问入口 # ====================================================================== def get_diag_state() -> Dict[str, Any]: """获取诊断套件运行时状态（给页面/API 查询用）。""" cache = {k: {"result": v["result"].to_dict() if isinstance(v.get("result"), RuleResult) else v.get("result"), "ts": v.get("ts"), "run_by": v.get("run_by")} for k, v in _DIAG_STATE["rules_cache"].items()} return { "enabled": _DIAG_STATE["enabled"], "docker_status": _DIAG_STATE["docker_status"], "docker_info": _DIAG_STATE["docker_info"], "started_at": _DIAG_STATE["started_at"], "rules_cache": cache, "latest": _DIAG_STATE["latest"], } # ====================================================================== # 后台入口（在 app.main.startup 中调用） # ====================================================================== def start_timers(docker_status: str, docker_info: Dict[str, Any]) -> None: """启动三个定时器循环（后台 daemon task）。""" cfg = get_diag_config() _DIAG_STATE["enabled"] = cfg.enabled _DIAG_STATE["docker_status"] = docker_status _DIAG_STATE["docker_info"] = docker_info _DIAG_STATE["started_at"] = datetime.now(timezone.utc).isoformat() if not cfg.enabled: logger.info("MeshSafe Diagnostics: disabled") return loop = asyncio.get_event_loop() if cfg.error_scan_interval > 0: _TIMER_TASKS.append(loop.create_task( _timer_loop("error_scan", cfg.error_scan_interval, lambda eng: eng.run_timer_error_scan(), docker_status, cfg.docker_log_tail), name="diag-error-scan", )) logger.info("诊断定时器启动: error_scan (%ss)", cfg.error_scan_interval) if cfg.db_aggregate_interval > 0: _TIMER_TASKS.append(loop.create_task( _timer_loop("db_aggregate", cfg.db_aggregate_interval, lambda eng: eng.run_timer_db_aggregate(), docker_status, cfg.docker_log_tail), name="diag-db-aggregate", )) logger.info("诊断定时器启动: db_aggregate (%ss)", cfg.db_aggregate_interval) if cfg.container_check_interval > 0: _TIMER_TASKS.append(loop.create_task( _timer_loop("container_check", cfg.container_check_interval, lambda eng: eng.run_timer_container(), docker_status, cfg.docker_log_tail), name="diag-container-check", )) logger.info("诊断定时器启动: container_check (%ss)", cfg.container_check_interval) # ====================================================================== # 定时器循环（async） # ====================================================================== async def _timer_loop( name: str, interval: int, runner_fn, # 接受 DiagnosticEngine 返回 List[RuleResult] docker_status: str, docker_log_tail: int, ) -> None: """通用定时器循环。每次迭代独立 try/except + 超时保护。""" # 等一个 interval 再开始第一次（避免 startup 时 DB 还没 ready） await asyncio.sleep(min(interval, 15)) while True: try: # 分布式锁（时间槽位）：多 worker 部署时，每个时间槽只有 # 一个 worker 能拿到锁并执行，保证执行频率精确等于 interval。 if not _acquire_timer_lock(name, interval): logger.debug("诊断定时器 %s 本轮跳过（其他 worker 持锁）", name) await asyncio.sleep(interval) continue results = await asyncio.wait_for( asyncio.get_event_loop().run_in_executor( None, lambda: _run_sync(name, runner_fn, docker_status, docker_log_tail), ), timeout=30.0, ) _cache_results(name, results, docker_status) logger.debug("诊断定时器 %s 跑完: %d 条规则", name, len(results)) except asyncio.CancelledError: logger.info("诊断定时器 %s 被取消", name) break except asyncio.TimeoutError: logger.warning("诊断定时器 %s 超时（30s），本轮跳过", name) except Exception as e: # noqa: BLE001 logger.exception("诊断定时器 %s 异常: %s，%ss 后重试", name, e, interval) await asyncio.sleep(interval) # ====================================================================== # 定时器分布式锁（Redis） # # 多 worker 部署（UVICORN_WORKERS > 1）时，每个 worker 都会启动定时器， # 导致：同一轮规则被跑 N 次、历史数据量 ×N、并发写 JSONL 造成文件损坏。 # 用 Redis SET NX 抢锁，保证同一时刻只有一个 worker 在跑。 # ====================================================================== _TIMER_LOCK_PREFIX = "meshsafe:diag:timer_lock:" @functools.lru_cache(maxsize=1) def _lock_redis(): """定时器锁专用 Redis 客户端（进程内复用，避免每次调用重建连接）。""" import redis # type: ignore from app.config import get_settings return redis.from_url( get_settings().redis_url, decode_responses=True, socket_timeout=3 ) def _acquire_timer_lock(timer_name: str, interval: int) -> bool: """尝试获取定时器锁（时间槽位设计）。 key 里带上当前时间槽位号：``meshsafe:diag:timer_lock::``， slot = 当前时间 // interval。这样无论多少 worker、谁先谁后， 同一个槽位内最多只有一个 worker 能拿到锁并执行， 下一个槽位自然放行 —— 执行频率精确等于 interval，不多不少。 Redis 不可用时返回 True（降级为单 worker 行为，不阻塞主流程）。 """ try: client = _lock_redis() slot = int(time.time() // interval) key = f"{_TIMER_LOCK_PREFIX}{timer_name}:{slot}" # EX 只做兜底（槽位 key 不会互相干扰），设 2 倍 interval 足够 return bool(client.set(key, "1", nx=True, ex=interval * 2)) except Exception as e: # noqa: BLE001 # Redis 不可用 → 不加锁（单机单 worker 场景行为不变） logger.debug("定时器锁获取失败，降级为无锁执行: %s", e) return True def _run_sync( name: str, runner_fn, docker_status: str, docker_log_tail: int, ) -> List[RuleResult]: """在线程池里同步跑规则（避免阻塞 event loop）。""" db = SessionLocal() try: # 延迟 import 避免循环依赖 from app.main import scanner, policy_store eng = DiagnosticEngine( db=db, scanner=scanner, policy_store=policy_store, docker_status=docker_status, docker_log_tail=docker_log_tail, ) results = runner_fn(eng) return results finally: db.close() def _cache_results( timer_name: str, results: List[RuleResult], docker_status: str, ) -> None: """把本轮规则结果写入内存缓存（dict）+ 持久化存储。""" now = time.time() for r in results: _DIAG_STATE["rules_cache"][r.id] = { "result": r, "ts": now, "run_by": r.run_by or f"timer:{timer_name}", } # 构建 latest 快照（合并所有缓存 + 本轮结果） all_results = [] seen = set() for rid, entry in _DIAG_STATE["rules_cache"].items(): if rid not in seen and entry.get("result"): all_results.append(entry["result"]) seen.add(rid) _DIAG_STATE["latest"] = { "rules": sorted(all_results, key=lambda x: x.id), "generated_at": now, "docker_status": docker_status, } # ---------- 持久化 ---------- try: from app.diagnostic.config import get_diag_config from app.diagnostic.storage import append_rule_results cfg = get_diag_config() if cfg.storage_enabled: append_rule_results(results, timer_name) # 可选：定时器也存 docker 日志快照（省磁盘所以用较小 tail） if cfg.log_snapshot_on_timer and docker_status != "none": from app.diagnostic.storage import snapshot_container_logs containers = [ "meshsafe-api", "meshsafe-clamav", "meshsafe-postgres", "meshsafe-redis", "meshsafe-nginx", ] # 不要每次都存 — 只在 container_check（5 分钟周期）里存一次 if timer_name == "container_check": snapshot_container_logs( containers, cfg.log_snapshot_tail, filename_prefix="timer", ) except Exception as e: # noqa: BLE001 logger.debug("诊断持久化跳过（非致命）: %s", e)