use std::collections::{HashMap, VecDeque};
use std::net::SocketAddr;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use arc_swap::ArcSwap;
pub use crate::protocol::NodeResources;
const NODE_HISTORY_CAP: usize = 32;
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum NodeStatus {
#[default]
Active,
Degraded,
Failed,
}
#[derive(Clone, Debug)]
pub struct NodeMetrics {
pub name: String,
pub status: NodeStatus,
pub vram_gb: f32,
pub total_vram_gb: f32,
pub system_memory_gb: f32,
pub compute_capability: String,
pub gpu_name: Option<String>,
pub ip_address: Option<SocketAddr>,
pub last_heartbeat: Instant,
pub heartbeat_timeout: Duration,
pub avg_latency_us: f32,
pub min_latency_us: f32,
pub max_latency_us: f32,
pub latency_samples: u64,
pub delivery_ratio: f32,
pub throughput_gbps: f32,
pub latency_history_us: VecDeque<f32>,
pub throughput_history_gbps: VecDeque<f32>,
pub used_vram_gb: f32,
pub available_vram_gb: f32,
pub streaming_layers: Option<(usize, usize)>,
pub af_xdp_gbps: f32,
pub latency_micros: f32,
pub(crate) delivery_ratio_initialized: bool,
}
impl Default for NodeMetrics {
fn default() -> Self {
Self {
name: String::new(),
status: NodeStatus::Active,
vram_gb: 0.0,
total_vram_gb: 0.0,
system_memory_gb: 0.0,
compute_capability: String::new(),
gpu_name: None,
last_heartbeat: Instant::now(),
heartbeat_timeout: Duration::from_secs(5),
avg_latency_us: 0.0,
min_latency_us: f32::MAX,
max_latency_us: 0.0,
latency_samples: 0,
delivery_ratio: 1.0,
throughput_gbps: 0.0,
latency_history_us: VecDeque::with_capacity(NODE_HISTORY_CAP),
throughput_history_gbps: VecDeque::with_capacity(NODE_HISTORY_CAP),
used_vram_gb: 0.0,
available_vram_gb: 0.0,
streaming_layers: None,
af_xdp_gbps: 0.0,
latency_micros: 0.0,
delivery_ratio_initialized: false,
ip_address: None,
}
}
}
impl NodeMetrics {
pub fn new(
vram_gb: f32,
system_memory_gb: f32,
compute_capability: String,
heartbeat_timeout: Duration,
) -> Self {
Self {
name: String::new(),
status: NodeStatus::Active,
vram_gb,
total_vram_gb: vram_gb,
system_memory_gb,
compute_capability,
gpu_name: None,
ip_address: None,
last_heartbeat: Instant::now(),
heartbeat_timeout,
avg_latency_us: 0.0,
min_latency_us: f32::MAX,
max_latency_us: 0.0,
latency_samples: 0,
delivery_ratio: 1.0,
throughput_gbps: 0.0,
latency_history_us: VecDeque::with_capacity(NODE_HISTORY_CAP),
throughput_history_gbps: VecDeque::with_capacity(NODE_HISTORY_CAP),
used_vram_gb: 0.0,
available_vram_gb: vram_gb,
streaming_layers: None,
af_xdp_gbps: 0.0,
latency_micros: 0.0,
delivery_ratio_initialized: false,
}
}
pub fn record_latency(&mut self, latency_us: f32) {
if latency_us < self.min_latency_us {
self.min_latency_us = latency_us;
}
if latency_us > self.max_latency_us {
self.max_latency_us = latency_us;
}
self.latency_samples += 1;
if self.latency_samples == 1 {
self.avg_latency_us = latency_us;
} else {
self.avg_latency_us = self.avg_latency_us * 0.9 + latency_us * 0.1;
}
if self.latency_history_us.len() >= NODE_HISTORY_CAP {
self.latency_history_us.pop_front();
}
self.latency_history_us.push_back(latency_us);
}
pub fn record_delivery_ratio(&mut self, ratio: f32) {
if !self.delivery_ratio_initialized {
self.delivery_ratio = ratio;
self.delivery_ratio_initialized = true;
} else {
self.delivery_ratio = self.delivery_ratio * 0.9 + ratio * 0.1;
}
}
pub fn record_throughput(&mut self, throughput_gbps: f32) {
self.throughput_gbps = self.throughput_gbps * 0.9 + throughput_gbps * 0.1;
if self.throughput_history_gbps.len() >= NODE_HISTORY_CAP {
self.throughput_history_gbps.pop_front();
}
self.throughput_history_gbps.push_back(throughput_gbps);
}
pub fn record_vram_usage(&mut self, used_vram_gb: f32) {
self.used_vram_gb = used_vram_gb;
self.available_vram_gb = self.total_vram_gb - used_vram_gb;
}
pub fn set_streaming_layers(&mut self, start: usize, end: usize) {
self.streaming_layers = Some((start, end));
}
pub fn clear_streaming_layers(&mut self) {
self.streaming_layers = None;
}
}
#[derive(Debug)]
pub struct ClusterState {
nodes: Arc<Mutex<HashMap<String, NodeResources>>>,
nodes_snapshot: Arc<ArcSwap<Vec<NodeResources>>>,
nodes_snapshot_dirty: Arc<AtomicBool>,
pub(crate) metrics: Arc<Mutex<HashMap<String, NodeMetrics>>>,
last_update: Arc<AtomicU64>,
total_vram_cache: Arc<AtomicU64>,
total_system_memory_cache: Arc<AtomicU64>,
circuit_breakers: Arc<Mutex<HashMap<String, crate::circuit_breaker::CircuitBreaker>>>,
}
impl Clone for ClusterState {
fn clone(&self) -> Self {
Self {
nodes: Arc::clone(&self.nodes),
nodes_snapshot: Arc::clone(&self.nodes_snapshot),
nodes_snapshot_dirty: Arc::clone(&self.nodes_snapshot_dirty),
metrics: Arc::clone(&self.metrics),
last_update: Arc::clone(&self.last_update),
total_vram_cache: Arc::clone(&self.total_vram_cache),
total_system_memory_cache: Arc::clone(&self.total_system_memory_cache),
circuit_breakers: Arc::clone(&self.circuit_breakers),
}
}
}
impl Default for ClusterState {
fn default() -> Self {
Self::new()
}
}
impl ClusterState {
pub fn new() -> Self {
Self {
nodes: Arc::new(Mutex::new(HashMap::new())),
nodes_snapshot: Arc::new(ArcSwap::from_pointee(Vec::<NodeResources>::new())),
nodes_snapshot_dirty: Arc::new(AtomicBool::new(false)),
metrics: Arc::new(Mutex::new(HashMap::new())),
last_update: Arc::new(AtomicU64::new(0)),
total_vram_cache: Arc::new(AtomicU64::new(0.0_f64.to_bits())),
total_system_memory_cache: Arc::new(AtomicU64::new(0.0_f64.to_bits())),
circuit_breakers: Arc::new(Mutex::new(HashMap::new())),
}
}
pub fn circuit_breaker_for(&self, node_id: &str) -> crate::circuit_breaker::CircuitBreaker {
let mut breakers = self
.circuit_breakers
.lock()
.unwrap_or_else(|poison| poison.into_inner());
breakers.entry(node_id.to_string()).or_default().clone()
}
pub fn register(&self, node: NodeResources) {
self.register_with_addr(node, None);
}
pub fn register_with_addr(&self, node: NodeResources, addr: Option<SocketAddr>) {
let mut nodes = self
.nodes
.lock()
.unwrap_or_else(|poison| poison.into_inner());
let mut metrics = self
.metrics
.lock()
.unwrap_or_else(|poison| poison.into_inner());
let id = node.id.clone();
let gpu_name = node.gpu_name.clone();
let vram_gb = node.vram_gb;
let system_memory_gb = node.system_memory_gb;
let compute_capability = node.compute_capability.clone();
let rpc_port = node.rpc_port;
let mut vram_delta = vram_gb;
let mut system_memory_delta = system_memory_gb;
if let Some(existing) = nodes.get_mut(&id) {
vram_delta = vram_gb - existing.vram_gb;
system_memory_delta = system_memory_gb - existing.system_memory_gb;
existing.vram_gb = vram_gb;
existing.system_memory_gb = system_memory_gb;
existing.compute_capability = compute_capability.clone();
existing.gpu_name = gpu_name.clone();
existing.rpc_port = rpc_port;
} else {
nodes.insert(id.clone(), node);
}
if let Some(existing_metrics) = metrics.get_mut(&id) {
existing_metrics.name = id.clone();
existing_metrics.vram_gb = vram_gb;
existing_metrics.total_vram_gb = vram_gb;
existing_metrics.system_memory_gb = system_memory_gb;
existing_metrics.compute_capability = compute_capability;
existing_metrics.gpu_name = gpu_name;
existing_metrics.heartbeat_timeout = Duration::from_secs(5);
if addr.is_some() {
existing_metrics.ip_address = addr;
}
} else {
let mut node_metrics = NodeMetrics::new(
vram_gb,
system_memory_gb,
compute_capability,
Duration::from_secs(5),
);
node_metrics.name = id.clone();
node_metrics.gpu_name = gpu_name;
node_metrics.ip_address = addr;
metrics.insert(id, node_metrics);
}
self.nodes_snapshot_dirty.store(true, Ordering::Release);
let current_total_vram = f64::from_bits(self.total_vram_cache.load(Ordering::Acquire));
self.total_vram_cache.store(
(current_total_vram + vram_delta as f64).to_bits(),
Ordering::Release,
);
let current_total_system_mem =
f64::from_bits(self.total_system_memory_cache.load(Ordering::Acquire));
self.total_system_memory_cache.store(
(current_total_system_mem + system_memory_delta as f64).to_bits(),
Ordering::Release,
);
self.last_update.store(
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or(Duration::ZERO)
.as_millis() as u64,
Ordering::Release,
);
}
pub fn nodes(&self) -> Vec<NodeResources> {
self.nodes_snapshot().as_ref().to_vec()
}
pub fn node_count(&self) -> usize {
self.nodes_snapshot().len()
}
pub fn nodes_snapshot(&self) -> Arc<Vec<NodeResources>> {
if self.nodes_snapshot_dirty.load(Ordering::Relaxed)
&& self.nodes_snapshot_dirty.swap(false, Ordering::AcqRel)
{
let nodes_map = self
.nodes
.lock()
.unwrap_or_else(|poison| poison.into_inner());
let mut nodes: Vec<_> = nodes_map.values().cloned().collect();
nodes.sort_unstable_by(|a, b| a.id.cmp(&b.id));
self.nodes_snapshot.store(Arc::new(nodes));
}
self.nodes_snapshot.load_full()
}
pub fn get_metrics(&self, node_id: &str) -> Option<NodeMetrics> {
let metrics = self
.metrics
.lock()
.unwrap_or_else(|poison| poison.into_inner());
metrics.get(node_id).cloned()
}
pub fn update_heartbeat(&self, node_id: &str) {
self.get_metrics_mut(node_id, |metrics| {
metrics.last_heartbeat = Instant::now();
});
}
pub fn get_metrics_mut<F, R>(&self, node_id: &str, f: F) -> Option<R>
where
F: FnOnce(&mut NodeMetrics) -> R,
{
let mut metrics = self
.metrics
.lock()
.unwrap_or_else(|poison| poison.into_inner());
metrics.get_mut(node_id).map(f)
}
pub fn check_heartbeat_timeout(&self, node_id: &str) -> bool {
if let Some(metrics) = self.get_metrics(node_id) {
let elapsed = Instant::now().duration_since(metrics.last_heartbeat);
elapsed >= metrics.heartbeat_timeout
} else {
false
}
}
pub fn mark_failed(&self, node_id: &str) {
self.get_metrics_mut(node_id, |metrics| {
metrics.status = NodeStatus::Failed;
});
}
pub fn recover_node(&self, node_id: &str) {
self.get_metrics_mut(node_id, |metrics| {
metrics.status = NodeStatus::Active;
});
}
pub fn active_nodes(&self) -> Vec<NodeMetrics> {
let metrics = self
.metrics
.lock()
.unwrap_or_else(|poison| poison.into_inner());
metrics
.values()
.filter(|m| m.status == NodeStatus::Active)
.cloned()
.collect()
}
pub fn active_nodes_count(&self) -> usize {
let metrics = self
.metrics
.lock()
.unwrap_or_else(|poison| poison.into_inner());
metrics
.values()
.filter(|m| m.status == NodeStatus::Active)
.count()
}
pub fn total_vram_gb(&self) -> f32 {
f64::from_bits(self.total_vram_cache.load(Ordering::Acquire)) as f32
}
pub fn total_system_memory_gb(&self) -> f32 {
f64::from_bits(self.total_system_memory_cache.load(Ordering::Acquire)) as f32
}
}
#[derive(Clone, Debug)]
pub struct ClusterHealthMonitor {
cluster: Arc<ClusterState>,
check_interval: Duration,
}
impl ClusterHealthMonitor {
pub fn new(cluster: Arc<ClusterState>, check_interval: Duration) -> Self {
Self {
cluster,
check_interval,
}
}
pub fn check_health(&self) {
let failed_nodes: Vec<String> = self
.cluster
.nodes_snapshot()
.iter()
.filter(|n| self.cluster.check_heartbeat_timeout(&n.id))
.map(|n| n.id.clone())
.collect();
for node_id in &failed_nodes {
self.cluster.mark_failed(node_id);
}
self.cluster.last_update.store(
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or(Duration::ZERO)
.as_millis() as u64,
Ordering::Release,
);
}
pub fn health_report(&self) -> String {
let active_count = self.cluster.active_nodes_count();
let total_nodes = self.cluster.nodes_snapshot().len();
format!(
"Cluster Health Report\n\
=================\n\
Active nodes: {}/{}\n\
Total VRAM: {:.1} GB\n\
System memory: {:.1} GB\n",
active_count,
total_nodes,
self.cluster.total_vram_gb(),
self.cluster.total_system_memory_gb()
)
}
pub fn start_periodic_checks(&self) {
let this = self.clone();
std::thread::spawn(move || loop {
this.check_health();
std::thread::sleep(this.check_interval);
});
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
use std::thread;
use std::time::Duration;
#[test]
fn register_replaces_existing_nodes() {
let cluster = ClusterState::new();
cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
cluster.register(NodeResources::new("node-a", 48.0, 128.0, "9.0", None));
assert_eq!(cluster.nodes().len(), 1);
assert_eq!(cluster.nodes()[0].vram_gb, 48.0);
}
#[test]
fn heartbeat_timeout_detection() {
let cluster = ClusterState::new();
cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
thread::sleep(Duration::from_secs(6));
assert!(cluster.check_heartbeat_timeout("node-a"));
}
#[test]
fn node_metrics_keep_recent_latency_and_throughput_history() {
let mut metrics = NodeMetrics::default();
for idx in 0..40 {
metrics.record_latency(idx as f32);
metrics.record_throughput(idx as f32 / 10.0);
}
assert_eq!(metrics.latency_history_us.len(), NODE_HISTORY_CAP);
assert_eq!(metrics.throughput_history_gbps.len(), NODE_HISTORY_CAP);
assert_eq!(metrics.latency_history_us.front().copied(), Some(8.0));
assert_eq!(metrics.throughput_history_gbps.back().copied(), Some(3.9));
}
#[test]
fn health_monitor_reports_active_count() {
let cluster = Arc::new(ClusterState::new());
cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
cluster.register(NodeResources::new("node-b", 12.0, 32.0, "8.6", None));
let monitor = ClusterHealthMonitor::new(cluster.clone(), Duration::from_secs(1));
monitor.check_health();
let report = monitor.health_report();
assert!(report.contains("Active nodes: 2/2"));
}
#[test]
fn metrics_record_latency() {
let mut metrics = NodeMetrics::new(24.0, 64.0, "8.9".to_string(), Duration::from_secs(5));
metrics.record_latency(1.0);
assert_eq!(metrics.avg_latency_us, 1.0);
assert_eq!(metrics.min_latency_us, 1.0);
assert_eq!(metrics.max_latency_us, 1.0);
metrics.record_latency(2.0);
assert!((metrics.avg_latency_us - 1.1).abs() < 1e-6);
}
#[test]
fn metrics_record_delivery_ratio() {
let mut metrics = NodeMetrics::new(24.0, 64.0, "8.9".to_string(), Duration::from_secs(5));
metrics.record_delivery_ratio(0.98);
assert!((metrics.delivery_ratio - 0.98).abs() < 1e-6);
metrics.record_delivery_ratio(0.90);
assert!((metrics.delivery_ratio - 0.972).abs() < 1e-6);
}
#[test]
fn cluster_tracks_total_vram() {
let cluster = ClusterState::new();
cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
cluster.register(NodeResources::new("node-b", 12.0, 32.0, "8.6", None));
assert_eq!(cluster.total_vram_gb(), 36.0);
}
#[test]
fn cluster_health_monitor_checks_all_nodes() {
let cluster = Arc::new(ClusterState::new());
cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
cluster.register(NodeResources::new("node-b", 12.0, 32.0, "8.6", None));
let monitor = ClusterHealthMonitor::new(cluster.clone(), Duration::from_millis(100));
monitor.check_health();
let report = monitor.health_report();
assert!(
report.contains("Active nodes: 2/2"),
"Report should show both nodes as active"
);
}
#[test]
fn cluster_metrics_use_exponential_moving_average() {
let mut metrics = NodeMetrics::new(24.0, 64.0, "8.9".to_string(), Duration::from_secs(5));
metrics.record_latency(1.0);
let avg1 = metrics.avg_latency_us;
assert_eq!(avg1, 1.0);
metrics.record_latency(2.0);
let avg2 = metrics.avg_latency_us;
assert!(
(avg2 - 1.1).abs() < 1e-6,
"EMA calculation: expected 1.1, got {avg2}"
);
metrics.record_latency(3.0);
let avg3 = metrics.avg_latency_us;
assert!(
(avg3 - 1.29).abs() < 1e-5,
"EMA calculation: expected 1.29, got {avg3}"
);
}
#[test]
fn cluster_delivery_ratio_tracks_network_quality() {
let mut metrics = NodeMetrics::new(24.0, 64.0, "8.9".to_string(), Duration::from_secs(5));
metrics.record_delivery_ratio(0.98);
assert!((metrics.delivery_ratio - 0.98).abs() < 1e-6);
metrics.record_delivery_ratio(0.95);
assert!((metrics.delivery_ratio - 0.977).abs() < 1e-6);
metrics.record_delivery_ratio(0.80);
assert!(
metrics.delivery_ratio < 0.98,
"Delivery ratio should degrade"
);
assert!(metrics.delivery_ratio > 0.90, "But recover somewhat");
}
#[test]
fn cluster_concurrent_metric_updates() {
let cluster = Arc::new(ClusterState::new());
cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
let mut handles = vec![];
for _ in 0..10 {
let cluster_clone = Arc::clone(&cluster);
let handle = thread::spawn(move || {
let nodes = cluster_clone.nodes();
for _ in 0..100 {
let _ = nodes.len();
}
});
handles.push(handle);
}
for handle in handles {
handle.join().unwrap();
}
assert_eq!(cluster.nodes().len(), 1);
assert_eq!(cluster.total_vram_gb(), 24.0);
}
#[test]
fn cluster_handles_rapid_registration_churn() {
let cluster = Arc::new(ClusterState::new());
for iteration in 0..5 {
for i in 0..10 {
cluster.register(NodeResources::new(
format!("node-{i}"),
24.0 + (iteration as f32),
64.0,
"8.9",
None,
));
}
let nodes = cluster.nodes();
assert_eq!(
nodes.len(),
10,
"Should have 10 nodes at iteration {iteration}"
);
let expected_vram: f32 = (0..10).map(|_| 24.0 + iteration as f32).sum();
let actual = cluster.total_vram_gb();
assert!((actual - expected_vram).abs() < 0.1);
}
}
#[test]
fn cluster_health_monitor_reports_accurate_status() {
let cluster = Arc::new(ClusterState::new());
cluster.register(NodeResources::new("node-a", 24.0, 64.0, "8.9", None));
cluster.register(NodeResources::new("node-b", 12.0, 32.0, "8.6", None));
let monitor = ClusterHealthMonitor::new(cluster.clone(), Duration::from_millis(100));
monitor.check_health();
let report = monitor.health_report();
assert!(
report.contains("Active nodes: 2"),
"Should report 2 active nodes"
);
assert!(report.contains("2/2"), "Should show 2 of 2 nodes");
}
#[test]
fn node_metrics_track_min_max_latency() {
let mut metrics = NodeMetrics::new(24.0, 64.0, "8.9".to_string(), Duration::from_secs(5));
metrics.record_latency(2.0);
assert_eq!(metrics.min_latency_us, 2.0);
assert_eq!(metrics.max_latency_us, 2.0);
metrics.record_latency(1.0);
assert_eq!(metrics.min_latency_us, 1.0);
assert_eq!(metrics.max_latency_us, 2.0);
metrics.record_latency(5.0);
assert_eq!(metrics.min_latency_us, 1.0);
assert_eq!(metrics.max_latency_us, 5.0);
metrics.record_latency(3.0);
assert_eq!(metrics.min_latency_us, 1.0);
assert_eq!(metrics.max_latency_us, 5.0);
}
}