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: Zero-allocation node re-registration path.
311        // Avoid cloning `node` or its string fields when capability resources are unchanged,
312        // move owned `node` directly into existing map entries on resource updates,
313        // and skip atomic cache operations when VRAM and memory deltas are zero.
314        let vram_gb = node.vram_gb;
315        let system_memory_gb = node.system_memory_gb;
316        let mut vram_delta = vram_gb;
317        let mut system_memory_delta = system_memory_gb;
318        let mut resources_changed = true;
319
320        // Optimization: Zero-allocation, single-pass node registration path.
321        // Updates metrics and resources in-place using zero-cost ownership moves (no node.clone()),
322        // eliminating secondary HashMap lookups while avoiding atomic cache invalidations
323        // when resource capabilities remain unchanged.
324        if let Some(existing_metrics) = metrics.get_mut(&node.id) {
325            existing_metrics.vram_gb = vram_gb;
326            existing_metrics.total_vram_gb = vram_gb;
327            existing_metrics.system_memory_gb = system_memory_gb;
328            if existing_metrics.compute_capability != node.compute_capability {
329                existing_metrics
330                    .compute_capability
331                    .clone_from(&node.compute_capability);
332            }
333            if existing_metrics.gpu_name != node.gpu_name {
334                existing_metrics.gpu_name.clone_from(&node.gpu_name);
335            }
336            existing_metrics.heartbeat_timeout = Duration::from_secs(5);
337            if addr.is_some() {
338                existing_metrics.ip_address = addr;
339            }
340        } else {
341            let node_id = node.id.clone();
342            let mut node_metrics = NodeMetrics::new(
343                vram_gb,
344                system_memory_gb,
345                node.compute_capability.clone(),
346                Duration::from_secs(5),
347            );
348            node_metrics.name = node_id.clone();
349            node_metrics.gpu_name = node.gpu_name.clone();
350            node_metrics.ip_address = addr;
351            metrics.insert(node_id, node_metrics);
352        }
353
354        if let Some(existing) = nodes.get_mut(&node.id) {
355            vram_delta = vram_gb - existing.vram_gb;
356            system_memory_delta = system_memory_gb - existing.system_memory_gb;
357
358            // Check if node resource fields actually changed to avoid snapshot invalidation churn
359            if existing.vram_gb == node.vram_gb
360                && existing.system_memory_gb == node.system_memory_gb
361                && existing.compute_capability == node.compute_capability
362                && existing.gpu_name == node.gpu_name
363                && existing.rpc_port == node.rpc_port
364                && existing.rpc_build_id == node.rpc_build_id
365            {
366                resources_changed = false;
367            } else {
368                *existing = node; // Zero-cost ownership move into existing entry
369            }
370        } else {
371            let node_id = node.id.clone();
372            nodes.insert(node_id, node); // Zero-cost ownership move into new map entry
373        }
374
375        if resources_changed {
376            self.nodes_snapshot_dirty.store(true, Ordering::Release);
377        }
378
379        if vram_delta != 0.0 {
380            let current_total_vram = f64::from_bits(self.total_vram_cache.load(Ordering::Acquire));
381            self.total_vram_cache.store(
382                (current_total_vram + vram_delta as f64).to_bits(),
383                Ordering::Release,
384            );
385        }
386
387        if system_memory_delta != 0.0 {
388            let current_total_system_mem =
389                f64::from_bits(self.total_system_memory_cache.load(Ordering::Acquire));
390            self.total_system_memory_cache.store(
391                (current_total_system_mem + system_memory_delta as f64).to_bits(),
392                Ordering::Release,
393            );
394        }
395
396        self.last_update.store(
397            std::time::SystemTime::now()
398                .duration_since(std::time::UNIX_EPOCH)
399                .unwrap_or(Duration::ZERO)
400                .as_millis() as u64,
401            Ordering::Release,
402        );
403    }
404
405    /// Get all nodes
406    pub fn nodes(&self) -> Vec<NodeResources> {
407        self.nodes_snapshot().as_ref().to_vec()
408    }
409
410    /// Get total number of registered nodes.
411    pub fn node_count(&self) -> usize {
412        self.nodes_snapshot().len()
413    }
414
415    /// Get a shared snapshot of all nodes
416    pub fn nodes_snapshot(&self) -> Arc<Vec<NodeResources>> {
417        if self.nodes_snapshot_dirty.load(Ordering::Relaxed)
418            && self.nodes_snapshot_dirty.swap(false, Ordering::AcqRel)
419        {
420            let nodes_map = self
421                .nodes
422                .lock()
423                .unwrap_or_else(|poison| poison.into_inner());
424            let mut nodes: Vec<_> = nodes_map.values().cloned().collect();
425            // Sort by ID to ensure deterministic layer assignment across runs
426            nodes.sort_unstable_by(|a, b| a.id.cmp(&b.id));
427            self.nodes_snapshot.store(Arc::new(nodes));
428        }
429        self.nodes_snapshot.load_full()
430    }
431
432    /// Execute a closure with a read lock on the metrics map, avoiding clones and multiple acquisitions.
433    pub fn with_metrics<F, R>(&self, f: F) -> R
434    where
435        F: FnOnce(&HashMap<String, NodeMetrics>) -> R,
436    {
437        let metrics = self
438            .metrics
439            .lock()
440            .unwrap_or_else(|poison| poison.into_inner());
441        f(&metrics)
442    }
443
444    /// Get metrics for a specific node
445    pub fn get_metrics(&self, node_id: &str) -> Option<NodeMetrics> {
446        let metrics = self
447            .metrics
448            .lock()
449            .unwrap_or_else(|poison| poison.into_inner());
450        metrics.get(node_id).cloned()
451    }
452
453    /// Update last heartbeat for a node
454    pub fn update_heartbeat(&self, node_id: &str) {
455        self.get_metrics_mut(node_id, |metrics| {
456            metrics.last_heartbeat = Instant::now();
457        });
458    }
459
460    /// Get metrics mutable reference (for internal updates)
461    pub fn get_metrics_mut<F, R>(&self, node_id: &str, f: F) -> Option<R>
462    where
463        F: FnOnce(&mut NodeMetrics) -> R,
464    {
465        let mut metrics = self
466            .metrics
467            .lock()
468            .unwrap_or_else(|poison| poison.into_inner());
469        metrics.get_mut(node_id).map(f)
470    }
471
472    /// Check if a node has timed out
473    pub fn check_heartbeat_timeout(&self, node_id: &str) -> bool {
474        // Optimize: Inspect node metrics directly under lock without cloning NodeMetrics struct.
475        let metrics = self
476            .metrics
477            .lock()
478            .unwrap_or_else(|poison| poison.into_inner());
479        if let Some(m) = metrics.get(node_id) {
480            Instant::now().duration_since(m.last_heartbeat) >= m.heartbeat_timeout
481        } else {
482            false
483        }
484    }
485
486    /// Mark a node as failed due to timeout
487    pub fn mark_failed(&self, node_id: &str) {
488        self.get_metrics_mut(node_id, |metrics| {
489            metrics.status = NodeStatus::Failed;
490        });
491    }
492
493    /// Recover a failed node
494    pub fn recover_node(&self, node_id: &str) {
495        self.get_metrics_mut(node_id, |metrics| {
496            metrics.status = NodeStatus::Active;
497        });
498    }
499
500    /// Get all active nodes
501    pub fn active_nodes(&self) -> Vec<NodeMetrics> {
502        let metrics = self
503            .metrics
504            .lock()
505            .unwrap_or_else(|poison| poison.into_inner());
506        metrics
507            .values()
508            .filter(|m| m.status == NodeStatus::Active)
509            .cloned()
510            .collect()
511    }
512
513    /// Get the count of active nodes without cloning metrics
514    pub fn active_nodes_count(&self) -> usize {
515        let metrics = self
516            .metrics
517            .lock()
518            .unwrap_or_else(|poison| poison.into_inner());
519        metrics
520            .values()
521            .filter(|m| m.status == NodeStatus::Active)
522            .count()
523    }
524
525    /// Get total cluster VRAM
526    pub fn total_vram_gb(&self) -> f32 {
527        f64::from_bits(self.total_vram_cache.load(Ordering::Acquire)) as f32
528    }
529
530    /// Get total system memory
531    pub fn total_system_memory_gb(&self) -> f32 {
532        f64::from_bits(self.total_system_memory_cache.load(Ordering::Acquire)) as f32
533    }
534}
535
536/// Cluster health monitor with periodic checks
537#[derive(Clone, Debug)]
538pub struct ClusterHealthMonitor {
539    cluster: Arc<ClusterState>,
540    /// Health check interval
541    check_interval: Duration,
542}
543
544impl ClusterHealthMonitor {
545    /// Create new health monitor
546    pub fn new(cluster: Arc<ClusterState>, check_interval: Duration) -> Self {
547        Self {
548            cluster,
549            check_interval,
550        }
551    }
552
553    /// Run health check on all nodes
554    pub fn check_health(&self) {
555        let failed_nodes: Vec<String> = self
556            .cluster
557            .nodes_snapshot()
558            .iter()
559            .filter(|n| self.cluster.check_heartbeat_timeout(&n.id))
560            .map(|n| n.id.clone())
561            .collect();
562
563        // Mark timed-out nodes as failed
564        for node_id in &failed_nodes {
565            self.cluster.mark_failed(node_id);
566        }
567
568        // Update last update timestamp
569        self.cluster.last_update.store(
570            std::time::SystemTime::now()
571                .duration_since(std::time::UNIX_EPOCH)
572                .unwrap_or(Duration::ZERO)
573                .as_millis() as u64,
574            Ordering::Release,
575        );
576    }
577
578    /// Get health report
579    pub fn health_report(&self) -> String {
580        let active_count = self.cluster.active_nodes_count();
581        let total_nodes = self.cluster.nodes_snapshot().len();
582
583        format!(
584            "Cluster Health Report\n\
585             =================\n\
586             Active nodes: {}/{}\n\
587             Total VRAM: {:.1} GB\n\
588             System memory: {:.1} GB\n",
589            active_count,
590            total_nodes,
591            self.cluster.total_vram_gb(),
592            self.cluster.total_system_memory_gb()
593        )
594    }
595
596    /// Run periodic health checks in background
597    pub fn start_periodic_checks(&self) {
598        let this = self.clone();
599
600        // Spawn health check task
601        std::thread::spawn(move || loop {
602            this.check_health();
603            std::thread::sleep(this.check_interval);
604        });
605    }
606}
607
608#[cfg(test)]
609mod tests {
610    use super::*;
611    use std::sync::Arc;
612    use std::thread;
613    use std::time::Duration;
614
615    #[test]
616    fn register_replaces_existing_nodes() {
617        let cluster = ClusterState::new();
618        cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
619        cluster.register(NodeResources::new("node-a", 48.0, 128.0, "9.0", None));
620
621        assert_eq!(cluster.nodes().len(), 1);
622        assert_eq!(cluster.nodes()[0].vram_gb, 48.0);
623    }
624
625    #[test]
626    fn heartbeat_timeout_detection() {
627        let cluster = ClusterState::new();
628        cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
629
630        // Simulate timeout by waiting longer than default heartbeat timeout
631        thread::sleep(Duration::from_secs(6));
632
633        assert!(cluster.check_heartbeat_timeout("node-a"));
634    }
635
636    #[test]
637    fn node_metrics_keep_recent_latency_and_throughput_history() {
638        let mut metrics = NodeMetrics::default();
639
640        for idx in 0..40 {
641            metrics.record_latency(idx as f32);
642            metrics.record_throughput(idx as f32 / 10.0);
643        }
644
645        assert_eq!(metrics.latency_history_us.len(), NODE_HISTORY_CAP);
646        assert_eq!(metrics.throughput_history_gbps.len(), NODE_HISTORY_CAP);
647        assert_eq!(metrics.latency_history_us.front().copied(), Some(8.0));
648        assert_eq!(metrics.throughput_history_gbps.back().copied(), Some(3.9));
649    }
650
651    #[test]
652    fn health_monitor_reports_active_count() {
653        let cluster = Arc::new(ClusterState::new());
654        cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
655        cluster.register(NodeResources::new("node-b", 12.0, 32.0, "8.6", None));
656
657        let monitor = ClusterHealthMonitor::new(cluster.clone(), Duration::from_secs(1));
658        monitor.check_health();
659
660        let report = monitor.health_report();
661        assert!(report.contains("Active nodes: 2/2"));
662    }
663
664    #[test]
665    fn metrics_record_latency() {
666        let mut metrics = NodeMetrics::new(24.0, 64.0, "8.9".to_string(), Duration::from_secs(5));
667
668        metrics.record_latency(1.0);
669        assert_eq!(metrics.avg_latency_us, 1.0);
670        assert_eq!(metrics.min_latency_us, 1.0);
671        assert_eq!(metrics.max_latency_us, 1.0);
672
673        metrics.record_latency(2.0);
674        // EMA with alpha=0.1: (1.0 * 0.9) + 2.0 * 0.1 = 0.9 + 0.2 = 1.1
675        assert!((metrics.avg_latency_us - 1.1).abs() < 1e-6);
676    }
677
678    #[test]
679    fn metrics_record_delivery_ratio() {
680        let mut metrics = NodeMetrics::new(24.0, 64.0, "8.9".to_string(), Duration::from_secs(5));
681
682        metrics.record_delivery_ratio(0.98);
683        assert!((metrics.delivery_ratio - 0.98).abs() < 1e-6);
684
685        metrics.record_delivery_ratio(0.90);
686        // EMA: 0.98 * 0.9 + 0.90 * 0.1 = 0.882 + 0.09 = 0.972
687        assert!((metrics.delivery_ratio - 0.972).abs() < 1e-6);
688    }
689
690    #[test]
691    fn with_metrics_allows_single_pass_inspection() {
692        let cluster = ClusterState::new();
693        cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
694        cluster.register(NodeResources::new("node-b", 12.0, 32.0, "8.6", None));
695
696        let total_nodes = cluster.with_metrics(|metrics_map| {
697            assert!(metrics_map.contains_key("node-a"));
698            assert!(metrics_map.contains_key("node-b"));
699            metrics_map.len()
700        });
701
702        assert_eq!(total_nodes, 2);
703    }
704
705    #[test]
706    fn cluster_tracks_total_vram() {
707        let cluster = ClusterState::new();
708        cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
709        cluster.register(NodeResources::new("node-b", 12.0, 32.0, "8.6", None));
710
711        assert_eq!(cluster.total_vram_gb(), 36.0);
712    }
713
714    // ========================================================================
715    // NEW: Comprehensive Health Monitoring & Failure Recovery Tests
716    // ========================================================================
717
718    #[test]
719    fn cluster_health_monitor_checks_all_nodes() {
720        let cluster = Arc::new(ClusterState::new());
721        cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
722        cluster.register(NodeResources::new("node-b", 12.0, 32.0, "8.6", None));
723
724        let monitor = ClusterHealthMonitor::new(cluster.clone(), Duration::from_millis(100));
725        monitor.check_health();
726
727        let report = monitor.health_report();
728        assert!(
729            report.contains("Active nodes: 2/2"),
730            "Report should show both nodes as active"
731        );
732    }
733
734    #[test]
735    fn cluster_metrics_use_exponential_moving_average() {
736        let mut metrics = NodeMetrics::new(24.0, 64.0, "8.9".to_string(), Duration::from_secs(5));
737
738        // Record sequence of latencies: 1, 2, 3
739        metrics.record_latency(1.0);
740        let avg1 = metrics.avg_latency_us;
741        assert_eq!(avg1, 1.0);
742
743        metrics.record_latency(2.0);
744        let avg2 = metrics.avg_latency_us;
745        // EMA: 1.0 * 0.9 + 2.0 * 0.1 = 1.1
746        assert!(
747            (avg2 - 1.1).abs() < 1e-6,
748            "EMA calculation: expected 1.1, got {avg2}"
749        );
750
751        metrics.record_latency(3.0);
752        let avg3 = metrics.avg_latency_us;
753        // EMA: 1.1 * 0.9 + 3.0 * 0.1 = 0.99 + 0.3 = 1.29
754        assert!(
755            (avg3 - 1.29).abs() < 1e-5,
756            "EMA calculation: expected 1.29, got {avg3}"
757        );
758    }
759
760    #[test]
761    fn cluster_delivery_ratio_tracks_network_quality() {
762        let mut metrics = NodeMetrics::new(24.0, 64.0, "8.9".to_string(), Duration::from_secs(5));
763
764        metrics.record_delivery_ratio(0.98);
765        assert!((metrics.delivery_ratio - 0.98).abs() < 1e-6);
766
767        metrics.record_delivery_ratio(0.95);
768        // EMA: 0.98 * 0.9 + 0.95 * 0.1 = 0.882 + 0.095 = 0.977
769        assert!((metrics.delivery_ratio - 0.977).abs() < 1e-6);
770
771        // Simulate degradation
772        metrics.record_delivery_ratio(0.80);
773        // EMA: 0.977 * 0.9 + 0.80 * 0.1 = 0.8793 + 0.08 = 0.9593
774        assert!(
775            metrics.delivery_ratio < 0.98,
776            "Delivery ratio should degrade"
777        );
778        assert!(metrics.delivery_ratio > 0.90, "But recover somewhat");
779    }
780
781    #[test]
782    fn cluster_concurrent_metric_updates() {
783        let cluster = Arc::new(ClusterState::new());
784        cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
785
786        let mut handles = vec![];
787
788        // Spawn 10 threads updating metrics concurrently
789        for _ in 0..10 {
790            let cluster_clone = Arc::clone(&cluster);
791            let handle = thread::spawn(move || {
792                let nodes = cluster_clone.nodes();
793                for _ in 0..100 {
794                    // Update would happen here in real usage
795                    let _ = nodes.len();
796                }
797            });
798            handles.push(handle);
799        }
800
801        for handle in handles {
802            handle.join().unwrap();
803        }
804
805        // Cluster should remain consistent
806        assert_eq!(cluster.nodes().len(), 1);
807        assert_eq!(cluster.total_vram_gb(), 24.0);
808    }
809
810    #[test]
811    fn cluster_handles_rapid_registration_churn() {
812        let cluster = Arc::new(ClusterState::new());
813
814        // Rapidly register and reregister nodes
815        for iteration in 0..5 {
816            for i in 0..10 {
817                cluster.register(NodeResources::new(
818                    format!("node-{i}"),
819                    24.0 + (iteration as f32),
820                    64.0,
821                    "8.9",
822                    None,
823                ));
824            }
825
826            let nodes = cluster.nodes();
827            assert_eq!(
828                nodes.len(),
829                10,
830                "Should have 10 nodes at iteration {iteration}"
831            );
832
833            // Verify VRAM reflects latest registration
834            let expected_vram: f32 = (0..10).map(|_| 24.0 + iteration as f32).sum();
835            let actual = cluster.total_vram_gb();
836            assert!((actual - expected_vram).abs() < 0.1);
837        }
838    }
839
840    #[test]
841    fn cluster_health_monitor_reports_accurate_status() {
842        let cluster = Arc::new(ClusterState::new());
843        cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
844        cluster.register(NodeResources::new("node-b", 12.0, 32.0, "8.6", None));
845
846        let monitor = ClusterHealthMonitor::new(cluster.clone(), Duration::from_millis(100));
847        monitor.check_health();
848
849        let report = monitor.health_report();
850        assert!(
851            report.contains("Active nodes: 2"),
852            "Should report 2 active nodes"
853        );
854        assert!(report.contains("2/2"), "Should show 2 of 2 nodes");
855    }
856
857    #[test]
858    fn node_metrics_track_min_max_latency() {
859        let mut metrics = NodeMetrics::new(24.0, 64.0, "8.9".to_string(), Duration::from_secs(5));
860
861        metrics.record_latency(2.0);
862        assert_eq!(metrics.min_latency_us, 2.0);
863        assert_eq!(metrics.max_latency_us, 2.0);
864
865        metrics.record_latency(1.0);
866        assert_eq!(metrics.min_latency_us, 1.0);
867        assert_eq!(metrics.max_latency_us, 2.0);
868
869        metrics.record_latency(5.0);
870        assert_eq!(metrics.min_latency_us, 1.0);
871        assert_eq!(metrics.max_latency_us, 5.0);
872
873        metrics.record_latency(3.0);
874        assert_eq!(metrics.min_latency_us, 1.0);
875        assert_eq!(metrics.max_latency_us, 5.0);
876    }
877}