1use std::collections::HashMap;
11use std::net::{SocketAddr, TcpStream};
12use std::sync::{Arc, Mutex};
13use std::time::{Duration, Instant};
14
15use crate::cluster::{ClusterState, NodeStatus};
16use crate::host::{AccelerationMode, RuntimeProfile};
17use crate::system_profile::SystemProfile;
18use crate::watcher::ProfileChange;
19
20#[derive(Clone, Copy, Debug)]
22pub struct HealthConfig {
23 pub check_interval: Duration,
25 pub timeout: Duration,
27 pub min_successes: usize,
29 pub max_failures: usize,
31 pub healthy_latency_us: f32,
33 pub degraded_latency_us: f32,
35 pub healthy_delivery_ratio: f32,
37 pub degraded_delivery_ratio: f32,
39}
40
41impl Default for HealthConfig {
42 fn default() -> Self {
43 Self {
44 check_interval: Duration::from_secs(5),
45 timeout: Duration::from_secs(3),
46 min_successes: 2,
47 max_failures: 3,
48 healthy_latency_us: 10.0,
49 degraded_latency_us: 20.0,
50 healthy_delivery_ratio: 0.95,
51 degraded_delivery_ratio: 0.80,
52 }
53 }
54}
55
56impl HealthConfig {
57 pub fn autotuned(profile: &RuntimeProfile) -> Self {
59 let (healthy_latency_us, degraded_latency_us, timeout, max_failures) =
60 match profile.acceleration_mode {
61 AccelerationMode::Gpu => (6.0, 14.0, Duration::from_secs(2), 4),
62 AccelerationMode::Avx512 => (8.0, 18.0, Duration::from_secs(2), 3),
63 _ => (10.0, 20.0, Duration::from_secs(3), 3),
64 };
65
66 Self {
67 check_interval: Duration::from_secs(5),
68 timeout,
69 min_successes: if profile.recommended_workers >= 8 {
70 3
71 } else {
72 2
73 },
74 max_failures,
75 healthy_latency_us,
76 degraded_latency_us,
77 healthy_delivery_ratio: if profile.node_resources.vram_gb > 0.0 {
78 0.97
79 } else {
80 0.95
81 },
82 degraded_delivery_ratio: if profile.node_resources.vram_gb > 0.0 {
83 0.88
84 } else {
85 0.80
86 },
87 }
88 }
89}
90
91#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
93pub enum HealthStatus {
94 Healthy,
96 Degraded,
98 Failed,
100 #[default]
102 Unknown,
103}
104
105#[derive(Clone, Debug)]
107pub struct HealthCheckResult {
108 pub node_id: String,
110 pub latency_us: f32,
112 pub delivery_ratio: f32,
114 pub status: HealthStatus,
116 pub timestamp: Instant,
118}
119
120#[derive(Clone, Debug)]
122pub struct NetworkHealthMonitor {
123 cluster: Arc<ClusterState>,
125 config: Arc<Mutex<HealthConfig>>,
127 last_check: Arc<Mutex<Option<Instant>>>,
129 recent_checks: Arc<Mutex<HashMap<String, Vec<HealthCheckResult>>>>,
131 tcp_probe_targets: Arc<Mutex<HashMap<String, SocketAddr>>>,
133}
134
135impl NetworkHealthMonitor {
136 pub fn new(cluster: Arc<ClusterState>, config: HealthConfig) -> Self {
138 Self {
139 cluster,
140 config: Arc::new(Mutex::new(config)),
141 last_check: Arc::new(Mutex::new(None)),
142 recent_checks: Arc::new(Mutex::new(HashMap::new())),
143 tcp_probe_targets: Arc::new(Mutex::new(HashMap::new())),
144 }
145 }
146
147 pub fn with_runtime_profile(cluster: Arc<ClusterState>, profile: &RuntimeProfile) -> Self {
149 Self::new(cluster, HealthConfig::autotuned(profile))
150 }
151
152 pub fn config(&self) -> HealthConfig {
154 self.config
155 .lock()
156 .ok()
157 .as_deref()
158 .copied()
159 .unwrap_or_default()
160 }
161
162 pub fn register_tcp_probe_target(&self, node_id: impl Into<String>, target: SocketAddr) {
164 self.tcp_probe_targets
165 .lock()
166 .unwrap_or_else(|poison| poison.into_inner())
167 .insert(node_id.into(), target);
168 }
169
170 pub fn clear_tcp_probe_target(&self, node_id: &str) -> bool {
172 self.tcp_probe_targets
173 .lock()
174 .unwrap_or_else(|poison| poison.into_inner())
175 .remove(node_id)
176 .is_some()
177 }
178
179 fn run_active_tcp_probe(&self, node_id: &str) -> Option<bool> {
180 let cfg = self.config();
181 let target = self
182 .tcp_probe_targets
183 .lock()
184 .unwrap_or_else(|poison| poison.into_inner())
185 .get(node_id)
186 .copied();
187
188 target.map(|addr| TcpStream::connect_timeout(&addr, cfg.timeout).is_ok())
189 }
190
191 pub fn check_all(&self) {
193 let cfg = self.config();
194 let now = Instant::now();
195
196 for node in self.cluster.nodes_snapshot().iter() {
197 let timeout = self.cluster.check_heartbeat_timeout(&node.id);
198 let tcp_probe_ok = self.run_active_tcp_probe(&node.id);
199 let result = if let Some(metrics) = self.cluster.get_metrics(&node.id) {
200 let latency_us = if metrics.latency_samples > 0 {
201 metrics.avg_latency_us
202 } else {
203 0.0
204 };
205
206 let delivery_ratio = if metrics.delivery_ratio_initialized {
207 metrics.delivery_ratio
208 } else {
209 1.0
210 };
211
212 let status = if tcp_probe_ok == Some(false)
213 || timeout
214 || metrics.status == NodeStatus::Failed
215 {
216 HealthStatus::Failed
217 } else if metrics.latency_samples == 0 && !metrics.delivery_ratio_initialized {
218 HealthStatus::Unknown
219 } else {
220 self.get_health_status(latency_us, delivery_ratio)
221 };
222
223 let measured_latency_us = if tcp_probe_ok == Some(false) {
224 cfg.timeout.as_secs_f32() * 1_000_000.0
225 } else {
226 latency_us
227 };
228
229 HealthCheckResult {
230 node_id: node.id.clone(),
231 latency_us: measured_latency_us,
232 delivery_ratio,
233 status,
234 timestamp: now,
235 }
236 } else {
237 HealthCheckResult {
238 node_id: node.id.clone(),
239 latency_us: 0.0,
240 delivery_ratio: 1.0,
241 status: HealthStatus::Unknown,
242 timestamp: now,
243 }
244 };
245
246 self.cluster.get_metrics_mut(&node.id, |m| {
248 m.status = match result.status {
249 HealthStatus::Healthy => NodeStatus::Active,
250 HealthStatus::Degraded => NodeStatus::Degraded,
251 HealthStatus::Failed => NodeStatus::Failed,
252 HealthStatus::Unknown => m.status,
253 };
254 });
255
256 let mut checks = self
258 .recent_checks
259 .lock()
260 .unwrap_or_else(|poison| poison.into_inner());
261 if let Some(node_checks) = checks.get_mut(&node.id) {
262 node_checks.push(result);
263 if node_checks.len() > 10 {
264 node_checks.remove(0);
265 }
266 } else {
267 checks.insert(node.id.clone(), vec![result]);
268 }
269 }
270
271 *self
272 .last_check
273 .lock()
274 .unwrap_or_else(|poison| poison.into_inner()) = Some(now);
275 }
276
277 fn get_health_status(&self, latency_us: f32, delivery_ratio: f32) -> HealthStatus {
287 let cfg = self.config();
288 if delivery_ratio >= cfg.healthy_delivery_ratio && latency_us <= cfg.healthy_latency_us {
289 HealthStatus::Healthy
290 } else if delivery_ratio >= cfg.degraded_delivery_ratio
291 && latency_us <= cfg.degraded_latency_us
292 {
293 HealthStatus::Degraded
294 } else {
295 HealthStatus::Failed
296 }
297 }
298
299 pub fn get_node_health(&self, node_id: &str) -> Option<HealthCheckResult> {
301 let checks = self
302 .recent_checks
303 .lock()
304 .unwrap_or_else(|poison| poison.into_inner());
305 checks
306 .get(node_id)
307 .and_then(|checks| checks.last())
308 .cloned()
309 }
310
311 pub fn get_all_health(&self) -> Vec<HealthCheckResult> {
313 let checks = self
314 .recent_checks
315 .lock()
316 .unwrap_or_else(|poison| poison.into_inner());
317 checks.values().flat_map(|checks| checks.clone()).collect()
318 }
319
320 pub fn healthy_node_count(&self) -> usize {
322 let checks = self
323 .recent_checks
324 .lock()
325 .unwrap_or_else(|poison| poison.into_inner());
326 checks
327 .values()
328 .filter(|node_checks| {
329 node_checks
330 .last()
331 .map(|c| c.status == HealthStatus::Healthy)
332 .unwrap_or(false)
333 })
334 .count()
335 }
336
337 pub fn needs_quantization_fallback(&self, node_id: &str) -> bool {
339 let cfg = self.config();
340 let checks = self
341 .recent_checks
342 .lock()
343 .unwrap_or_else(|poison| poison.into_inner());
344
345 if let Some(node_checks) = checks.get(node_id) {
346 let recent: Vec<_> = node_checks.iter().rev().take(3).collect();
348
349 if recent.len() < 3 {
350 return false;
351 }
352
353 let avg_delivery_ratio: f32 =
355 recent.iter().map(|r| r.delivery_ratio).sum::<f32>() / 3.0;
356
357 let avg_latency_us: f32 = recent.iter().map(|r| r.latency_us).sum::<f32>() / 3.0;
359
360 avg_delivery_ratio < cfg.healthy_delivery_ratio
362 || avg_latency_us > cfg.healthy_latency_us
363 } else {
364 false
365 }
366 }
367
368 pub fn get_health_summary(&self) -> String {
370 let checks = self
371 .recent_checks
372 .lock()
373 .unwrap_or_else(|poison| poison.into_inner());
374
375 let total_nodes = self.cluster.node_count();
376 let healthy_count = checks
377 .values()
378 .filter(|node_checks| {
379 node_checks
380 .last()
381 .map(|c| c.status == HealthStatus::Healthy)
382 .unwrap_or(false)
383 })
384 .count();
385
386 let degraded_count = checks
387 .values()
388 .filter(|node_checks| {
389 node_checks
390 .last()
391 .map(|c| c.status == HealthStatus::Degraded)
392 .unwrap_or(false)
393 })
394 .count();
395
396 let failed_count = checks
397 .values()
398 .filter(|node_checks| {
399 node_checks
400 .last()
401 .map(|c| c.status == HealthStatus::Failed)
402 .unwrap_or(false)
403 })
404 .count();
405
406 format!(
407 "Network Health Summary\n\
408 =================\n\
409 Total nodes: {total_nodes}\n\
410 Healthy: {healthy_count}\n\
411 Degraded: {degraded_count}\n\
412 Failed: {failed_count}\n"
413 )
414 }
415
416 pub fn start_periodic_checks(&self) {
418 let check_interval = self.config().check_interval;
419 let this = self.clone();
420
421 std::thread::spawn(move || loop {
422 this.check_all();
423 std::thread::sleep(check_interval);
424 });
425 }
426
427 pub fn reconfigure_from_system_profile(&self, profile: &SystemProfile) {
437 let rp: RuntimeProfile = profile.into();
438 if let Ok(mut guard) = self.config.lock() {
439 *guard = HealthConfig::autotuned(&rp);
440 }
441 }
442
443 pub fn subscribe_to_watcher(
448 self: &Arc<Self>,
449 watcher: &crate::watcher::SystemProfileWatcher,
450 ) -> tokio::task::JoinHandle<()> {
451 let this = self.clone();
452 let mut rx = watcher.subscribe();
453 tokio::spawn(async move {
454 while let Ok(event) = rx.recv().await {
455 if let ProfileChange::Updated(profile) = event {
456 this.reconfigure_from_system_profile(&profile);
457 tracing::info!(
458 "health monitor re-tuned from profile: {} workers, {:?} accel",
459 profile.recommended_workers,
460 profile.acceleration_mode,
461 );
462 }
463 }
464 })
465 }
466}
467
468#[derive(Clone, Debug, Default)]
470pub struct HealthMetrics {
471 pub avg_latency_us: f32,
473 pub min_latency_us: f32,
475 pub max_latency_us: f32,
477 pub avg_delivery_ratio: f32,
479 pub min_delivery_ratio: f32,
481 pub sample_count: usize,
483}
484
485impl HealthMetrics {
486 pub fn new() -> Self {
488 Self::default()
489 }
490
491 pub fn record(&mut self, latency_us: f32, delivery_ratio: f32) {
493 if self.sample_count == 0 {
494 self.min_latency_us = latency_us;
495 self.max_latency_us = latency_us;
496 self.min_delivery_ratio = delivery_ratio;
497 } else {
498 if latency_us < self.min_latency_us {
499 self.min_latency_us = latency_us;
500 }
501 if latency_us > self.max_latency_us {
502 self.max_latency_us = latency_us;
503 }
504 if delivery_ratio < self.min_delivery_ratio {
505 self.min_delivery_ratio = delivery_ratio;
506 }
507 }
508
509 let n = self.sample_count as f32;
511 self.avg_latency_us = (self.avg_latency_us * n + latency_us) / (n + 1.0);
512 self.avg_delivery_ratio = (self.avg_delivery_ratio * n + delivery_ratio) / (n + 1.0);
513
514 self.sample_count += 1;
515 }
516
517 pub fn get_recommendation(&self) -> String {
519 if self.sample_count == 0 {
520 return "No data available".to_string();
521 }
522
523 match (self.avg_delivery_ratio, self.avg_latency_us) {
524 (ratio, latency) if ratio >= 0.95 && latency <= 10.0 => {
525 "Cluster is healthy. No action needed.".to_string()
526 }
527 (ratio, latency) if ratio >= 0.90 && latency <= 20.0 => {
528 "Cluster is degraded. Consider load balancing.".to_string()
529 }
530 (ratio, latency) if ratio >= 0.80 && latency <= 50.0 => {
531 "Cluster performance is poor. Recommend quantization fallback.".to_string()
532 }
533 _ => "Cluster has failed nodes. Investigate and recover.".to_string(),
534 }
535 }
536}
537
538#[derive(Clone, Debug)]
540pub struct FaultDetector {
541 cluster: Arc<ClusterState>,
543 detection_interval: Duration,
545}
546
547impl FaultDetector {
548 pub fn new(cluster: Arc<ClusterState>, detection_interval: Duration) -> Self {
550 Self {
551 cluster,
552 detection_interval,
553 }
554 }
555
556 pub fn detect_failures(&self) -> Vec<String> {
558 let mut failed_nodes: Vec<String> = Vec::new();
559
560 for node in self.cluster.nodes_snapshot().iter() {
561 if self.cluster.check_heartbeat_timeout(&node.id) {
562 if let Some(metrics) = self.cluster.get_metrics(&node.id) {
563 if metrics.status != NodeStatus::Failed {
570 self.cluster.mark_failed(&node.id);
571 failed_nodes.push(node.id.clone());
572 }
573 } else {
574 self.cluster.mark_failed(&node.id);
576 failed_nodes.push(node.id.clone());
577 }
578 }
579 }
580
581 failed_nodes
582 }
583
584 pub fn recover_node(&self, node_id: &str) {
586 self.cluster.get_metrics_mut(node_id, |metrics| {
587 metrics.status = NodeStatus::Active;
588 metrics.last_heartbeat = Instant::now();
589
590 metrics.avg_latency_us = 0.0;
592 metrics.min_latency_us = f32::MAX;
593 metrics.max_latency_us = 0.0;
594 metrics.latency_samples = 0;
595 metrics.delivery_ratio = 1.0;
596 metrics.throughput_gbps = 0.0;
597 });
598 }
599
600 pub fn start_periodic_detection(&self) {
602 let this = self.clone();
603
604 std::thread::spawn(move || loop {
605 let failed = this.detect_failures();
606 if !failed.is_empty() {
607 tracing::warn!(
608 "Detected {} failed nodes: {:?}",
609 failed.len(),
610 failed.iter().map(|s| s.as_str()).collect::<Vec<_>>()
611 );
612 }
613 std::thread::sleep(this.detection_interval);
614 });
615 }
616}
617
618#[cfg(test)]
619mod tests {
620 use super::*;
621 use crate::cluster::ClusterState;
622 use crate::protocol::NodeResources;
623 use std::net::TcpListener;
624 use std::sync::Arc;
625 use std::time::Duration;
626
627 #[test]
628 fn health_monitor_checks_all_nodes() {
629 let cluster = Arc::new(ClusterState::new());
630 cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
631 cluster.register(NodeResources::new("node-b", 12.0, 32.0, "8.6", None));
632
633 let monitor = NetworkHealthMonitor::new(cluster.clone(), HealthConfig::default());
634 monitor.check_all();
635
636 let summary = monitor.get_health_summary();
637 assert!(summary.contains("Total nodes: 2"));
638 }
639
640 #[test]
641 fn health_monitor_marks_fresh_nodes_unknown_without_samples() {
642 let cluster = Arc::new(ClusterState::new());
643 cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
644
645 let monitor = NetworkHealthMonitor::new(cluster, HealthConfig::default());
646 monitor.check_all();
647
648 let node_health = monitor
649 .get_node_health("node-a")
650 .expect("expected health check result for node-a");
651 assert_eq!(node_health.status, HealthStatus::Unknown);
652 }
653
654 #[test]
655 fn health_monitor_detects_degraded_nodes() {
656 let cluster = Arc::new(ClusterState::new());
657 cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
658
659 let monitor = NetworkHealthMonitor::new(cluster.clone(), HealthConfig::default());
660
661 for _ in 0..3 {
664 let result = HealthCheckResult {
665 node_id: "node-a".to_string(),
666 latency_us: 15.0,
667 delivery_ratio: 0.85,
668 status: HealthStatus::Degraded,
669 timestamp: std::time::Instant::now(),
670 };
671 monitor
672 .recent_checks
673 .lock()
674 .unwrap()
675 .entry("node-a".to_string())
676 .or_default()
677 .push(result);
678 }
679
680 assert!(monitor.needs_quantization_fallback("node-a"));
681 }
682
683 #[test]
684 fn health_metrics_aggregates() {
685 let mut metrics = HealthMetrics::new();
686
687 metrics.record(1.0, 0.98);
688 metrics.record(2.0, 0.95);
689 metrics.record(1.5, 0.97);
690
691 assert!((metrics.avg_latency_us - 1.5).abs() < 0.1);
692 assert!(metrics.min_latency_us <= 1.0);
693 assert!(metrics.max_latency_us >= 2.0);
694 }
695
696 #[test]
697 fn fault_detector_detects_failures() {
698 let cluster = Arc::new(ClusterState::new());
699 cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
700
701 std::thread::sleep(Duration::from_secs(6));
703
704 let detector = FaultDetector::new(cluster.clone(), Duration::from_secs(1));
705 let failed = detector.detect_failures();
706
707 assert!(failed.contains(&"node-a".to_string()));
708 }
709
710 #[test]
711 fn fault_detector_detects_failures_for_degraded_nodes() {
712 let cluster = Arc::new(ClusterState::new());
718 cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
719 cluster.get_metrics_mut("node-a", |metrics| {
720 metrics.status = NodeStatus::Degraded;
721 });
722
723 std::thread::sleep(Duration::from_secs(6));
724
725 let detector = FaultDetector::new(cluster.clone(), Duration::from_secs(1));
726 let failed = detector.detect_failures();
727
728 assert!(failed.contains(&"node-a".to_string()));
729 assert_eq!(
730 cluster.get_metrics("node-a").unwrap().status,
731 NodeStatus::Failed
732 );
733 }
734
735 #[test]
736 fn fault_detector_recovers_nodes() {
737 let cluster = Arc::new(ClusterState::new());
738 cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
739
740 std::thread::sleep(Duration::from_secs(6));
741
742 let detector = FaultDetector::new(cluster.clone(), Duration::from_secs(1));
743 detector.detect_failures();
744
745 assert!(cluster.get_metrics("node-a").unwrap().status == NodeStatus::Failed);
746
747 detector.recover_node("node-a");
748 assert_eq!(
749 cluster.get_metrics("node-a").unwrap().status,
750 NodeStatus::Active
751 );
752 }
753
754 #[test]
755 fn health_monitor_gets_recommendation() {
756 let mut metrics = HealthMetrics::new();
757
758 metrics.record(1.0, 0.98);
760 metrics.record(2.0, 0.95);
761
762 let recommendation = metrics.get_recommendation();
763 assert!(recommendation.contains("healthy"));
764 }
765
766 #[test]
767 fn get_recommendation_uses_average_latency_not_minimum() {
768 let mut metrics = HealthMetrics::new();
774 metrics.record(1.0, 0.98);
775 metrics.record(200.0, 0.98);
776
777 let recommendation = metrics.get_recommendation();
778 assert!(!recommendation.contains("healthy"));
779 }
780
781 #[test]
782 fn health_config_autotunes_for_gpu_profiles() {
783 let profile = RuntimeProfile {
784 node_resources: NodeResources::new("node-a", 24.0, 64.0, "8.9", None),
785 logical_cores: 16,
786 recommended_workers: 8,
787 acceleration_mode: AccelerationMode::Gpu,
788 gpu_backend: crate::host::GpuBackend::Cuda,
789 xdp_supported: true,
790 detection_source: String::from("test"),
791 probe_mode: crate::host::ProbeMode::Fast,
792 };
793
794 let config = HealthConfig::autotuned(&profile);
795 assert!(config.healthy_latency_us < 10.0);
796 assert!(config.healthy_delivery_ratio > 0.95);
797 }
798
799 #[test]
800 fn get_health_status_fails_on_either_metric_crossing_floor() {
801 let cluster = Arc::new(ClusterState::new());
806 let monitor = NetworkHealthMonitor::new(cluster, HealthConfig::default());
807 let cfg = HealthConfig::default();
808
809 assert_eq!(
811 monitor.get_health_status(cfg.degraded_latency_us * 10.0, 0.99),
812 HealthStatus::Failed
813 );
814 assert_eq!(
816 monitor.get_health_status(1.0, cfg.degraded_delivery_ratio / 2.0),
817 HealthStatus::Failed
818 );
819 assert_eq!(
821 monitor.get_health_status(cfg.degraded_latency_us, cfg.degraded_delivery_ratio),
822 HealthStatus::Degraded
823 );
824 }
825
826 #[test]
827 fn tcp_probe_target_can_mark_node_failed_when_unreachable() {
828 let cluster = Arc::new(ClusterState::new());
829 cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
830
831 let listener = TcpListener::bind("127.0.0.1:0").expect("bind probe listener");
832 let addr = listener.local_addr().expect("listener local addr");
833 drop(listener);
834
835 let monitor = NetworkHealthMonitor::new(cluster.clone(), HealthConfig::default());
836 monitor.register_tcp_probe_target("node-a", addr);
837 monitor.check_all();
838
839 let node_health = monitor
840 .get_node_health("node-a")
841 .expect("expected health check result for node-a");
842 assert_eq!(node_health.status, HealthStatus::Failed);
843 }
844
845 #[test]
846 fn tcp_probe_target_healthy_when_reachable() {
847 let cluster = Arc::new(ClusterState::new());
848 cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
849
850 cluster.get_metrics_mut("node-a", |metrics| {
851 metrics.record_latency(2.0);
852 metrics.record_delivery_ratio(0.99);
853 });
854
855 let listener = TcpListener::bind("127.0.0.1:0").expect("bind probe listener");
856 let addr = listener.local_addr().expect("listener local addr");
857
858 let monitor = NetworkHealthMonitor::new(cluster.clone(), HealthConfig::default());
859 monitor.register_tcp_probe_target("node-a", addr);
860 monitor.check_all();
861
862 let node_health = monitor
863 .get_node_health("node-a")
864 .expect("expected health check result for node-a");
865 assert_eq!(node_health.status, HealthStatus::Healthy);
866 }
867}