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