"""Background watch probe loop — 5s sampling for pinned servers.""" from __future__ import annotations import asyncio import json import logging from collections import defaultdict from datetime import datetime, timedelta, timezone from typing import Any from server.domain.models import Server from server.infrastructure.database.session import AsyncSessionLocal from server.infrastructure.database.server_repo import ServerRepositoryImpl from server.infrastructure.database.watch_repo import WatchRepositoryImpl from server.infrastructure.redis.client import get_redis from server.infrastructure.ssh.remote_probe import ssh_watch_probe, ssh_watch_processes, ssh_load_avg_probe from server.utils.watch_alerts import process_watch_probe_alerts from server.utils.watch_probe_errors import watch_probe_error_zh from server.utils.watch_state import clear_watch_redis_snapshot from server.utils.watch_metrics import ( WATCH_PROBE_INTERVAL_SEC, WATCH_PROBE_TIMEOUT_SEC, WATCH_REDIS_FRESH_SEC, WatchProbeSample, agent_watch_sample_usable, failure_sample, merge_redis_and_ssh, parse_redis_watch_payload, parse_ssh_watch_payload, fill_hardware_gaps, apply_load_avg_to_sample, sample_needs_load_supplement, ) logger = logging.getLogger("nexus.watch_probe") REDIS_HEARTBEAT_PREFIX = "heartbeat:" REDIS_WATCH_LIVE_PREFIX = "watch:live:" REDIS_WATCH_PROC_PREFIX = "watch:proc:" WATCH_LIVE_MAX_POINTS = 360 WATCH_RECORDS_RETENTION_DAYS = 7 _process_cycle: dict[int, int] = defaultdict(int) async def _supplement_load_if_needed(server: Server, sample: WatchProbeSample) -> WatchProbeSample: """Agent heartbeat omits load_avg; lightweight SSH when full probe was skipped.""" if not sample_needs_load_supplement(sample): return sample load_result = await ssh_load_avg_probe(server, timeout=5) if load_result.get("ok"): payload = load_result.get("payload") or {} apply_load_avg_to_sample(sample, payload.get("load_avg")) return sample def _utc_naive_now() -> datetime: return datetime.now(timezone.utc).replace(tzinfo=None) def _heartbeat_fresh(heartbeat: dict[str, str]) -> bool: if not heartbeat or heartbeat.get("is_online") != "True": return False ts = heartbeat.get("last_heartbeat") or "" if not ts: return False try: agent_ts = datetime.fromisoformat(ts.replace("Z", "+00:00")) if agent_ts.tzinfo is None: agent_ts = agent_ts.replace(tzinfo=timezone.utc) age = (datetime.now(timezone.utc) - agent_ts).total_seconds() return age <= WATCH_REDIS_FRESH_SEC except Exception: return False async def _append_live_point(redis, server_id: int, point: dict[str, Any]) -> None: key = f"{REDIS_WATCH_LIVE_PREFIX}{server_id}" await redis.lpush(key, json.dumps(point)) await redis.ltrim(key, 0, WATCH_LIVE_MAX_POINTS - 1) await redis.expire(key, 3600) async def _save_processes(redis, server_id: int, processes: list[dict[str, Any]] | None) -> None: if not processes: return await redis.set(f"{REDIS_WATCH_PROC_PREFIX}{server_id}", json.dumps(processes), ex=60) async def probe_server_metrics( server: Server, heartbeat: dict[str, str] | None, *, fetch_processes: bool = False, ) -> WatchProbeSample: redis_sample: WatchProbeSample | None = None if heartbeat and _heartbeat_fresh(heartbeat): redis_sample = parse_redis_watch_payload(heartbeat) if agent_watch_sample_usable(heartbeat, redis_sample) and redis_sample is not None: sample = redis_sample sample.source = "agent" sample.is_online = True if fetch_processes: procs = await ssh_watch_processes(server, timeout=WATCH_PROBE_TIMEOUT_SEC) if procs: sample.processes_json = procs fill_hardware_gaps(sample) return sample # Agent 未进 watch 模式或心跳过期 → SSH psutil（与宝塔同源） ssh_result = await ssh_watch_probe(server, timeout=WATCH_PROBE_TIMEOUT_SEC) if ssh_result.get("ok"): ssh_sample = parse_ssh_watch_payload(ssh_result["payload"]) ssh_sample.duration_ms = int(ssh_result.get("duration_ms") or 0) sample = merge_redis_and_ssh(redis_sample, ssh_sample) if redis_sample else ssh_sample elif redis_sample is not None: redis_sample.duration_ms = int(ssh_result.get("duration_ms") or 0) sample = await _supplement_load_if_needed(server, redis_sample) else: status = ssh_result.get("probe_status") or "offline" if not heartbeat and status == "offline": return failure_sample( status="no_cred", source="ssh", error=watch_probe_error_zh( ssh_result.get("error") or "无 Agent 且 SSH 探针失败", status="no_cred", ) or "无 Agent 且 SSH 探针失败", duration_ms=int(ssh_result.get("duration_ms") or 0), ) return failure_sample( status=status, source="ssh", error=watch_probe_error_zh(ssh_result.get("error") or "探针失败", status=status) or "探针失败", duration_ms=int(ssh_result.get("duration_ms") or 0), ) if fetch_processes and sample.probe_status == "ok": procs = await ssh_watch_processes(server, timeout=WATCH_PROBE_TIMEOUT_SEC) if procs: sample.processes_json = procs if sample.probe_status == "ok": fill_hardware_gaps(sample) return sample async def _probe_server_with_pins( *, watch_repo: WatchRepositoryImpl, server_repo: ServerRepositoryImpl, redis, server_id: int, pin_rows: list[dict[str, Any]], now: datetime, fetch_processes: bool, ) -> None: """Probe one server, persist Redis snapshot, WS push, and DB records for all its pins.""" server = await server_repo.get_by_id(server_id) if not server: sample = failure_sample(status="offline", source="ssh", error="服务器不存在") server_name = f"server-{server_id}" else: server_name = server.name or f"server-{server_id}" heartbeat = await redis.hgetall(f"{REDIS_HEARTBEAT_PREFIX}{server_id}") sample = await probe_server_metrics( server, heartbeat or None, fetch_processes=fetch_processes, ) live = sample.to_live_dict(server_id=server_id) live["ts"] = now.isoformat() await _append_live_point(redis, server_id, live) if sample.processes_json: await _save_processes(redis, server_id, sample.processes_json) if sample.probe_status == "ok" and server: try: await process_watch_probe_alerts( server_id=server_id, server_name=server_name, sample=sample, ) except Exception: logger.debug("watch alert processing failed", exc_info=True) admin_ids = [r["admin_id"] for r in pin_rows] try: from server.api.websocket import publish_watch_metrics await publish_watch_metrics( { "type": "watch_metrics", "server_id": server_id, "admin_ids": admin_ids, "metrics": live, } ) except Exception: logger.debug("watch WS publish failed", exc_info=True) for row in pin_rows: await watch_repo.insert_probe_record( session_id=row["session_id"], server_id=server_id, admin_id=row["admin_id"], recorded_at=now, source=sample.source, probe_status=sample.probe_status, duration_ms=sample.duration_ms, is_online=sample.is_online, cpu_pct=sample.cpu_pct, mem_pct=sample.mem_pct, disk_pct=sample.disk_pct, disk_mount=sample.disk_mount, load_1=sample.load_1, load_5=sample.load_5, load_15=sample.load_15, processes_json=sample.processes_json, error=sample.error, ) async def trigger_immediate_watch_probe(server_id: int) -> None: """Eager probe after pin add or monitoring resume — fills Redis before list_slots.""" try: async with AsyncSessionLocal() as db: watch_repo = WatchRepositoryImpl(db) pin_rows = [ p for p in await watch_repo.list_active_pins_all(monitoring_only=True) if p["server_id"] == server_id ] if not pin_rows: return server_repo = ServerRepositoryImpl(db) redis = get_redis() now = _utc_naive_now() _process_cycle[server_id] += 1 fetch_procs = _process_cycle[server_id] % 2 == 0 await _probe_server_with_pins( watch_repo=watch_repo, server_repo=server_repo, redis=redis, server_id=server_id, pin_rows=pin_rows, now=now, fetch_processes=fetch_procs, ) await db.commit() except Exception: logger.exception("immediate watch probe failed for server %s", server_id) async def _run_probe_cycle() -> None: async with AsyncSessionLocal() as db: watch_repo = WatchRepositoryImpl(db) server_repo = ServerRepositoryImpl(db) expired_servers = await watch_repo.expire_due_pins() for server_id in expired_servers: if await watch_repo.count_monitoring_pins_for_server(server_id) == 0: await clear_watch_redis_snapshot(server_id) pins = await watch_repo.list_active_pins_all(monitoring_only=True) if not pins: await db.commit() return by_server: dict[int, list[dict[str, Any]]] = defaultdict(list) for p in pins: by_server[p["server_id"]].append(p) redis = get_redis() now = _utc_naive_now() for server_id, pin_rows in by_server.items(): _process_cycle[server_id] += 1 fetch_procs = _process_cycle[server_id] % 2 == 0 await _probe_server_with_pins( watch_repo=watch_repo, server_repo=server_repo, redis=redis, server_id=server_id, pin_rows=pin_rows, now=now, fetch_processes=fetch_procs, ) await db.commit() async def _purge_old_records() -> None: cutoff = _utc_naive_now() - timedelta(days=WATCH_RECORDS_RETENTION_DAYS) async with AsyncSessionLocal() as db: repo = WatchRepositoryImpl(db) deleted = await repo.purge_records_before(cutoff) await db.commit() if deleted: logger.info("Purged %s watch_probe_records older than %s days", deleted, WATCH_RECORDS_RETENTION_DAYS) async def watch_probe_loop() -> None: """Run watch probes every 5 seconds (primary worker only).""" purge_counter = 0 while True: try: await _run_probe_cycle() purge_counter += 1 if purge_counter >= 720: # ~1h at 5s purge_counter = 0 await _purge_old_records() except asyncio.CancelledError: raise except Exception: logger.exception("watch_probe cycle failed") await asyncio.sleep(WATCH_PROBE_INTERVAL_SEC)