474 lines
16 KiB
Python
474 lines
16 KiB
Python
import asyncio
|
||
import gc
|
||
import logging
|
||
import time
|
||
from dataclasses import dataclass, field
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
DEFAULT_HEARTBEAT_INTERVAL = 30
|
||
STUCK_SESSION_TIMEOUT = 600
|
||
MAX_CONSECUTIVE_FAILURES = 5
|
||
|
||
MB = 1024 * 1024
|
||
DEFAULT_RSS_WARNING_MB = 1536
|
||
DEFAULT_RSS_CRITICAL_MB = 3072
|
||
DEFAULT_DATA_WARNING_MB = 1024
|
||
DEFAULT_DATA_CRITICAL_MB = 2048
|
||
DEFAULT_GC_OBJECTS_WARNING = 500000
|
||
DEFAULT_GC_SAMPLE_INTERVAL = 5
|
||
DEFAULT_RSS_GROWTH_WARNING_MB = 512
|
||
DEFAULT_RSS_GROWTH_CRITICAL_MB = 1024
|
||
DEFAULT_GROWTH_WINDOW_S = 600
|
||
DEFAULT_PRESSURE_COOLDOWN_S = 300
|
||
|
||
|
||
@dataclass
|
||
class MemorySnapshot:
|
||
rss_mb: float = 0.0
|
||
vms_mb: float = 0.0
|
||
data_mb: float = 0.0
|
||
gc_objects: int = 0
|
||
gc_g0_count: int = 0
|
||
gc_g1_count: int = 0
|
||
gc_g2_count: int = 0
|
||
timestamp: float = 0.0
|
||
|
||
|
||
@dataclass
|
||
class MemoryPressureEvent:
|
||
level: str # "warning" | "critical"
|
||
reason: str # "rss_threshold" | "data_threshold" | "rss_growth"
|
||
rss_mb: float = 0.0
|
||
threshold_mb: float = 0.0
|
||
rss_growth_mb: float = 0.0
|
||
window_s: float = 0.0
|
||
timestamp: float = 0.0
|
||
|
||
|
||
@dataclass
|
||
class HeartbeatReport:
|
||
active_sessions: int = 0
|
||
stuck_sessions: int = 0
|
||
stuck_details: list[dict] = field(default_factory=list)
|
||
tool_loop_sessions: int = 0
|
||
tool_loop_details: list[dict] = field(default_factory=list)
|
||
memory: MemorySnapshot = field(default_factory=MemorySnapshot)
|
||
active_tasks: int = 0
|
||
uptime_seconds: float = 0.0
|
||
process_start_timestamp: float = 0.0
|
||
timestamp: float = 0.0
|
||
warnings: list[str] = field(default_factory=list)
|
||
memory_pressure: MemoryPressureEvent | None = None
|
||
|
||
def to_dict(self) -> dict:
|
||
result = {
|
||
"active_sessions": self.active_sessions,
|
||
"stuck_sessions": self.stuck_sessions,
|
||
"stuck_details": self.stuck_details,
|
||
"tool_loop_sessions": self.tool_loop_sessions,
|
||
"tool_loop_details": self.tool_loop_details,
|
||
"memory": {
|
||
"rss_mb": self.memory.rss_mb,
|
||
"vms_mb": self.memory.vms_mb,
|
||
"data_mb": self.memory.data_mb,
|
||
"gc_objects": self.memory.gc_objects,
|
||
"gc_g0_count": self.memory.gc_g0_count,
|
||
"gc_g1_count": self.memory.gc_g1_count,
|
||
"gc_g2_count": self.memory.gc_g2_count,
|
||
"timestamp": self.memory.timestamp,
|
||
},
|
||
"active_tasks": self.active_tasks,
|
||
"uptime_seconds": self.uptime_seconds,
|
||
"process_start_timestamp": self.process_start_timestamp,
|
||
"timestamp": self.timestamp,
|
||
"warnings": self.warnings,
|
||
}
|
||
if self.memory_pressure:
|
||
result["memory_pressure"] = {
|
||
"level": self.memory_pressure.level,
|
||
"reason": self.memory_pressure.reason,
|
||
"rss_mb": self.memory_pressure.rss_mb,
|
||
"threshold_mb": self.memory_pressure.threshold_mb,
|
||
"rss_growth_mb": self.memory_pressure.rss_growth_mb,
|
||
"window_s": self.memory_pressure.window_s,
|
||
"timestamp": self.memory_pressure.timestamp,
|
||
}
|
||
return result
|
||
|
||
|
||
class DiagnosticHeartbeat:
|
||
"""诊断心跳：30s 间隔检测卡住 session、内存快照 + 内存压力。
|
||
|
||
集成到 FastAPI lifespan 后启动，可被 health endpoint 消费。
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
interval: int = DEFAULT_HEARTBEAT_INTERVAL,
|
||
stuck_timeout: int = STUCK_SESSION_TIMEOUT,
|
||
rss_warning_mb: float = DEFAULT_RSS_WARNING_MB,
|
||
rss_critical_mb: float = DEFAULT_RSS_CRITICAL_MB,
|
||
data_warning_mb: float = DEFAULT_DATA_WARNING_MB,
|
||
data_critical_mb: float = DEFAULT_DATA_CRITICAL_MB,
|
||
rss_growth_warning_mb: float = DEFAULT_RSS_GROWTH_WARNING_MB,
|
||
rss_growth_critical_mb: float = DEFAULT_RSS_GROWTH_CRITICAL_MB,
|
||
growth_window_s: float = DEFAULT_GROWTH_WINDOW_S,
|
||
pressure_cooldown_s: float = DEFAULT_PRESSURE_COOLDOWN_S,
|
||
gc_objects_warning: int = DEFAULT_GC_OBJECTS_WARNING,
|
||
gc_sample_interval: int = DEFAULT_GC_SAMPLE_INTERVAL,
|
||
):
|
||
self._interval = interval
|
||
self._stuck_timeout = stuck_timeout
|
||
self._rss_warning_mb = rss_warning_mb
|
||
self._rss_critical_mb = rss_critical_mb
|
||
self._data_warning_mb = data_warning_mb
|
||
self._data_critical_mb = data_critical_mb
|
||
self._rss_growth_warning_mb = rss_growth_warning_mb
|
||
self._rss_growth_critical_mb = rss_growth_critical_mb
|
||
self._growth_window_s = growth_window_s
|
||
self._pressure_cooldown_s = pressure_cooldown_s
|
||
self._gc_objects_warning = gc_objects_warning
|
||
self._gc_sample_interval = gc_sample_interval
|
||
self._running = False
|
||
self._task: asyncio.Task | None = None
|
||
self._last_report: HeartbeatReport | None = None
|
||
self._last_memory: MemorySnapshot | None = None
|
||
self._last_pressure_at: dict[str, float] = {}
|
||
self._start_time: float | None = None
|
||
self._start_monotonic: float | None = None
|
||
self._consecutive_failures = 0
|
||
self._cycle_count = 0
|
||
|
||
@property
|
||
def last_report(self) -> HeartbeatReport | None:
|
||
return self._last_report
|
||
|
||
@staticmethod
|
||
def _resolve_start_time() -> float:
|
||
try:
|
||
import psutil
|
||
|
||
return psutil.Process().create_time()
|
||
except Exception:
|
||
return time.time()
|
||
|
||
async def start(self):
|
||
self._running = True
|
||
self._start_time = self._resolve_start_time()
|
||
self._start_monotonic = time.monotonic()
|
||
self._task = asyncio.create_task(self._loop())
|
||
logger.info(
|
||
"DiagnosticHeartbeat started (interval=%ds, stuck_timeout=%ds)",
|
||
self._interval,
|
||
self._stuck_timeout,
|
||
)
|
||
|
||
async def stop(self):
|
||
self._running = False
|
||
if self._task:
|
||
self._task.cancel()
|
||
try:
|
||
await self._task
|
||
except asyncio.CancelledError:
|
||
pass
|
||
self._task = None
|
||
logger.info("DiagnosticHeartbeat stopped")
|
||
|
||
async def _loop(self):
|
||
while self._running:
|
||
await asyncio.sleep(self._interval)
|
||
try:
|
||
self._last_report = await self._collect()
|
||
self._consecutive_failures = 0
|
||
await self._push_metrics(self._last_report)
|
||
if self._last_report.warnings:
|
||
for w in self._last_report.warnings:
|
||
logger.warning("DiagnosticHeartbeat: %s", w)
|
||
if self._last_report.memory_pressure:
|
||
mp = self._last_report.memory_pressure
|
||
logger.warning(
|
||
"DiagnosticHeartbeat: 内存压力 %s (reason=%s, rss=%.0fMB, threshold=%.0fMB)",
|
||
mp.level,
|
||
mp.reason,
|
||
mp.rss_mb,
|
||
mp.threshold_mb,
|
||
)
|
||
except Exception:
|
||
self._consecutive_failures += 1
|
||
logger.exception(
|
||
"DiagnosticHeartbeat: collect failed (consecutive=%d/%d)",
|
||
self._consecutive_failures,
|
||
MAX_CONSECUTIVE_FAILURES,
|
||
)
|
||
if self._consecutive_failures >= MAX_CONSECUTIVE_FAILURES:
|
||
logger.critical(
|
||
"DiagnosticHeartbeat: %d consecutive collect failures, heartbeat may be broken",
|
||
self._consecutive_failures,
|
||
)
|
||
|
||
async def _collect(self) -> HeartbeatReport:
|
||
now = time.monotonic()
|
||
self._cycle_count += 1
|
||
memory = await self._capture_memory(now)
|
||
self._last_memory = memory
|
||
|
||
snapshots = _get_all_snapshots()
|
||
stuck_details, stuck_count = await self._detect_stuck_sessions(now, snapshots)
|
||
active_tasks = len(asyncio.all_tasks())
|
||
active_sessions = _count_active_sessions_from(snapshots)
|
||
|
||
uptime_seconds = now - self._start_monotonic if self._start_monotonic else 0.0
|
||
process_start_timestamp = self._start_time or 0.0
|
||
|
||
memory_pressure = self._detect_memory_pressure(memory, now)
|
||
|
||
warnings: list[str] = []
|
||
if stuck_count > 0:
|
||
warnings.append(f"检测到 {stuck_count} 个卡住 session")
|
||
if memory.rss_mb > self._rss_warning_mb:
|
||
warnings.append(f"内存 RSS {memory.rss_mb:.0f}MB > {self._rss_warning_mb:.0f}MB")
|
||
if memory.data_mb > self._data_warning_mb:
|
||
warnings.append(f"内存 DATA {memory.data_mb:.0f}MB > {self._data_warning_mb:.0f}MB")
|
||
if memory.gc_objects > self._gc_objects_warning:
|
||
warnings.append(f"GC 对象数 {memory.gc_objects} > {self._gc_objects_warning}")
|
||
if memory_pressure:
|
||
level_label = "严重" if memory_pressure.level == "critical" else ""
|
||
warnings.append(f"内存压力{level_label}: {memory_pressure.reason} (RSS={memory_pressure.rss_mb:.0f}MB)")
|
||
|
||
return HeartbeatReport(
|
||
active_sessions=active_sessions,
|
||
stuck_sessions=stuck_count,
|
||
stuck_details=stuck_details,
|
||
tool_loop_sessions=0,
|
||
tool_loop_details=[],
|
||
memory=memory,
|
||
active_tasks=active_tasks,
|
||
uptime_seconds=uptime_seconds,
|
||
process_start_timestamp=process_start_timestamp,
|
||
timestamp=now,
|
||
warnings=warnings,
|
||
memory_pressure=memory_pressure,
|
||
)
|
||
|
||
async def _push_metrics(self, report: HeartbeatReport):
|
||
try:
|
||
from yuxi.channel.monitoring.metrics import (
|
||
active_sessions,
|
||
active_tasks,
|
||
memory_gc_objects,
|
||
memory_rss_mb,
|
||
stuck_sessions,
|
||
)
|
||
|
||
active_sessions.set(report.active_sessions)
|
||
stuck_sessions.set(report.stuck_sessions)
|
||
memory_rss_mb.set(report.memory.rss_mb)
|
||
memory_gc_objects.set(report.memory.gc_objects)
|
||
active_tasks.set(report.active_tasks)
|
||
except Exception:
|
||
pass
|
||
|
||
def _detect_memory_pressure(self, current: MemorySnapshot, now: float) -> MemoryPressureEvent | None:
|
||
pressure = _check_threshold_pressure(
|
||
current,
|
||
rss_warning_mb=self._rss_warning_mb,
|
||
rss_critical_mb=self._rss_critical_mb,
|
||
)
|
||
if pressure is None:
|
||
pressure = _check_data_threshold_pressure(
|
||
current,
|
||
data_warning_mb=self._data_warning_mb,
|
||
data_critical_mb=self._data_critical_mb,
|
||
)
|
||
if pressure is None:
|
||
pressure = _check_growth_pressure(
|
||
self._last_memory,
|
||
current,
|
||
now,
|
||
growth_window_s=self._growth_window_s,
|
||
rss_growth_warning_mb=self._rss_growth_warning_mb,
|
||
rss_growth_critical_mb=self._rss_growth_critical_mb,
|
||
)
|
||
|
||
if pressure is None:
|
||
return None
|
||
|
||
key = pressure.reason
|
||
last_at = self._last_pressure_at.get(key, 0)
|
||
if now - last_at < self._pressure_cooldown_s:
|
||
return None
|
||
|
||
self._last_pressure_at[key] = now
|
||
return pressure
|
||
|
||
async def _detect_stuck_sessions(self, now: float, snapshots: dict) -> tuple[list[dict], int]:
|
||
stuck: list[dict] = []
|
||
for key, snap in snapshots.items():
|
||
if snap.last_event_at and (now - snap.last_event_at) > self._stuck_timeout:
|
||
state_value = _get_state_value(snap.state)
|
||
if snap.connected and state_value in ("running", "retrying"):
|
||
stuck.append(
|
||
{
|
||
"key": key,
|
||
"state": state_value,
|
||
"last_event_at": snap.last_event_at,
|
||
"idle_seconds": round(now - snap.last_event_at, 1),
|
||
}
|
||
)
|
||
return stuck, len(stuck)
|
||
|
||
async def _capture_memory(self, now: float) -> MemorySnapshot:
|
||
snapshot = MemorySnapshot(timestamp=now)
|
||
try:
|
||
import psutil
|
||
|
||
proc = psutil.Process()
|
||
mem = proc.memory_info()
|
||
snapshot.rss_mb = round(mem.rss / 1024 / 1024, 2)
|
||
snapshot.vms_mb = round(mem.vms / 1024 / 1024, 2)
|
||
snapshot.data_mb = round(getattr(mem, "data", 0) / 1024 / 1024, 2)
|
||
except Exception:
|
||
pass
|
||
|
||
gc_counts = gc.get_count()
|
||
snapshot.gc_g0_count = gc_counts[0]
|
||
snapshot.gc_g1_count = gc_counts[1]
|
||
snapshot.gc_g2_count = gc_counts[2]
|
||
|
||
if self._cycle_count % self._gc_sample_interval == 0:
|
||
snapshot.gc_objects = len(gc.get_objects())
|
||
|
||
return snapshot
|
||
|
||
def get_memory_trend(self) -> dict:
|
||
if self._last_memory:
|
||
return {
|
||
"rss_mb": self._last_memory.rss_mb,
|
||
"vms_mb": self._last_memory.vms_mb,
|
||
"data_mb": self._last_memory.data_mb,
|
||
"gc_objects": self._last_memory.gc_objects,
|
||
"timestamp": self._last_memory.timestamp,
|
||
}
|
||
return {}
|
||
|
||
def reset(self):
|
||
self._last_memory = None
|
||
self._last_pressure_at.clear()
|
||
self._last_report = None
|
||
self._cycle_count = 0
|
||
|
||
|
||
def _check_threshold_pressure(
|
||
memory: MemorySnapshot,
|
||
*,
|
||
rss_warning_mb: float,
|
||
rss_critical_mb: float,
|
||
) -> MemoryPressureEvent | None:
|
||
if memory.rss_mb >= rss_critical_mb:
|
||
return MemoryPressureEvent(
|
||
level="critical",
|
||
reason="rss_threshold",
|
||
rss_mb=memory.rss_mb,
|
||
threshold_mb=rss_critical_mb,
|
||
timestamp=memory.timestamp,
|
||
)
|
||
if memory.rss_mb >= rss_warning_mb:
|
||
return MemoryPressureEvent(
|
||
level="warning",
|
||
reason="rss_threshold",
|
||
rss_mb=memory.rss_mb,
|
||
threshold_mb=rss_warning_mb,
|
||
timestamp=memory.timestamp,
|
||
)
|
||
return None
|
||
|
||
|
||
def _check_data_threshold_pressure(
|
||
memory: MemorySnapshot,
|
||
*,
|
||
data_warning_mb: float,
|
||
data_critical_mb: float,
|
||
) -> MemoryPressureEvent | None:
|
||
if memory.data_mb >= data_critical_mb:
|
||
return MemoryPressureEvent(
|
||
level="critical",
|
||
reason="data_threshold",
|
||
rss_mb=memory.rss_mb,
|
||
threshold_mb=data_critical_mb,
|
||
timestamp=memory.timestamp,
|
||
)
|
||
if memory.data_mb >= data_warning_mb:
|
||
return MemoryPressureEvent(
|
||
level="warning",
|
||
reason="data_threshold",
|
||
rss_mb=memory.rss_mb,
|
||
threshold_mb=data_warning_mb,
|
||
timestamp=memory.timestamp,
|
||
)
|
||
return None
|
||
|
||
|
||
def _check_growth_pressure(
|
||
previous: MemorySnapshot | None,
|
||
current: MemorySnapshot,
|
||
now: float,
|
||
*,
|
||
growth_window_s: float,
|
||
rss_growth_warning_mb: float,
|
||
rss_growth_critical_mb: float,
|
||
) -> MemoryPressureEvent | None:
|
||
if previous is None or previous.rss_mb <= 0:
|
||
return None
|
||
|
||
window_s = now - previous.timestamp
|
||
if window_s <= 0 or window_s > growth_window_s:
|
||
return None
|
||
|
||
rss_growth_mb = current.rss_mb - previous.rss_mb
|
||
if rss_growth_mb <= 0:
|
||
return None
|
||
|
||
if rss_growth_mb >= rss_growth_critical_mb:
|
||
return MemoryPressureEvent(
|
||
level="critical",
|
||
reason="rss_growth",
|
||
rss_mb=current.rss_mb,
|
||
threshold_mb=rss_growth_critical_mb,
|
||
rss_growth_mb=round(rss_growth_mb, 2),
|
||
window_s=round(window_s, 1),
|
||
timestamp=current.timestamp,
|
||
)
|
||
if rss_growth_mb >= rss_growth_warning_mb:
|
||
return MemoryPressureEvent(
|
||
level="warning",
|
||
reason="rss_growth",
|
||
rss_mb=current.rss_mb,
|
||
threshold_mb=rss_growth_warning_mb,
|
||
rss_growth_mb=round(rss_growth_mb, 2),
|
||
window_s=round(window_s, 1),
|
||
timestamp=current.timestamp,
|
||
)
|
||
return None
|
||
|
||
|
||
def _get_state_value(state) -> str:
|
||
return state.value if hasattr(state, "value") else str(state)
|
||
|
||
|
||
def _get_all_snapshots() -> dict:
|
||
try:
|
||
from yuxi.channel.runtime.manager import gateway # 延迟导入以避免循环依赖；Python import 缓存保证无重复加载开销
|
||
|
||
return gateway.get_all_snapshots()
|
||
except Exception:
|
||
logger.warning("DiagnosticHeartbeat: get_all_snapshots failed", exc_info=True)
|
||
return {}
|
||
|
||
|
||
def _count_active_sessions_from(snapshots: dict) -> int:
|
||
return sum(1 for s in snapshots.values() if _get_state_value(s.state) == "running")
|
||
|
||
|
||
diagnostic_heartbeat = DiagnosticHeartbeat()
|