Skip to main content

ghostlink_core/
cluster.rs

1//! Thread-Safe Cluster State with Metrics Collection
2//!
3//! This module provides a lock-free, thread-safe cluster state tracking:
4//! - Node capabilities (VRAM, system memory, compute capability)
5//! - Live metrics (latency, delivery ratio, throughput)
6//! - Fault detection and recovery
7
8use std::collections::{HashMap, VecDeque};
9use std::net::SocketAddr;
10use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
11use std::sync::{Arc, Mutex};
12use std::time::{Duration, Instant};
13
14use arc_swap::ArcSwap;
15
16pub use crate::protocol::NodeResources;
17
18const NODE_HISTORY_CAP: usize = 32;
19
20/// Node status enumeration
21#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
22pub enum NodeStatus {
23    /// Node is healthy and accepting traffic
24    #[default]
25    Active,
26    /// Node is degraded (below threshold)
27    Degraded,
28    /// Node has failed or timed out
29    Failed,
30}
31
32/// Metrics for a single node
33#[derive(Clone, Debug)]
34pub struct NodeMetrics {
35    /// Node display name
36    pub name: String,
37    /// Current status
38    pub status: NodeStatus,
39    /// Total VRAM in GB
40    pub vram_gb: f32,
41    /// Total VRAM (alias for vram_gb, used by display/balancer)
42    pub total_vram_gb: f32,
43    /// System memory in GB
44    pub system_memory_gb: f32,
45    /// Compute capability
46    pub compute_capability: String,
47    /// GPU name/model
48    pub gpu_name: Option<String>,
49    /// Resolved IP address of the node
50    pub ip_address: Option<SocketAddr>,
51
52    /// Last heartbeat time
53    pub last_heartbeat: Instant,
54    /// Heartbeat interval threshold
55    pub heartbeat_timeout: Duration,
56
57    /// Average latency in microseconds
58    pub avg_latency_us: f32,
59    /// Minimum observed latency
60    pub min_latency_us: f32,
61    /// Maximum observed latency
62    pub max_latency_us: f32,
63    /// Number of latency samples
64    pub latency_samples: u64,
65
66    /// Delivery ratio (0.0 to 1.0)
67    pub delivery_ratio: f32,
68    /// Throughput in GB/s
69    pub throughput_gbps: f32,
70
71    /// Recent latency history in microseconds
72    pub latency_history_us: VecDeque<f32>,
73    /// Recent throughput history in GB/s
74    pub throughput_history_gbps: VecDeque<f32>,
75
76    /// Current used VRAM in GB
77    pub used_vram_gb: f32,
78    /// Available VRAM in GB
79    pub available_vram_gb: f32,
80
81    /// Number of layers streaming on this node
82    pub streaming_layers: Option<(usize, usize)>,
83
84    /// AF_XDP throughput in Gbps (for display)
85    pub af_xdp_gbps: f32,
86    /// Per-packet latency in microseconds (for display)
87    pub latency_micros: f32,
88    /// Whether delivery_ratio has been initialized (first sample sets directly)
89    pub(crate) delivery_ratio_initialized: bool,
90}
91
92impl Default for NodeMetrics {
93    fn default() -> Self {
94        Self {
95            name: String::new(),
96            status: NodeStatus::Active,
97            vram_gb: 0.0,
98            total_vram_gb: 0.0,
99            system_memory_gb: 0.0,
100            compute_capability: String::new(),
101            gpu_name: None,
102            last_heartbeat: Instant::now(),
103            heartbeat_timeout: Duration::from_secs(5),
104            avg_latency_us: 0.0,
105            min_latency_us: f32::MAX,
106            max_latency_us: 0.0,
107            latency_samples: 0,
108            delivery_ratio: 1.0,
109            throughput_gbps: 0.0,
110            latency_history_us: VecDeque::with_capacity(NODE_HISTORY_CAP),
111            throughput_history_gbps: VecDeque::with_capacity(NODE_HISTORY_CAP),
112            used_vram_gb: 0.0,
113            available_vram_gb: 0.0,
114            streaming_layers: None,
115            af_xdp_gbps: 0.0,
116            latency_micros: 0.0,
117            delivery_ratio_initialized: false,
118            ip_address: None,
119        }
120    }
121}
122
123impl NodeMetrics {
124    /// Create new metrics for a node
125    pub fn new(
126        vram_gb: f32,
127        system_memory_gb: f32,
128        compute_capability: String,
129        heartbeat_timeout: Duration,
130    ) -> Self {
131        Self {
132            name: String::new(),
133            status: NodeStatus::Active,
134            vram_gb,
135            total_vram_gb: vram_gb,
136            system_memory_gb,
137            compute_capability,
138            gpu_name: None,
139            ip_address: None,
140            last_heartbeat: Instant::now(),
141            heartbeat_timeout,
142            avg_latency_us: 0.0,
143            min_latency_us: f32::MAX,
144            max_latency_us: 0.0,
145            latency_samples: 0,
146            delivery_ratio: 1.0,
147            throughput_gbps: 0.0,
148            latency_history_us: VecDeque::with_capacity(NODE_HISTORY_CAP),
149            throughput_history_gbps: VecDeque::with_capacity(NODE_HISTORY_CAP),
150            used_vram_gb: 0.0,
151            available_vram_gb: vram_gb,
152            streaming_layers: None,
153            af_xdp_gbps: 0.0,
154            latency_micros: 0.0,
155            delivery_ratio_initialized: false,
156        }
157    }
158
159    /// Update metrics with new latency sample (EMA with alpha=0.1)
160    pub fn record_latency(&mut self, latency_us: f32) {
161        if latency_us < self.min_latency_us {
162            self.min_latency_us = latency_us;
163        }
164        if latency_us > self.max_latency_us {
165            self.max_latency_us = latency_us;
166        }
167
168        self.latency_samples += 1;
169        if self.latency_samples == 1 {
170            self.avg_latency_us = latency_us;
171        } else {
172            self.avg_latency_us = self.avg_latency_us * 0.9 + latency_us * 0.1;
173        }
174
175        if self.latency_history_us.len() >= NODE_HISTORY_CAP {
176            self.latency_history_us.pop_front();
177        }
178        self.latency_history_us.push_back(latency_us);
179    }
180
181    /// Update metrics with new delivery ratio sample
182    pub fn record_delivery_ratio(&mut self, ratio: f32) {
183        if !self.delivery_ratio_initialized {
184            self.delivery_ratio = ratio;
185            self.delivery_ratio_initialized = true;
186        } else {
187            // Exponential moving average with alpha=0.1
188            self.delivery_ratio = self.delivery_ratio * 0.9 + ratio * 0.1;
189        }
190    }
191
192    /// Update metrics with new throughput sample
193    pub fn record_throughput(&mut self, throughput_gbps: f32) {
194        // Exponential moving average with alpha=0.1
195        self.throughput_gbps = self.throughput_gbps * 0.9 + throughput_gbps * 0.1;
196
197        if self.throughput_history_gbps.len() >= NODE_HISTORY_CAP {
198            self.throughput_history_gbps.pop_front();
199        }
200        self.throughput_history_gbps.push_back(throughput_gbps);
201    }
202
203    /// Update used VRAM
204    pub fn record_vram_usage(&mut self, used_vram_gb: f32) {
205        self.used_vram_gb = used_vram_gb;
206        self.available_vram_gb = self.total_vram_gb - used_vram_gb;
207    }
208
209    /// Set streaming layers
210    pub fn set_streaming_layers(&mut self, start: usize, end: usize) {
211        self.streaming_layers = Some((start, end));
212    }
213
214    /// Clear streaming layers
215    pub fn clear_streaming_layers(&mut self) {
216        self.streaming_layers = None;
217    }
218}
219
220/// Cluster state with thread-safe metrics collection
221#[derive(Debug)]
222pub struct ClusterState {
223    /// Map of node ID to resources and metrics
224    nodes: Arc<Mutex<HashMap<String, NodeResources>>>,
225    /// Cached shared snapshot of nodes for read-heavy paths
226    nodes_snapshot: Arc<ArcSwap<Vec<NodeResources>>>,
227    /// Indicates whether the shared snapshot needs to be refreshed
228    nodes_snapshot_dirty: Arc<AtomicBool>,
229    /// Map of node ID to live metrics
230    pub(crate) metrics: Arc<Mutex<HashMap<String, NodeMetrics>>>,
231    /// Last cluster update timestamp
232    last_update: Arc<AtomicU64>,
233    /// Cached total VRAM across all registered nodes
234    total_vram_cache: Arc<AtomicU64>,
235    /// Cached total system memory across all registered nodes
236    total_system_memory_cache: Arc<AtomicU64>,
237    /// Per-node circuit breakers for the TCP transport bridge (`runtime::spawn_tcp_bridge`).
238    /// Lazily created on first use via `circuit_breaker_for` and kept here — rather than
239    /// inside a single pipeline execution — so failures accumulate *across* pipeline runs
240    /// targeting the same node: a node that's been chronically unreachable trips open and
241    /// stays fail-fast instead of re-running the full connect/backoff dance every call.
242    circuit_breakers: Arc<Mutex<HashMap<String, crate::circuit_breaker::CircuitBreaker>>>,
243}
244
245impl Clone for ClusterState {
246    fn clone(&self) -> Self {
247        Self {
248            nodes: Arc::clone(&self.nodes),
249            nodes_snapshot: Arc::clone(&self.nodes_snapshot),
250            nodes_snapshot_dirty: Arc::clone(&self.nodes_snapshot_dirty),
251            metrics: Arc::clone(&self.metrics),
252            last_update: Arc::clone(&self.last_update),
253            total_vram_cache: Arc::clone(&self.total_vram_cache),
254            total_system_memory_cache: Arc::clone(&self.total_system_memory_cache),
255            circuit_breakers: Arc::clone(&self.circuit_breakers),
256        }
257    }
258}
259
260impl Default for ClusterState {
261    fn default() -> Self {
262        Self::new()
263    }
264}
265
266impl ClusterState {
267    /// Create new cluster state
268    pub fn new() -> Self {
269        Self {
270            nodes: Arc::new(Mutex::new(HashMap::new())),
271            nodes_snapshot: Arc::new(ArcSwap::from_pointee(Vec::<NodeResources>::new())),
272            nodes_snapshot_dirty: Arc::new(AtomicBool::new(false)),
273            metrics: Arc::new(Mutex::new(HashMap::new())),
274            last_update: Arc::new(AtomicU64::new(0)),
275            total_vram_cache: Arc::new(AtomicU64::new(0.0_f64.to_bits())),
276            total_system_memory_cache: Arc::new(AtomicU64::new(0.0_f64.to_bits())),
277            circuit_breakers: Arc::new(Mutex::new(HashMap::new())),
278        }
279    }
280
281    /// Gets (creating with default config on first use) the circuit breaker
282    /// tracking TCP transport failures for one node. The returned handle is
283    /// a cheap `Arc`-backed clone sharing state with every other caller for
284    /// the same `node_id` — recording a failure through one clone is visible
285    /// to all of them.
286    pub fn circuit_breaker_for(&self, node_id: &str) -> crate::circuit_breaker::CircuitBreaker {
287        let mut breakers = self
288            .circuit_breakers
289            .lock()
290            .unwrap_or_else(|poison| poison.into_inner());
291        breakers.entry(node_id.to_string()).or_default().clone()
292    }
293
294    /// Register a new node with the cluster
295    pub fn register(&self, node: NodeResources) {
296        self.register_with_addr(node, None);
297    }
298
299    /// Register a new node with a specific socket address
300    pub fn register_with_addr(&self, node: NodeResources, addr: Option<SocketAddr>) {
301        let mut nodes = self
302            .nodes
303            .lock()
304            .unwrap_or_else(|poison| poison.into_inner());
305        let mut metrics = self
306            .metrics
307            .lock()
308            .unwrap_or_else(|poison| poison.into_inner());
309
310        // Optimization: Operate on references/moves and avoid redundant cloning of NodeResources fields.
311        // Also avoid marking nodes_snapshot_dirty when updating a node with unchanged resources.
312        let vram_gb = node.vram_gb;
313        let system_memory_gb = node.system_memory_gb;
314        let mut vram_delta = vram_gb;
315        let mut system_memory_delta = system_memory_gb;
316        let mut resources_changed = true;
317
318        if let Some(existing) = nodes.get_mut(&node.id) {
319            vram_delta = vram_gb - existing.vram_gb;
320            system_memory_delta = system_memory_gb - existing.system_memory_gb;
321
322            // Check if node resource fields actually changed to avoid snapshot invalidation churn
323            if existing.vram_gb == node.vram_gb
324                && existing.system_memory_gb == node.system_memory_gb
325                && existing.compute_capability == node.compute_capability
326                && existing.gpu_name == node.gpu_name
327                && existing.rpc_port == node.rpc_port
328                && existing.rpc_build_id == node.rpc_build_id
329            {
330                resources_changed = false;
331            } else {
332                *existing = node.clone();
333            }
334        } else {
335            nodes.insert(node.id.clone(), node.clone());
336        }
337
338        if let Some(existing_metrics) = metrics.get_mut(&node.id) {
339            existing_metrics.vram_gb = vram_gb;
340            existing_metrics.total_vram_gb = vram_gb;
341            existing_metrics.system_memory_gb = system_memory_gb;
342            existing_metrics
343                .compute_capability
344                .clone_from(&node.compute_capability);
345            existing_metrics.gpu_name.clone_from(&node.gpu_name);
346            existing_metrics.heartbeat_timeout = Duration::from_secs(5);
347            if addr.is_some() {
348                existing_metrics.ip_address = addr;
349            }
350        } else {
351            let mut node_metrics = NodeMetrics::new(
352                vram_gb,
353                system_memory_gb,
354                node.compute_capability.clone(),
355                Duration::from_secs(5),
356            );
357            node_metrics.name = node.id.clone();
358            node_metrics.gpu_name = node.gpu_name.clone();
359            node_metrics.ip_address = addr;
360            metrics.insert(node.id.clone(), node_metrics);
361        }
362
363        if resources_changed {
364            self.nodes_snapshot_dirty.store(true, Ordering::Release);
365        }
366
367        let current_total_vram = f64::from_bits(self.total_vram_cache.load(Ordering::Acquire));
368        self.total_vram_cache.store(
369            (current_total_vram + vram_delta as f64).to_bits(),
370            Ordering::Release,
371        );
372
373        let current_total_system_mem =
374            f64::from_bits(self.total_system_memory_cache.load(Ordering::Acquire));
375        self.total_system_memory_cache.store(
376            (current_total_system_mem + system_memory_delta as f64).to_bits(),
377            Ordering::Release,
378        );
379
380        self.last_update.store(
381            std::time::SystemTime::now()
382                .duration_since(std::time::UNIX_EPOCH)
383                .unwrap_or(Duration::ZERO)
384                .as_millis() as u64,
385            Ordering::Release,
386        );
387    }
388
389    /// Get all nodes
390    pub fn nodes(&self) -> Vec<NodeResources> {
391        self.nodes_snapshot().as_ref().to_vec()
392    }
393
394    /// Get total number of registered nodes.
395    pub fn node_count(&self) -> usize {
396        self.nodes_snapshot().len()
397    }
398
399    /// Get a shared snapshot of all nodes
400    pub fn nodes_snapshot(&self) -> Arc<Vec<NodeResources>> {
401        if self.nodes_snapshot_dirty.load(Ordering::Relaxed)
402            && self.nodes_snapshot_dirty.swap(false, Ordering::AcqRel)
403        {
404            let nodes_map = self
405                .nodes
406                .lock()
407                .unwrap_or_else(|poison| poison.into_inner());
408            let mut nodes: Vec<_> = nodes_map.values().cloned().collect();
409            // Sort by ID to ensure deterministic layer assignment across runs
410            nodes.sort_unstable_by(|a, b| a.id.cmp(&b.id));
411            self.nodes_snapshot.store(Arc::new(nodes));
412        }
413        self.nodes_snapshot.load_full()
414    }
415
416    /// Execute a closure with a read lock on the metrics map, avoiding clones and multiple acquisitions.
417    pub fn with_metrics<F, R>(&self, f: F) -> R
418    where
419        F: FnOnce(&HashMap<String, NodeMetrics>) -> R,
420    {
421        let metrics = self
422            .metrics
423            .lock()
424            .unwrap_or_else(|poison| poison.into_inner());
425        f(&metrics)
426    }
427
428    /// Get metrics for a specific node
429    pub fn get_metrics(&self, node_id: &str) -> Option<NodeMetrics> {
430        let metrics = self
431            .metrics
432            .lock()
433            .unwrap_or_else(|poison| poison.into_inner());
434        metrics.get(node_id).cloned()
435    }
436
437    /// Update last heartbeat for a node
438    pub fn update_heartbeat(&self, node_id: &str) {
439        self.get_metrics_mut(node_id, |metrics| {
440            metrics.last_heartbeat = Instant::now();
441        });
442    }
443
444    /// Get metrics mutable reference (for internal updates)
445    pub fn get_metrics_mut<F, R>(&self, node_id: &str, f: F) -> Option<R>
446    where
447        F: FnOnce(&mut NodeMetrics) -> R,
448    {
449        let mut metrics = self
450            .metrics
451            .lock()
452            .unwrap_or_else(|poison| poison.into_inner());
453        metrics.get_mut(node_id).map(f)
454    }
455
456    /// Check if a node has timed out
457    pub fn check_heartbeat_timeout(&self, node_id: &str) -> bool {
458        // Optimize: Inspect node metrics directly under lock without cloning NodeMetrics struct.
459        let metrics = self
460            .metrics
461            .lock()
462            .unwrap_or_else(|poison| poison.into_inner());
463        if let Some(m) = metrics.get(node_id) {
464            Instant::now().duration_since(m.last_heartbeat) >= m.heartbeat_timeout
465        } else {
466            false
467        }
468    }
469
470    /// Mark a node as failed due to timeout
471    pub fn mark_failed(&self, node_id: &str) {
472        self.get_metrics_mut(node_id, |metrics| {
473            metrics.status = NodeStatus::Failed;
474        });
475    }
476
477    /// Recover a failed node
478    pub fn recover_node(&self, node_id: &str) {
479        self.get_metrics_mut(node_id, |metrics| {
480            metrics.status = NodeStatus::Active;
481        });
482    }
483
484    /// Get all active nodes
485    pub fn active_nodes(&self) -> Vec<NodeMetrics> {
486        let metrics = self
487            .metrics
488            .lock()
489            .unwrap_or_else(|poison| poison.into_inner());
490        metrics
491            .values()
492            .filter(|m| m.status == NodeStatus::Active)
493            .cloned()
494            .collect()
495    }
496
497    /// Get the count of active nodes without cloning metrics
498    pub fn active_nodes_count(&self) -> usize {
499        let metrics = self
500            .metrics
501            .lock()
502            .unwrap_or_else(|poison| poison.into_inner());
503        metrics
504            .values()
505            .filter(|m| m.status == NodeStatus::Active)
506            .count()
507    }
508
509    /// Get total cluster VRAM
510    pub fn total_vram_gb(&self) -> f32 {
511        f64::from_bits(self.total_vram_cache.load(Ordering::Acquire)) as f32
512    }
513
514    /// Get total system memory
515    pub fn total_system_memory_gb(&self) -> f32 {
516        f64::from_bits(self.total_system_memory_cache.load(Ordering::Acquire)) as f32
517    }
518}
519
520/// Cluster health monitor with periodic checks
521#[derive(Clone, Debug)]
522pub struct ClusterHealthMonitor {
523    cluster: Arc<ClusterState>,
524    /// Health check interval
525    check_interval: Duration,
526}
527
528impl ClusterHealthMonitor {
529    /// Create new health monitor
530    pub fn new(cluster: Arc<ClusterState>, check_interval: Duration) -> Self {
531        Self {
532            cluster,
533            check_interval,
534        }
535    }
536
537    /// Run health check on all nodes
538    pub fn check_health(&self) {
539        let failed_nodes: Vec<String> = self
540            .cluster
541            .nodes_snapshot()
542            .iter()
543            .filter(|n| self.cluster.check_heartbeat_timeout(&n.id))
544            .map(|n| n.id.clone())
545            .collect();
546
547        // Mark timed-out nodes as failed
548        for node_id in &failed_nodes {
549            self.cluster.mark_failed(node_id);
550        }
551
552        // Update last update timestamp
553        self.cluster.last_update.store(
554            std::time::SystemTime::now()
555                .duration_since(std::time::UNIX_EPOCH)
556                .unwrap_or(Duration::ZERO)
557                .as_millis() as u64,
558            Ordering::Release,
559        );
560    }
561
562    /// Get health report
563    pub fn health_report(&self) -> String {
564        let active_count = self.cluster.active_nodes_count();
565        let total_nodes = self.cluster.nodes_snapshot().len();
566
567        format!(
568            "Cluster Health Report\n\
569             =================\n\
570             Active nodes: {}/{}\n\
571             Total VRAM: {:.1} GB\n\
572             System memory: {:.1} GB\n",
573            active_count,
574            total_nodes,
575            self.cluster.total_vram_gb(),
576            self.cluster.total_system_memory_gb()
577        )
578    }
579
580    /// Run periodic health checks in background
581    pub fn start_periodic_checks(&self) {
582        let this = self.clone();
583
584        // Spawn health check task
585        std::thread::spawn(move || loop {
586            this.check_health();
587            std::thread::sleep(this.check_interval);
588        });
589    }
590}
591
592#[cfg(test)]
593mod tests {
594    use super::*;
595    use std::sync::Arc;
596    use std::thread;
597    use std::time::Duration;
598
599    #[test]
600    fn register_replaces_existing_nodes() {
601        let cluster = ClusterState::new();
602        cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
603        cluster.register(NodeResources::new("node-a", 48.0, 128.0, "9.0", None));
604
605        assert_eq!(cluster.nodes().len(), 1);
606        assert_eq!(cluster.nodes()[0].vram_gb, 48.0);
607    }
608
609    #[test]
610    fn heartbeat_timeout_detection() {
611        let cluster = ClusterState::new();
612        cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
613
614        // Simulate timeout by waiting longer than default heartbeat timeout
615        thread::sleep(Duration::from_secs(6));
616
617        assert!(cluster.check_heartbeat_timeout("node-a"));
618    }
619
620    #[test]
621    fn node_metrics_keep_recent_latency_and_throughput_history() {
622        let mut metrics = NodeMetrics::default();
623
624        for idx in 0..40 {
625            metrics.record_latency(idx as f32);
626            metrics.record_throughput(idx as f32 / 10.0);
627        }
628
629        assert_eq!(metrics.latency_history_us.len(), NODE_HISTORY_CAP);
630        assert_eq!(metrics.throughput_history_gbps.len(), NODE_HISTORY_CAP);
631        assert_eq!(metrics.latency_history_us.front().copied(), Some(8.0));
632        assert_eq!(metrics.throughput_history_gbps.back().copied(), Some(3.9));
633    }
634
635    #[test]
636    fn health_monitor_reports_active_count() {
637        let cluster = Arc::new(ClusterState::new());
638        cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
639        cluster.register(NodeResources::new("node-b", 12.0, 32.0, "8.6", None));
640
641        let monitor = ClusterHealthMonitor::new(cluster.clone(), Duration::from_secs(1));
642        monitor.check_health();
643
644        let report = monitor.health_report();
645        assert!(report.contains("Active nodes: 2/2"));
646    }
647
648    #[test]
649    fn metrics_record_latency() {
650        let mut metrics = NodeMetrics::new(24.0, 64.0, "8.9".to_string(), Duration::from_secs(5));
651
652        metrics.record_latency(1.0);
653        assert_eq!(metrics.avg_latency_us, 1.0);
654        assert_eq!(metrics.min_latency_us, 1.0);
655        assert_eq!(metrics.max_latency_us, 1.0);
656
657        metrics.record_latency(2.0);
658        // EMA with alpha=0.1: (1.0 * 0.9) + 2.0 * 0.1 = 0.9 + 0.2 = 1.1
659        assert!((metrics.avg_latency_us - 1.1).abs() < 1e-6);
660    }
661
662    #[test]
663    fn metrics_record_delivery_ratio() {
664        let mut metrics = NodeMetrics::new(24.0, 64.0, "8.9".to_string(), Duration::from_secs(5));
665
666        metrics.record_delivery_ratio(0.98);
667        assert!((metrics.delivery_ratio - 0.98).abs() < 1e-6);
668
669        metrics.record_delivery_ratio(0.90);
670        // EMA: 0.98 * 0.9 + 0.90 * 0.1 = 0.882 + 0.09 = 0.972
671        assert!((metrics.delivery_ratio - 0.972).abs() < 1e-6);
672    }
673
674    #[test]
675    fn with_metrics_allows_single_pass_inspection() {
676        let cluster = ClusterState::new();
677        cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
678        cluster.register(NodeResources::new("node-b", 12.0, 32.0, "8.6", None));
679
680        let total_nodes = cluster.with_metrics(|metrics_map| {
681            assert!(metrics_map.contains_key("node-a"));
682            assert!(metrics_map.contains_key("node-b"));
683            metrics_map.len()
684        });
685
686        assert_eq!(total_nodes, 2);
687    }
688
689    #[test]
690    fn cluster_tracks_total_vram() {
691        let cluster = ClusterState::new();
692        cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
693        cluster.register(NodeResources::new("node-b", 12.0, 32.0, "8.6", None));
694
695        assert_eq!(cluster.total_vram_gb(), 36.0);
696    }
697
698    // ========================================================================
699    // NEW: Comprehensive Health Monitoring & Failure Recovery Tests
700    // ========================================================================
701
702    #[test]
703    fn cluster_health_monitor_checks_all_nodes() {
704        let cluster = Arc::new(ClusterState::new());
705        cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
706        cluster.register(NodeResources::new("node-b", 12.0, 32.0, "8.6", None));
707
708        let monitor = ClusterHealthMonitor::new(cluster.clone(), Duration::from_millis(100));
709        monitor.check_health();
710
711        let report = monitor.health_report();
712        assert!(
713            report.contains("Active nodes: 2/2"),
714            "Report should show both nodes as active"
715        );
716    }
717
718    #[test]
719    fn cluster_metrics_use_exponential_moving_average() {
720        let mut metrics = NodeMetrics::new(24.0, 64.0, "8.9".to_string(), Duration::from_secs(5));
721
722        // Record sequence of latencies: 1, 2, 3
723        metrics.record_latency(1.0);
724        let avg1 = metrics.avg_latency_us;
725        assert_eq!(avg1, 1.0);
726
727        metrics.record_latency(2.0);
728        let avg2 = metrics.avg_latency_us;
729        // EMA: 1.0 * 0.9 + 2.0 * 0.1 = 1.1
730        assert!(
731            (avg2 - 1.1).abs() < 1e-6,
732            "EMA calculation: expected 1.1, got {avg2}"
733        );
734
735        metrics.record_latency(3.0);
736        let avg3 = metrics.avg_latency_us;
737        // EMA: 1.1 * 0.9 + 3.0 * 0.1 = 0.99 + 0.3 = 1.29
738        assert!(
739            (avg3 - 1.29).abs() < 1e-5,
740            "EMA calculation: expected 1.29, got {avg3}"
741        );
742    }
743
744    #[test]
745    fn cluster_delivery_ratio_tracks_network_quality() {
746        let mut metrics = NodeMetrics::new(24.0, 64.0, "8.9".to_string(), Duration::from_secs(5));
747
748        metrics.record_delivery_ratio(0.98);
749        assert!((metrics.delivery_ratio - 0.98).abs() < 1e-6);
750
751        metrics.record_delivery_ratio(0.95);
752        // EMA: 0.98 * 0.9 + 0.95 * 0.1 = 0.882 + 0.095 = 0.977
753        assert!((metrics.delivery_ratio - 0.977).abs() < 1e-6);
754
755        // Simulate degradation
756        metrics.record_delivery_ratio(0.80);
757        // EMA: 0.977 * 0.9 + 0.80 * 0.1 = 0.8793 + 0.08 = 0.9593
758        assert!(
759            metrics.delivery_ratio < 0.98,
760            "Delivery ratio should degrade"
761        );
762        assert!(metrics.delivery_ratio > 0.90, "But recover somewhat");
763    }
764
765    #[test]
766    fn cluster_concurrent_metric_updates() {
767        let cluster = Arc::new(ClusterState::new());
768        cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
769
770        let mut handles = vec![];
771
772        // Spawn 10 threads updating metrics concurrently
773        for _ in 0..10 {
774            let cluster_clone = Arc::clone(&cluster);
775            let handle = thread::spawn(move || {
776                let nodes = cluster_clone.nodes();
777                for _ in 0..100 {
778                    // Update would happen here in real usage
779                    let _ = nodes.len();
780                }
781            });
782            handles.push(handle);
783        }
784
785        for handle in handles {
786            handle.join().unwrap();
787        }
788
789        // Cluster should remain consistent
790        assert_eq!(cluster.nodes().len(), 1);
791        assert_eq!(cluster.total_vram_gb(), 24.0);
792    }
793
794    #[test]
795    fn cluster_handles_rapid_registration_churn() {
796        let cluster = Arc::new(ClusterState::new());
797
798        // Rapidly register and reregister nodes
799        for iteration in 0..5 {
800            for i in 0..10 {
801                cluster.register(NodeResources::new(
802                    format!("node-{i}"),
803                    24.0 + (iteration as f32),
804                    64.0,
805                    "8.9",
806                    None,
807                ));
808            }
809
810            let nodes = cluster.nodes();
811            assert_eq!(
812                nodes.len(),
813                10,
814                "Should have 10 nodes at iteration {iteration}"
815            );
816
817            // Verify VRAM reflects latest registration
818            let expected_vram: f32 = (0..10).map(|_| 24.0 + iteration as f32).sum();
819            let actual = cluster.total_vram_gb();
820            assert!((actual - expected_vram).abs() < 0.1);
821        }
822    }
823
824    #[test]
825    fn cluster_health_monitor_reports_accurate_status() {
826        let cluster = Arc::new(ClusterState::new());
827        cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
828        cluster.register(NodeResources::new("node-b", 12.0, 32.0, "8.6", None));
829
830        let monitor = ClusterHealthMonitor::new(cluster.clone(), Duration::from_millis(100));
831        monitor.check_health();
832
833        let report = monitor.health_report();
834        assert!(
835            report.contains("Active nodes: 2"),
836            "Should report 2 active nodes"
837        );
838        assert!(report.contains("2/2"), "Should show 2 of 2 nodes");
839    }
840
841    #[test]
842    fn node_metrics_track_min_max_latency() {
843        let mut metrics = NodeMetrics::new(24.0, 64.0, "8.9".to_string(), Duration::from_secs(5));
844
845        metrics.record_latency(2.0);
846        assert_eq!(metrics.min_latency_us, 2.0);
847        assert_eq!(metrics.max_latency_us, 2.0);
848
849        metrics.record_latency(1.0);
850        assert_eq!(metrics.min_latency_us, 1.0);
851        assert_eq!(metrics.max_latency_us, 2.0);
852
853        metrics.record_latency(5.0);
854        assert_eq!(metrics.min_latency_us, 1.0);
855        assert_eq!(metrics.max_latency_us, 5.0);
856
857        metrics.record_latency(3.0);
858        assert_eq!(metrics.min_latency_us, 1.0);
859        assert_eq!(metrics.max_latency_us, 5.0);
860    }
861}