Codex Release Bot
319 lines
11 KiB
Python
319 lines
11 KiB
Python
"""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)
|