use std::collections::{HashMap, VecDeque};
use std::net::{SocketAddr, TcpStream};
use std::sync::{Arc, Mutex};
use arc_swap::ArcSwap;
use std::time::{Duration, Instant};
use crate::cluster::{ClusterState, NodeStatus};
use crate::host::{AccelerationMode, RuntimeProfile};
use crate::system_profile::SystemProfile;
use crate::watcher::ProfileChange;
#[derive(Clone, Copy, Debug)]
pub struct HealthConfig {
pub check_interval: Duration,
pub timeout: Duration,
pub min_successes: usize,
pub max_failures: usize,
pub healthy_latency_us: f32,
pub degraded_latency_us: f32,
pub healthy_delivery_ratio: f32,
pub degraded_delivery_ratio: f32,
}
impl Default for HealthConfig {
fn default() -> Self {
Self {
check_interval: Duration::from_secs(5),
timeout: Duration::from_secs(3),
min_successes: 2,
max_failures: 3,
healthy_latency_us: 10.0,
degraded_latency_us: 20.0,
healthy_delivery_ratio: 0.95,
degraded_delivery_ratio: 0.80,
}
}
}
impl HealthConfig {
pub fn autotuned(profile: &RuntimeProfile) -> Self {
let (healthy_latency_us, degraded_latency_us, timeout, max_failures) =
match profile.acceleration_mode {
AccelerationMode::Gpu => (6.0, 14.0, Duration::from_secs(2), 4),
AccelerationMode::Avx512 => (8.0, 18.0, Duration::from_secs(2), 3),
_ => (10.0, 20.0, Duration::from_secs(3), 3),
};
Self {
check_interval: Duration::from_secs(5),
timeout,
min_successes: if profile.recommended_workers >= 8 {
3
} else {
2
},
max_failures,
healthy_latency_us,
degraded_latency_us,
healthy_delivery_ratio: if profile.node_resources.vram_gb > 0.0 {
0.97
} else {
0.95
},
degraded_delivery_ratio: if profile.node_resources.vram_gb > 0.0 {
0.88
} else {
0.80
},
}
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum HealthStatus {
Healthy,
Degraded,
Failed,
#[default]
Unknown,
}
#[derive(Clone, Debug)]
pub struct HealthCheckResult {
pub node_id: String,
pub latency_us: f32,
pub delivery_ratio: f32,
pub status: HealthStatus,
pub timestamp: Instant,
}
#[derive(Clone, Debug)]
pub struct NetworkHealthMonitor {
cluster: Arc<ClusterState>,
config: Arc<ArcSwap<HealthConfig>>,
last_check: Arc<Mutex<Option<Instant>>>,
recent_checks: Arc<Mutex<HashMap<String, VecDeque<HealthCheckResult>>>>,
tcp_probe_targets: Arc<Mutex<HashMap<String, SocketAddr>>>,
}
impl NetworkHealthMonitor {
pub fn new(cluster: Arc<ClusterState>, config: HealthConfig) -> Self {
Self {
cluster,
config: Arc::new(ArcSwap::from_pointee(config)),
last_check: Arc::new(Mutex::new(None)),
recent_checks: Arc::new(Mutex::new(HashMap::new())),
tcp_probe_targets: Arc::new(Mutex::new(HashMap::new())),
}
}
pub fn with_runtime_profile(cluster: Arc<ClusterState>, profile: &RuntimeProfile) -> Self {
Self::new(cluster, HealthConfig::autotuned(profile))
}
pub fn config(&self) -> HealthConfig {
**self.config.load()
}
pub fn register_tcp_probe_target(&self, node_id: impl Into<String>, target: SocketAddr) {
self.tcp_probe_targets
.lock()
.unwrap_or_else(|poison| poison.into_inner())
.insert(node_id.into(), target);
}
pub fn clear_tcp_probe_target(&self, node_id: &str) -> bool {
self.tcp_probe_targets
.lock()
.unwrap_or_else(|poison| poison.into_inner())
.remove(node_id)
.is_some()
}
pub fn check_all(&self) {
let cfg = self.config();
let now = Instant::now();
let nodes_snapshot = self.cluster.nodes_snapshot();
let targets: Vec<(&str, SocketAddr)> = {
let probe_targets = self
.tcp_probe_targets
.lock()
.unwrap_or_else(|poison| poison.into_inner());
if probe_targets.is_empty() {
Vec::new()
} else {
nodes_snapshot
.iter()
.filter_map(|node| {
probe_targets
.get(&node.id)
.map(|addr| (node.id.as_str(), *addr))
})
.collect()
}
};
let mut probe_results = HashMap::with_capacity(targets.len());
for (node_id, addr) in targets {
let ok = TcpStream::connect_timeout(&addr, cfg.timeout).is_ok();
probe_results.insert(node_id, ok);
}
let mut metrics_guard = self
.cluster
.metrics
.lock()
.unwrap_or_else(|poison| poison.into_inner());
let mut checks = self
.recent_checks
.lock()
.unwrap_or_else(|poison| poison.into_inner());
let has_probe_results = !probe_results.is_empty();
for node in nodes_snapshot.iter() {
let tcp_probe_ok = if has_probe_results {
probe_results.get(node.id.as_str()).copied()
} else {
None
};
let node_checks = if let Some(node_checks) = checks.get_mut(&node.id) {
node_checks
} else {
checks.entry(node.id.clone()).or_default()
};
let mut result = if node_checks.len() >= 10 {
node_checks.pop_front().unwrap()
} else {
HealthCheckResult {
node_id: node.id.clone(),
latency_us: 0.0,
delivery_ratio: 1.0,
status: HealthStatus::Unknown,
timestamp: now,
}
};
if let Some(m) = metrics_guard.get_mut(&node.id) {
let timeout =
now.saturating_duration_since(m.last_heartbeat) >= m.heartbeat_timeout;
let latency_us = if m.latency_samples > 0 {
m.avg_latency_us
} else {
0.0
};
let delivery_ratio = if m.delivery_ratio_initialized {
m.delivery_ratio
} else {
1.0
};
let status =
if tcp_probe_ok == Some(false) || timeout || m.status == NodeStatus::Failed {
HealthStatus::Failed
} else if m.latency_samples == 0 && !m.delivery_ratio_initialized {
HealthStatus::Unknown
} else {
get_health_status_with_config(&cfg, latency_us, delivery_ratio)
};
let measured_latency_us = if tcp_probe_ok == Some(false) {
cfg.timeout.as_secs_f32() * 1_000_000.0
} else {
latency_us
};
m.status = match status {
HealthStatus::Healthy => NodeStatus::Active,
HealthStatus::Degraded => NodeStatus::Degraded,
HealthStatus::Failed => NodeStatus::Failed,
HealthStatus::Unknown => m.status,
};
result.latency_us = measured_latency_us;
result.delivery_ratio = delivery_ratio;
result.status = status;
result.timestamp = now;
} else {
result.latency_us = 0.0;
result.delivery_ratio = 1.0;
result.status = HealthStatus::Unknown;
result.timestamp = now;
}
node_checks.push_back(result);
}
*self
.last_check
.lock()
.unwrap_or_else(|poison| poison.into_inner()) = Some(now);
}
#[allow(dead_code)]
pub(crate) fn get_health_status(&self, latency_us: f32, delivery_ratio: f32) -> HealthStatus {
let cfg = self.config();
get_health_status_with_config(&cfg, latency_us, delivery_ratio)
}
pub fn get_node_health(&self, node_id: &str) -> Option<HealthCheckResult> {
let checks = self
.recent_checks
.lock()
.unwrap_or_else(|poison| poison.into_inner());
checks
.get(node_id)
.and_then(|checks| checks.back())
.cloned()
}
pub fn get_all_health(&self) -> Vec<HealthCheckResult> {
let checks = self
.recent_checks
.lock()
.unwrap_or_else(|poison| poison.into_inner());
let total_capacity: usize = checks.values().map(|c| c.len()).sum();
let mut results = Vec::with_capacity(total_capacity);
for node_checks in checks.values() {
results.extend(node_checks.iter().cloned());
}
results
}
pub fn healthy_node_count(&self) -> usize {
let checks = self
.recent_checks
.lock()
.unwrap_or_else(|poison| poison.into_inner());
checks
.values()
.filter(|node_checks| {
node_checks
.back()
.map(|c| c.status == HealthStatus::Healthy)
.unwrap_or(false)
})
.count()
}
pub fn needs_quantization_fallback(&self, node_id: &str) -> bool {
let cfg = self.config();
let checks = self
.recent_checks
.lock()
.unwrap_or_else(|poison| poison.into_inner());
if let Some(node_checks) = checks.get(node_id) {
if node_checks.len() < 3 {
return false;
}
let mut sum_delivery_ratio = 0.0f32;
let mut sum_latency_us = 0.0f32;
for r in node_checks.iter().rev().take(3) {
sum_delivery_ratio += r.delivery_ratio;
sum_latency_us += r.latency_us;
}
let avg_delivery_ratio = sum_delivery_ratio / 3.0;
let avg_latency_us = sum_latency_us / 3.0;
avg_delivery_ratio < cfg.healthy_delivery_ratio
|| avg_latency_us > cfg.healthy_latency_us
} else {
false
}
}
pub fn get_health_summary(&self) -> String {
let checks = self
.recent_checks
.lock()
.unwrap_or_else(|poison| poison.into_inner());
let total_nodes = self.cluster.node_count();
let mut healthy_count = 0usize;
let mut degraded_count = 0usize;
let mut failed_count = 0usize;
for node_checks in checks.values() {
if let Some(c) = node_checks.back() {
match c.status {
HealthStatus::Healthy => healthy_count += 1,
HealthStatus::Degraded => degraded_count += 1,
HealthStatus::Failed => failed_count += 1,
HealthStatus::Unknown => {}
}
}
}
format!(
"Network Health Summary\n\
=================\n\
Total nodes: {total_nodes}\n\
Healthy: {healthy_count}\n\
Degraded: {degraded_count}\n\
Failed: {failed_count}\n"
)
}
pub fn start_periodic_checks(&self) {
let check_interval = self.config().check_interval;
let this = self.clone();
std::thread::spawn(move || loop {
this.check_all();
std::thread::sleep(check_interval);
});
}
pub fn reconfigure_from_system_profile(&self, profile: &SystemProfile) {
let rp: RuntimeProfile = profile.into();
self.config.store(Arc::new(HealthConfig::autotuned(&rp)));
}
pub fn subscribe_to_watcher(
self: &Arc<Self>,
watcher: &crate::watcher::SystemProfileWatcher,
) -> tokio::task::JoinHandle<()> {
let this = self.clone();
let mut rx = watcher.subscribe();
tokio::spawn(async move {
while let Ok(event) = rx.recv().await {
if let ProfileChange::Updated(profile) = event {
this.reconfigure_from_system_profile(&profile);
tracing::info!(
"health monitor re-tuned from profile: {} workers, {:?} accel",
profile.recommended_workers,
profile.acceleration_mode,
);
}
}
})
}
}
#[inline]
fn get_health_status_with_config(
cfg: &HealthConfig,
latency_us: f32,
delivery_ratio: f32,
) -> HealthStatus {
if delivery_ratio >= cfg.healthy_delivery_ratio && latency_us <= cfg.healthy_latency_us {
HealthStatus::Healthy
} else if delivery_ratio >= cfg.degraded_delivery_ratio && latency_us <= cfg.degraded_latency_us
{
HealthStatus::Degraded
} else {
HealthStatus::Failed
}
}
#[derive(Clone, Debug, Default)]
pub struct HealthMetrics {
pub avg_latency_us: f32,
pub min_latency_us: f32,
pub max_latency_us: f32,
pub avg_delivery_ratio: f32,
pub min_delivery_ratio: f32,
pub sample_count: usize,
}
impl HealthMetrics {
pub fn new() -> Self {
Self::default()
}
pub fn record(&mut self, latency_us: f32, delivery_ratio: f32) {
if self.sample_count == 0 {
self.min_latency_us = latency_us;
self.max_latency_us = latency_us;
self.min_delivery_ratio = delivery_ratio;
} else {
if latency_us < self.min_latency_us {
self.min_latency_us = latency_us;
}
if latency_us > self.max_latency_us {
self.max_latency_us = latency_us;
}
if delivery_ratio < self.min_delivery_ratio {
self.min_delivery_ratio = delivery_ratio;
}
}
let n = self.sample_count as f32;
self.avg_latency_us = (self.avg_latency_us * n + latency_us) / (n + 1.0);
self.avg_delivery_ratio = (self.avg_delivery_ratio * n + delivery_ratio) / (n + 1.0);
self.sample_count += 1;
}
pub fn get_recommendation(&self) -> String {
if self.sample_count == 0 {
return "No data available".to_string();
}
match (self.avg_delivery_ratio, self.avg_latency_us) {
(ratio, latency) if ratio >= 0.95 && latency <= 10.0 => {
"Cluster is healthy. No action needed.".to_string()
}
(ratio, latency) if ratio >= 0.90 && latency <= 20.0 => {
"Cluster is degraded. Consider load balancing.".to_string()
}
(ratio, latency) if ratio >= 0.80 && latency <= 50.0 => {
"Cluster performance is poor. Recommend quantization fallback.".to_string()
}
_ => "Cluster has failed nodes. Investigate and recover.".to_string(),
}
}
}
#[derive(Clone, Debug)]
pub struct FaultDetector {
cluster: Arc<ClusterState>,
detection_interval: Duration,
}
impl FaultDetector {
pub fn new(cluster: Arc<ClusterState>, detection_interval: Duration) -> Self {
Self {
cluster,
detection_interval,
}
}
pub fn detect_failures(&self) -> Vec<String> {
let now = Instant::now();
let mut failed_nodes = Vec::new();
let mut metrics = self
.cluster
.metrics
.lock()
.unwrap_or_else(|poison| poison.into_inner());
let total_nodes = self.cluster.node_count();
if metrics.len() == total_nodes {
for m in metrics.values_mut() {
if now.saturating_duration_since(m.last_heartbeat) >= m.heartbeat_timeout
&& m.status != NodeStatus::Failed
{
m.status = NodeStatus::Failed;
failed_nodes.push(m.name.clone());
}
}
} else {
let nodes_snapshot = self.cluster.nodes_snapshot();
for node in nodes_snapshot.iter() {
if let Some(m) = metrics.get_mut(&node.id) {
if now.saturating_duration_since(m.last_heartbeat) >= m.heartbeat_timeout
&& m.status != NodeStatus::Failed
{
m.status = NodeStatus::Failed;
failed_nodes.push(node.id.clone());
}
} else {
let mut m = crate::cluster::NodeMetrics::new(
node.vram_gb,
node.system_memory_gb,
node.compute_capability.clone(),
Duration::from_secs(5),
);
m.name = node.id.clone();
m.gpu_name = node.gpu_name.clone();
m.status = NodeStatus::Failed;
metrics.insert(node.id.clone(), m);
failed_nodes.push(node.id.clone());
}
}
}
failed_nodes
}
pub fn recover_node(&self, node_id: &str) {
self.cluster.get_metrics_mut(node_id, |metrics| {
metrics.status = NodeStatus::Active;
metrics.last_heartbeat = Instant::now();
metrics.avg_latency_us = 0.0;
metrics.min_latency_us = f32::MAX;
metrics.max_latency_us = 0.0;
metrics.latency_samples = 0;
metrics.delivery_ratio = 1.0;
metrics.throughput_gbps = 0.0;
});
}
pub fn start_periodic_detection(&self) {
let this = self.clone();
std::thread::spawn(move || loop {
let failed = this.detect_failures();
if !failed.is_empty() {
tracing::warn!(
"Detected {} failed nodes: {:?}",
failed.len(),
failed.iter().map(|s| s.as_str()).collect::<Vec<_>>()
);
}
std::thread::sleep(this.detection_interval);
});
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::cluster::ClusterState;
use crate::protocol::NodeResources;
use std::net::TcpListener;
use std::sync::Arc;
use std::time::Duration;
#[test]
fn health_monitor_checks_all_nodes() {
let cluster = Arc::new(ClusterState::new());
cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
cluster.register(NodeResources::new("node-b", 12.0, 32.0, "8.6", None));
let monitor = NetworkHealthMonitor::new(cluster.clone(), HealthConfig::default());
monitor.check_all();
let summary = monitor.get_health_summary();
assert!(summary.contains("Total nodes: 2"));
}
#[test]
fn health_monitor_marks_fresh_nodes_unknown_without_samples() {
let cluster = Arc::new(ClusterState::new());
cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
let monitor = NetworkHealthMonitor::new(cluster, HealthConfig::default());
monitor.check_all();
let node_health = monitor
.get_node_health("node-a")
.expect("expected health check result for node-a");
assert_eq!(node_health.status, HealthStatus::Unknown);
}
#[test]
fn health_monitor_detects_degraded_nodes() {
let cluster = Arc::new(ClusterState::new());
cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
let monitor = NetworkHealthMonitor::new(cluster.clone(), HealthConfig::default());
for _ in 0..3 {
let result = HealthCheckResult {
node_id: "node-a".to_string(),
latency_us: 15.0,
delivery_ratio: 0.85,
status: HealthStatus::Degraded,
timestamp: std::time::Instant::now(),
};
monitor
.recent_checks
.lock()
.unwrap()
.entry("node-a".to_string())
.or_default()
.push_back(result);
}
assert!(monitor.needs_quantization_fallback("node-a"));
}
#[test]
fn health_metrics_aggregates() {
let mut metrics = HealthMetrics::new();
metrics.record(1.0, 0.98);
metrics.record(2.0, 0.95);
metrics.record(1.5, 0.97);
assert!((metrics.avg_latency_us - 1.5).abs() < 0.1);
assert!(metrics.min_latency_us <= 1.0);
assert!(metrics.max_latency_us >= 2.0);
}
#[test]
fn fault_detector_detects_failures() {
let cluster = Arc::new(ClusterState::new());
cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
std::thread::sleep(Duration::from_secs(6));
let detector = FaultDetector::new(cluster.clone(), Duration::from_secs(1));
let failed = detector.detect_failures();
assert!(failed.contains(&"node-a".to_string()));
}
#[test]
fn fault_detector_detects_failures_for_degraded_nodes() {
let cluster = Arc::new(ClusterState::new());
cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
cluster.get_metrics_mut("node-a", |metrics| {
metrics.status = NodeStatus::Degraded;
});
std::thread::sleep(Duration::from_secs(6));
let detector = FaultDetector::new(cluster.clone(), Duration::from_secs(1));
let failed = detector.detect_failures();
assert!(failed.contains(&"node-a".to_string()));
assert_eq!(
cluster.get_metrics("node-a").unwrap().status,
NodeStatus::Failed
);
}
#[test]
fn fault_detector_recovers_nodes() {
let cluster = Arc::new(ClusterState::new());
cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
std::thread::sleep(Duration::from_secs(6));
let detector = FaultDetector::new(cluster.clone(), Duration::from_secs(1));
detector.detect_failures();
assert!(cluster.get_metrics("node-a").unwrap().status == NodeStatus::Failed);
detector.recover_node("node-a");
assert_eq!(
cluster.get_metrics("node-a").unwrap().status,
NodeStatus::Active
);
}
#[test]
fn health_monitor_gets_recommendation() {
let mut metrics = HealthMetrics::new();
metrics.record(1.0, 0.98);
metrics.record(2.0, 0.95);
let recommendation = metrics.get_recommendation();
assert!(recommendation.contains("healthy"));
}
#[test]
fn get_recommendation_uses_average_latency_not_minimum() {
let mut metrics = HealthMetrics::new();
metrics.record(1.0, 0.98);
metrics.record(200.0, 0.98);
let recommendation = metrics.get_recommendation();
assert!(!recommendation.contains("healthy"));
}
#[test]
fn health_config_autotunes_for_gpu_profiles() {
let profile = RuntimeProfile {
node_resources: NodeResources::new("node-a", 24.0, 64.0, "8.9", None),
logical_cores: 16,
recommended_workers: 8,
acceleration_mode: AccelerationMode::Gpu,
gpu_backend: crate::host::GpuBackend::Cuda,
xdp_supported: true,
detection_source: String::from("test"),
probe_mode: crate::host::ProbeMode::Fast,
};
let config = HealthConfig::autotuned(&profile);
assert!(config.healthy_latency_us < 10.0);
assert!(config.healthy_delivery_ratio > 0.95);
}
#[test]
fn get_health_status_fails_on_either_metric_crossing_floor() {
let cluster = Arc::new(ClusterState::new());
let monitor = NetworkHealthMonitor::new(cluster, HealthConfig::default());
let cfg = HealthConfig::default();
assert_eq!(
monitor.get_health_status(cfg.degraded_latency_us * 10.0, 0.99),
HealthStatus::Failed
);
assert_eq!(
monitor.get_health_status(1.0, cfg.degraded_delivery_ratio / 2.0),
HealthStatus::Failed
);
assert_eq!(
monitor.get_health_status(cfg.degraded_latency_us, cfg.degraded_delivery_ratio),
HealthStatus::Degraded
);
}
#[test]
fn tcp_probe_target_can_mark_node_failed_when_unreachable() {
let cluster = Arc::new(ClusterState::new());
cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
let listener = TcpListener::bind("127.0.0.1:0").expect("bind probe listener");
let addr = listener.local_addr().expect("listener local addr");
drop(listener);
let monitor = NetworkHealthMonitor::new(cluster.clone(), HealthConfig::default());
monitor.register_tcp_probe_target("node-a", addr);
monitor.check_all();
let node_health = monitor
.get_node_health("node-a")
.expect("expected health check result for node-a");
assert_eq!(node_health.status, HealthStatus::Failed);
}
#[test]
fn tcp_probe_target_healthy_when_reachable() {
let cluster = Arc::new(ClusterState::new());
cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
cluster.get_metrics_mut("node-a", |metrics| {
metrics.record_latency(2.0);
metrics.record_delivery_ratio(0.99);
});
let listener = TcpListener::bind("127.0.0.1:0").expect("bind probe listener");
let addr = listener.local_addr().expect("listener local addr");
let monitor = NetworkHealthMonitor::new(cluster.clone(), HealthConfig::default());
monitor.register_tcp_probe_target("node-a", addr);
monitor.check_all();
let node_health = monitor
.get_node_health("node-a")
.expect("expected health check result for node-a");
assert_eq!(node_health.status, HealthStatus::Healthy);
}
}