1use std::collections::HashMap;
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
18#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
20pub enum NodeStatus {
21 #[default]
23 Active,
24 Degraded,
26 Failed,
28}
29
30#[derive(Clone, Debug)]
32pub struct NodeMetrics {
33 pub name: String,
35 pub status: NodeStatus,
37 pub vram_gb: f32,
39 pub total_vram_gb: f32,
41 pub system_memory_gb: f32,
43 pub compute_capability: String,
45 pub gpu_name: Option<String>,
47 pub ip_address: Option<SocketAddr>,
49
50 pub last_heartbeat: Instant,
52 pub heartbeat_timeout: Duration,
54
55 pub avg_latency_us: f32,
57 pub min_latency_us: f32,
59 pub max_latency_us: f32,
61 pub latency_samples: u64,
63
64 pub delivery_ratio: f32,
66 pub throughput_gbps: f32,
68
69 pub used_vram_gb: f32,
71 pub available_vram_gb: f32,
73
74 pub streaming_layers: Option<(usize, usize)>,
76
77 pub af_xdp_gbps: f32,
79 pub latency_micros: f32,
81 pub(crate) delivery_ratio_initialized: bool,
83}
84
85impl Default for NodeMetrics {
86 fn default() -> Self {
87 Self {
88 name: String::new(),
89 status: NodeStatus::Active,
90 vram_gb: 0.0,
91 total_vram_gb: 0.0,
92 system_memory_gb: 0.0,
93 compute_capability: String::new(),
94 gpu_name: None,
95 last_heartbeat: Instant::now(),
96 heartbeat_timeout: Duration::from_secs(5),
97 avg_latency_us: 0.0,
98 min_latency_us: f32::MAX,
99 max_latency_us: 0.0,
100 latency_samples: 0,
101 delivery_ratio: 1.0,
102 throughput_gbps: 0.0,
103 used_vram_gb: 0.0,
104 available_vram_gb: 0.0,
105 streaming_layers: None,
106 af_xdp_gbps: 0.0,
107 latency_micros: 0.0,
108 delivery_ratio_initialized: false,
109 ip_address: None,
110 }
111 }
112}
113
114impl NodeMetrics {
115 pub fn new(
117 vram_gb: f32,
118 system_memory_gb: f32,
119 compute_capability: String,
120 heartbeat_timeout: Duration,
121 ) -> Self {
122 Self {
123 name: String::new(),
124 status: NodeStatus::Active,
125 vram_gb,
126 total_vram_gb: vram_gb,
127 system_memory_gb,
128 compute_capability,
129 gpu_name: None,
130 ip_address: None,
131 last_heartbeat: Instant::now(),
132 heartbeat_timeout,
133 avg_latency_us: 0.0,
134 min_latency_us: f32::MAX,
135 max_latency_us: 0.0,
136 latency_samples: 0,
137 delivery_ratio: 1.0,
138 throughput_gbps: 0.0,
139 used_vram_gb: 0.0,
140 available_vram_gb: vram_gb,
141 streaming_layers: None,
142 af_xdp_gbps: 0.0,
143 latency_micros: 0.0,
144 delivery_ratio_initialized: false,
145 }
146 }
147
148 pub fn record_latency(&mut self, latency_us: f32) {
150 if latency_us < self.min_latency_us {
151 self.min_latency_us = latency_us;
152 }
153 if latency_us > self.max_latency_us {
154 self.max_latency_us = latency_us;
155 }
156
157 self.latency_samples += 1;
158 if self.latency_samples == 1 {
159 self.avg_latency_us = latency_us;
160 } else {
161 self.avg_latency_us = self.avg_latency_us * 0.9 + latency_us * 0.1;
162 }
163 }
164
165 pub fn record_delivery_ratio(&mut self, ratio: f32) {
167 if !self.delivery_ratio_initialized {
168 self.delivery_ratio = ratio;
169 self.delivery_ratio_initialized = true;
170 } else {
171 self.delivery_ratio = self.delivery_ratio * 0.9 + ratio * 0.1;
173 }
174 }
175
176 pub fn record_throughput(&mut self, throughput_gbps: f32) {
178 self.throughput_gbps = self.throughput_gbps * 0.9 + throughput_gbps * 0.1;
180 }
181
182 pub fn record_vram_usage(&mut self, used_vram_gb: f32) {
184 self.used_vram_gb = used_vram_gb;
185 self.available_vram_gb = self.total_vram_gb - used_vram_gb;
186 }
187
188 pub fn set_streaming_layers(&mut self, start: usize, end: usize) {
190 self.streaming_layers = Some((start, end));
191 }
192
193 pub fn clear_streaming_layers(&mut self) {
195 self.streaming_layers = None;
196 }
197}
198
199#[derive(Debug)]
201pub struct ClusterState {
202 nodes: Arc<Mutex<HashMap<String, NodeResources>>>,
204 nodes_snapshot: Arc<ArcSwap<Vec<NodeResources>>>,
206 nodes_snapshot_dirty: Arc<AtomicBool>,
208 pub(crate) metrics: Arc<Mutex<HashMap<String, NodeMetrics>>>,
210 last_update: Arc<AtomicU64>,
212 total_vram_cache: Arc<AtomicU64>,
214}
215
216impl Clone for ClusterState {
217 fn clone(&self) -> Self {
218 Self {
219 nodes: Arc::clone(&self.nodes),
220 nodes_snapshot: Arc::clone(&self.nodes_snapshot),
221 nodes_snapshot_dirty: Arc::clone(&self.nodes_snapshot_dirty),
222 metrics: Arc::clone(&self.metrics),
223 last_update: Arc::clone(&self.last_update),
224 total_vram_cache: Arc::clone(&self.total_vram_cache),
225 }
226 }
227}
228
229impl Default for ClusterState {
230 fn default() -> Self {
231 Self::new()
232 }
233}
234
235impl ClusterState {
236 pub fn new() -> Self {
238 Self {
239 nodes: Arc::new(Mutex::new(HashMap::new())),
240 nodes_snapshot: Arc::new(ArcSwap::from_pointee(Vec::<NodeResources>::new())),
241 nodes_snapshot_dirty: Arc::new(AtomicBool::new(false)),
242 metrics: Arc::new(Mutex::new(HashMap::new())),
243 last_update: Arc::new(AtomicU64::new(0)),
244 total_vram_cache: Arc::new(AtomicU64::new(0.0_f64.to_bits())),
245 }
246 }
247
248 pub fn register(&self, node: NodeResources) {
250 self.register_with_addr(node, None);
251 }
252
253 pub fn register_with_addr(&self, node: NodeResources, addr: Option<SocketAddr>) {
255 let mut nodes = self
256 .nodes
257 .lock()
258 .unwrap_or_else(|poison| poison.into_inner());
259 let mut metrics = self
260 .metrics
261 .lock()
262 .unwrap_or_else(|poison| poison.into_inner());
263
264 let id = node.id.clone();
265 let gpu_name = node.gpu_name.clone();
266 let vram_gb = node.vram_gb;
267 let system_memory_gb = node.system_memory_gb;
268 let compute_capability = node.compute_capability.clone();
269 let rpc_port = node.rpc_port;
270 let mut vram_delta = vram_gb;
271
272 if let Some(existing) = nodes.get_mut(&id) {
273 vram_delta = vram_gb - existing.vram_gb;
274 existing.vram_gb = vram_gb;
275 existing.system_memory_gb = system_memory_gb;
276 existing.compute_capability = compute_capability.clone();
277 existing.gpu_name = gpu_name.clone();
278 existing.rpc_port = rpc_port;
284 } else {
285 nodes.insert(id.clone(), node);
286 }
287
288 if let Some(existing_metrics) = metrics.get_mut(&id) {
289 existing_metrics.name = id.clone();
290 existing_metrics.vram_gb = vram_gb;
291 existing_metrics.total_vram_gb = vram_gb;
292 existing_metrics.system_memory_gb = system_memory_gb;
293 existing_metrics.compute_capability = compute_capability;
294 existing_metrics.gpu_name = gpu_name;
295 existing_metrics.heartbeat_timeout = Duration::from_secs(5);
296 if addr.is_some() {
297 existing_metrics.ip_address = addr;
298 }
299 } else {
300 let mut node_metrics = NodeMetrics::new(
301 vram_gb,
302 system_memory_gb,
303 compute_capability,
304 Duration::from_secs(5),
305 );
306 node_metrics.name = id.clone();
307 node_metrics.gpu_name = gpu_name;
308 node_metrics.ip_address = addr;
309 metrics.insert(id, node_metrics);
310 }
311
312 self.nodes_snapshot_dirty.store(true, Ordering::Release);
313
314 let current_total_vram = f64::from_bits(self.total_vram_cache.load(Ordering::Acquire));
315 self.total_vram_cache.store(
316 (current_total_vram + vram_delta as f64).to_bits(),
317 Ordering::Release,
318 );
319
320 self.last_update.store(
321 std::time::SystemTime::now()
322 .duration_since(std::time::UNIX_EPOCH)
323 .unwrap_or(Duration::ZERO)
324 .as_millis() as u64,
325 Ordering::Release,
326 );
327 }
328
329 pub fn nodes(&self) -> Vec<NodeResources> {
331 self.nodes_snapshot().as_ref().to_vec()
332 }
333
334 pub fn node_count(&self) -> usize {
336 self.nodes_snapshot().len()
337 }
338
339 pub fn nodes_snapshot(&self) -> Arc<Vec<NodeResources>> {
341 if self.nodes_snapshot_dirty.load(Ordering::Relaxed)
342 && self.nodes_snapshot_dirty.swap(false, Ordering::AcqRel)
343 {
344 let nodes_map = self
345 .nodes
346 .lock()
347 .unwrap_or_else(|poison| poison.into_inner());
348 let mut nodes: Vec<_> = nodes_map.values().cloned().collect();
349 nodes.sort_unstable_by(|a, b| a.id.cmp(&b.id));
351 self.nodes_snapshot.store(Arc::new(nodes));
352 }
353 self.nodes_snapshot.load_full()
354 }
355
356 pub fn get_metrics(&self, node_id: &str) -> Option<NodeMetrics> {
358 let metrics = self
359 .metrics
360 .lock()
361 .unwrap_or_else(|poison| poison.into_inner());
362 metrics.get(node_id).cloned()
363 }
364
365 pub fn update_heartbeat(&self, node_id: &str) {
367 self.get_metrics_mut(node_id, |metrics| {
368 metrics.last_heartbeat = Instant::now();
369 });
370 }
371
372 pub fn get_metrics_mut<F, R>(&self, node_id: &str, f: F) -> Option<R>
374 where
375 F: FnOnce(&mut NodeMetrics) -> R,
376 {
377 let mut metrics = self
378 .metrics
379 .lock()
380 .unwrap_or_else(|poison| poison.into_inner());
381 metrics.get_mut(node_id).map(f)
382 }
383
384 pub fn check_heartbeat_timeout(&self, node_id: &str) -> bool {
386 if let Some(metrics) = self.get_metrics(node_id) {
387 let elapsed = Instant::now().duration_since(metrics.last_heartbeat);
388 elapsed >= metrics.heartbeat_timeout
389 } else {
390 false
391 }
392 }
393
394 pub fn mark_failed(&self, node_id: &str) {
396 self.get_metrics_mut(node_id, |metrics| {
397 metrics.status = NodeStatus::Failed;
398 });
399 }
400
401 pub fn recover_node(&self, node_id: &str) {
403 self.get_metrics_mut(node_id, |metrics| {
404 metrics.status = NodeStatus::Active;
405 });
406 }
407
408 pub fn active_nodes(&self) -> Vec<NodeMetrics> {
410 let metrics = self
411 .metrics
412 .lock()
413 .unwrap_or_else(|poison| poison.into_inner());
414 metrics
415 .values()
416 .filter(|m| m.status == NodeStatus::Active)
417 .cloned()
418 .collect()
419 }
420
421 pub fn active_nodes_count(&self) -> usize {
423 let metrics = self
424 .metrics
425 .lock()
426 .unwrap_or_else(|poison| poison.into_inner());
427 metrics
428 .values()
429 .filter(|m| m.status == NodeStatus::Active)
430 .count()
431 }
432
433 pub fn total_vram_gb(&self) -> f32 {
435 f64::from_bits(self.total_vram_cache.load(Ordering::Acquire)) as f32
436 }
437
438 pub fn total_system_memory_gb(&self) -> f32 {
440 let nodes = self
441 .nodes
442 .lock()
443 .unwrap_or_else(|poison| poison.into_inner());
444 nodes.values().map(|n| n.system_memory_gb).sum()
445 }
446}
447
448#[derive(Clone, Debug)]
450pub struct ClusterHealthMonitor {
451 cluster: Arc<ClusterState>,
452 check_interval: Duration,
454}
455
456impl ClusterHealthMonitor {
457 pub fn new(cluster: Arc<ClusterState>, check_interval: Duration) -> Self {
459 Self {
460 cluster,
461 check_interval,
462 }
463 }
464
465 pub fn check_health(&self) {
467 let failed_nodes: Vec<String> = self
468 .cluster
469 .nodes_snapshot()
470 .iter()
471 .filter(|n| self.cluster.check_heartbeat_timeout(&n.id))
472 .map(|n| n.id.clone())
473 .collect();
474
475 for node_id in &failed_nodes {
477 self.cluster.mark_failed(node_id);
478 }
479
480 self.cluster.last_update.store(
482 std::time::SystemTime::now()
483 .duration_since(std::time::UNIX_EPOCH)
484 .unwrap_or(Duration::ZERO)
485 .as_millis() as u64,
486 Ordering::Release,
487 );
488 }
489
490 pub fn health_report(&self) -> String {
492 let active_count = self.cluster.active_nodes_count();
493 let total_nodes = self.cluster.nodes_snapshot().len();
494
495 format!(
496 "Cluster Health Report\n\
497 =================\n\
498 Active nodes: {}/{}\n\
499 Total VRAM: {:.1} GB\n\
500 System memory: {:.1} GB\n",
501 active_count,
502 total_nodes,
503 self.cluster.total_vram_gb(),
504 self.cluster.total_system_memory_gb()
505 )
506 }
507
508 pub fn start_periodic_checks(&self) {
510 let this = self.clone();
511
512 std::thread::spawn(move || loop {
514 this.check_health();
515 std::thread::sleep(this.check_interval);
516 });
517 }
518}
519
520#[cfg(test)]
521mod tests {
522 use super::*;
523 use std::sync::Arc;
524 use std::thread;
525 use std::time::Duration;
526
527 #[test]
528 fn register_replaces_existing_nodes() {
529 let cluster = ClusterState::new();
530 cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
531 cluster.register(NodeResources::new("node-a", 48.0, 128.0, "9.0", None));
532
533 assert_eq!(cluster.nodes().len(), 1);
534 assert_eq!(cluster.nodes()[0].vram_gb, 48.0);
535 }
536
537 #[test]
538 fn heartbeat_timeout_detection() {
539 let cluster = ClusterState::new();
540 cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
541
542 thread::sleep(Duration::from_secs(6));
544
545 assert!(cluster.check_heartbeat_timeout("node-a"));
546 }
547
548 #[test]
549 fn health_monitor_reports_active_count() {
550 let cluster = Arc::new(ClusterState::new());
551 cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
552 cluster.register(NodeResources::new("node-b", 12.0, 32.0, "8.6", None));
553
554 let monitor = ClusterHealthMonitor::new(cluster.clone(), Duration::from_secs(1));
555 monitor.check_health();
556
557 let report = monitor.health_report();
558 assert!(report.contains("Active nodes: 2/2"));
559 }
560
561 #[test]
562 fn metrics_record_latency() {
563 let mut metrics = NodeMetrics::new(24.0, 64.0, "8.9".to_string(), Duration::from_secs(5));
564
565 metrics.record_latency(1.0);
566 assert_eq!(metrics.avg_latency_us, 1.0);
567 assert_eq!(metrics.min_latency_us, 1.0);
568 assert_eq!(metrics.max_latency_us, 1.0);
569
570 metrics.record_latency(2.0);
571 assert!((metrics.avg_latency_us - 1.1).abs() < 1e-6);
573 }
574
575 #[test]
576 fn metrics_record_delivery_ratio() {
577 let mut metrics = NodeMetrics::new(24.0, 64.0, "8.9".to_string(), Duration::from_secs(5));
578
579 metrics.record_delivery_ratio(0.98);
580 assert!((metrics.delivery_ratio - 0.98).abs() < 1e-6);
581
582 metrics.record_delivery_ratio(0.90);
583 assert!((metrics.delivery_ratio - 0.972).abs() < 1e-6);
585 }
586
587 #[test]
588 fn cluster_tracks_total_vram() {
589 let cluster = ClusterState::new();
590 cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
591 cluster.register(NodeResources::new("node-b", 12.0, 32.0, "8.6", None));
592
593 assert_eq!(cluster.total_vram_gb(), 36.0);
594 }
595
596 #[test]
601 fn cluster_health_monitor_checks_all_nodes() {
602 let cluster = Arc::new(ClusterState::new());
603 cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
604 cluster.register(NodeResources::new("node-b", 12.0, 32.0, "8.6", None));
605
606 let monitor = ClusterHealthMonitor::new(cluster.clone(), Duration::from_millis(100));
607 monitor.check_health();
608
609 let report = monitor.health_report();
610 assert!(
611 report.contains("Active nodes: 2/2"),
612 "Report should show both nodes as active"
613 );
614 }
615
616 #[test]
617 fn cluster_metrics_use_exponential_moving_average() {
618 let mut metrics = NodeMetrics::new(24.0, 64.0, "8.9".to_string(), Duration::from_secs(5));
619
620 metrics.record_latency(1.0);
622 let avg1 = metrics.avg_latency_us;
623 assert_eq!(avg1, 1.0);
624
625 metrics.record_latency(2.0);
626 let avg2 = metrics.avg_latency_us;
627 assert!(
629 (avg2 - 1.1).abs() < 1e-6,
630 "EMA calculation: expected 1.1, got {avg2}"
631 );
632
633 metrics.record_latency(3.0);
634 let avg3 = metrics.avg_latency_us;
635 assert!(
637 (avg3 - 1.29).abs() < 1e-5,
638 "EMA calculation: expected 1.29, got {avg3}"
639 );
640 }
641
642 #[test]
643 fn cluster_delivery_ratio_tracks_network_quality() {
644 let mut metrics = NodeMetrics::new(24.0, 64.0, "8.9".to_string(), Duration::from_secs(5));
645
646 metrics.record_delivery_ratio(0.98);
647 assert!((metrics.delivery_ratio - 0.98).abs() < 1e-6);
648
649 metrics.record_delivery_ratio(0.95);
650 assert!((metrics.delivery_ratio - 0.977).abs() < 1e-6);
652
653 metrics.record_delivery_ratio(0.80);
655 assert!(
657 metrics.delivery_ratio < 0.98,
658 "Delivery ratio should degrade"
659 );
660 assert!(metrics.delivery_ratio > 0.90, "But recover somewhat");
661 }
662
663 #[test]
664 fn cluster_concurrent_metric_updates() {
665 let cluster = Arc::new(ClusterState::new());
666 cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
667
668 let mut handles = vec![];
669
670 for _ in 0..10 {
672 let cluster_clone = Arc::clone(&cluster);
673 let handle = thread::spawn(move || {
674 let nodes = cluster_clone.nodes();
675 for _ in 0..100 {
676 let _ = nodes.len();
678 }
679 });
680 handles.push(handle);
681 }
682
683 for handle in handles {
684 handle.join().unwrap();
685 }
686
687 assert_eq!(cluster.nodes().len(), 1);
689 assert_eq!(cluster.total_vram_gb(), 24.0);
690 }
691
692 #[test]
693 fn cluster_handles_rapid_registration_churn() {
694 let cluster = Arc::new(ClusterState::new());
695
696 for iteration in 0..5 {
698 for i in 0..10 {
699 cluster.register(NodeResources::new(
700 format!("node-{i}"),
701 24.0 + (iteration as f32),
702 64.0,
703 "8.9",
704 None,
705 ));
706 }
707
708 let nodes = cluster.nodes();
709 assert_eq!(
710 nodes.len(),
711 10,
712 "Should have 10 nodes at iteration {iteration}"
713 );
714
715 let expected_vram: f32 = (0..10).map(|_| 24.0 + iteration as f32).sum();
717 let actual = cluster.total_vram_gb();
718 assert!((actual - expected_vram).abs() < 0.1);
719 }
720 }
721
722 #[test]
723 fn cluster_health_monitor_reports_accurate_status() {
724 let cluster = Arc::new(ClusterState::new());
725 cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
726 cluster.register(NodeResources::new("node-b", 12.0, 32.0, "8.6", None));
727
728 let monitor = ClusterHealthMonitor::new(cluster.clone(), Duration::from_millis(100));
729 monitor.check_health();
730
731 let report = monitor.health_report();
732 assert!(
733 report.contains("Active nodes: 2"),
734 "Should report 2 active nodes"
735 );
736 assert!(report.contains("2/2"), "Should show 2 of 2 nodes");
737 }
738
739 #[test]
740 fn node_metrics_track_min_max_latency() {
741 let mut metrics = NodeMetrics::new(24.0, 64.0, "8.9".to_string(), Duration::from_secs(5));
742
743 metrics.record_latency(2.0);
744 assert_eq!(metrics.min_latency_us, 2.0);
745 assert_eq!(metrics.max_latency_us, 2.0);
746
747 metrics.record_latency(1.0);
748 assert_eq!(metrics.min_latency_us, 1.0);
749 assert_eq!(metrics.max_latency_us, 2.0);
750
751 metrics.record_latency(5.0);
752 assert_eq!(metrics.min_latency_us, 1.0);
753 assert_eq!(metrics.max_latency_us, 5.0);
754
755 metrics.record_latency(3.0);
756 assert_eq!(metrics.min_latency_us, 1.0);
757 assert_eq!(metrics.max_latency_us, 5.0);
758 }
759}