use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::time::{Duration, SystemTime};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ConnectionQuality {
Excellent,
Good,
Fair,
Poor,
VeryPoor,
}
impl ConnectionQuality {
#[must_use]
pub const fn from_latency(latency_ms: u64) -> Self {
if latency_ms < 50 {
Self::Excellent
} else if latency_ms < 150 {
Self::Good
} else if latency_ms < 300 {
Self::Fair
} else if latency_ms < 500 {
Self::Poor
} else {
Self::VeryPoor
}
}
#[must_use]
pub const fn health_score(&self) -> f64 {
match self {
Self::Excellent => 1.0,
Self::Good => 0.8,
Self::Fair => 0.6,
Self::Poor => 0.4,
Self::VeryPoor => 0.2,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NetworkStats {
pub peer_id: String,
pub avg_latency_ms: f64,
pub min_latency_ms: u64,
pub max_latency_ms: u64,
pub latency_stddev: f64,
pub sample_count: u64,
pub packet_loss_rate: f64,
pub estimated_bandwidth_bps: f64,
pub quality: ConnectionQuality,
pub last_update: SystemTime,
}
impl NetworkStats {
fn new(peer_id: String) -> Self {
Self {
peer_id,
avg_latency_ms: 0.0,
min_latency_ms: u64::MAX,
max_latency_ms: 0,
latency_stddev: 0.0,
sample_count: 0,
packet_loss_rate: 0.0,
estimated_bandwidth_bps: 0.0,
quality: ConnectionQuality::Good,
last_update: SystemTime::now(),
}
}
#[must_use]
#[inline]
pub fn health_score(&self) -> f64 {
let latency_score = self.quality.health_score();
let packet_loss_penalty = self.packet_loss_rate;
let stability_score = if self.sample_count > 10 {
1.0 - (self.latency_stddev / 100.0).min(1.0)
} else {
0.5
};
(latency_score * 0.5 + (1.0 - packet_loss_penalty) * 0.3 + stability_score * 0.2)
.clamp(0.0, 1.0)
}
#[must_use]
#[inline]
pub const fn is_stable(&self) -> bool {
self.latency_stddev < 50.0 && self.packet_loss_rate < 0.05
}
#[must_use]
#[inline]
pub fn is_healthy(&self) -> bool {
self.health_score() >= 0.6
}
}
pub struct NetworkMonitor {
stats: HashMap<String, NetworkStats>,
max_history: usize,
latency_history: HashMap<String, Vec<u64>>,
}
impl NetworkMonitor {
#[must_use]
#[inline]
pub fn new() -> Self {
Self::with_history_size(100)
}
#[must_use]
#[inline]
pub fn with_history_size(max_history: usize) -> Self {
Self {
stats: HashMap::new(),
max_history,
latency_history: HashMap::new(),
}
}
pub fn record_latency(&mut self, peer_id: String, latency_ms: u64) {
let stats = self
.stats
.entry(peer_id.clone())
.or_insert_with(|| NetworkStats::new(peer_id.clone()));
stats.sample_count += 1;
stats.min_latency_ms = stats.min_latency_ms.min(latency_ms);
stats.max_latency_ms = stats.max_latency_ms.max(latency_ms);
stats.last_update = SystemTime::now();
let alpha = 0.3;
if stats.sample_count == 1 {
stats.avg_latency_ms = latency_ms as f64;
} else {
stats.avg_latency_ms = alpha * latency_ms as f64 + (1.0 - alpha) * stats.avg_latency_ms;
}
stats.quality = ConnectionQuality::from_latency(stats.avg_latency_ms as u64);
let history = self.latency_history.entry(peer_id).or_default();
history.push(latency_ms);
if history.len() > self.max_history {
history.remove(0);
}
if history.len() > 1 {
let mean = history.iter().sum::<u64>() as f64 / history.len() as f64;
let variance = history
.iter()
.map(|&x| {
let diff = x as f64 - mean;
diff * diff
})
.sum::<f64>()
/ history.len() as f64;
stats.latency_stddev = variance.sqrt();
}
}
pub fn record_bandwidth(&mut self, peer_id: &str, bytes_per_sec: f64) {
if let Some(stats) = self.stats.get_mut(peer_id) {
let alpha = 0.2;
if stats.estimated_bandwidth_bps == 0.0 {
stats.estimated_bandwidth_bps = bytes_per_sec;
} else {
stats.estimated_bandwidth_bps =
alpha * bytes_per_sec + (1.0 - alpha) * stats.estimated_bandwidth_bps;
}
}
}
pub fn record_packet_loss(&mut self, peer_id: &str, lost_packets: u64, total_packets: u64) {
if let Some(stats) = self.stats.get_mut(peer_id) {
let loss_rate = lost_packets as f64 / total_packets as f64;
let alpha = 0.3;
stats.packet_loss_rate = alpha * loss_rate + (1.0 - alpha) * stats.packet_loss_rate;
}
}
#[must_use]
#[inline]
pub fn get_stats(&self, peer_id: &str) -> Option<&NetworkStats> {
self.stats.get(peer_id)
}
#[must_use]
#[inline]
pub fn get_quality(&self, peer_id: &str) -> ConnectionQuality {
self.stats
.get(peer_id)
.map(|s| s.quality)
.unwrap_or(ConnectionQuality::Good)
}
#[must_use]
#[inline]
pub fn health_score(&self, peer_id: &str) -> f64 {
self.stats
.get(peer_id)
.map(|s| s.health_score())
.unwrap_or(0.5)
}
#[must_use]
#[inline]
pub fn get_healthy_peers(&self) -> Vec<String> {
self.stats
.values()
.filter(|s| s.is_healthy())
.map(|s| s.peer_id.clone())
.collect()
}
#[must_use]
#[inline]
pub fn get_excellent_peers(&self) -> Vec<String> {
self.stats
.values()
.filter(|s| s.quality == ConnectionQuality::Excellent)
.map(|s| s.peer_id.clone())
.collect()
}
#[must_use]
#[inline]
pub fn average_health(&self) -> f64 {
if self.stats.is_empty() {
return 0.5;
}
let sum: f64 = self.stats.values().map(|s| s.health_score()).sum();
sum / self.stats.len() as f64
}
pub fn cleanup_old_peers(&mut self, max_age_secs: u64) {
let now = SystemTime::now();
self.stats.retain(|peer_id, stats| {
let age = now
.duration_since(stats.last_update)
.unwrap_or(Duration::from_secs(0))
.as_secs();
if age >= max_age_secs {
self.latency_history.remove(peer_id);
false
} else {
true
}
});
}
#[must_use]
#[inline]
pub fn peer_count(&self) -> usize {
self.stats.len()
}
#[must_use]
#[inline]
pub fn get_all_peer_ids(&self) -> Vec<String> {
self.stats.keys().cloned().collect()
}
}
impl Default for NetworkMonitor {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_connection_quality_from_latency() {
assert_eq!(
ConnectionQuality::from_latency(30),
ConnectionQuality::Excellent
);
assert_eq!(
ConnectionQuality::from_latency(100),
ConnectionQuality::Good
);
assert_eq!(
ConnectionQuality::from_latency(200),
ConnectionQuality::Fair
);
assert_eq!(
ConnectionQuality::from_latency(400),
ConnectionQuality::Poor
);
assert_eq!(
ConnectionQuality::from_latency(600),
ConnectionQuality::VeryPoor
);
}
#[test]
fn test_connection_quality_health_score() {
assert_eq!(ConnectionQuality::Excellent.health_score(), 1.0);
assert_eq!(ConnectionQuality::Good.health_score(), 0.8);
assert_eq!(ConnectionQuality::Fair.health_score(), 0.6);
assert_eq!(ConnectionQuality::Poor.health_score(), 0.4);
assert_eq!(ConnectionQuality::VeryPoor.health_score(), 0.2);
}
#[test]
fn test_network_monitor_basic() {
let mut monitor = NetworkMonitor::new();
monitor.record_latency("peer1".to_string(), 50);
monitor.record_latency("peer1".to_string(), 55);
monitor.record_latency("peer1".to_string(), 45);
let stats = monitor.get_stats("peer1").unwrap();
assert_eq!(stats.sample_count, 3);
assert_eq!(stats.min_latency_ms, 45);
assert_eq!(stats.max_latency_ms, 55);
}
#[test]
fn test_network_monitor_quality() {
let mut monitor = NetworkMonitor::new();
monitor.record_latency("peer1".to_string(), 30);
assert_eq!(monitor.get_quality("peer1"), ConnectionQuality::Excellent);
monitor.record_latency("peer2".to_string(), 200);
assert_eq!(monitor.get_quality("peer2"), ConnectionQuality::Fair);
}
#[test]
fn test_bandwidth_tracking() {
let mut monitor = NetworkMonitor::new();
monitor.record_latency("peer1".to_string(), 50);
monitor.record_bandwidth("peer1", 1_000_000.0);
monitor.record_bandwidth("peer1", 1_500_000.0);
let stats = monitor.get_stats("peer1").unwrap();
assert!(stats.estimated_bandwidth_bps > 0.0);
}
#[test]
fn test_packet_loss() {
let mut monitor = NetworkMonitor::new();
monitor.record_latency("peer1".to_string(), 50);
monitor.record_packet_loss("peer1", 5, 100);
let stats = monitor.get_stats("peer1").unwrap();
assert!(stats.packet_loss_rate > 0.0);
}
#[test]
fn test_health_score() {
let mut monitor = NetworkMonitor::new();
for _ in 0..10 {
monitor.record_latency("peer1".to_string(), 40);
}
let score = monitor.health_score("peer1");
assert!(score > 0.8);
}
#[test]
fn test_get_healthy_peers() {
let mut monitor = NetworkMonitor::new();
monitor.record_latency("peer1".to_string(), 40);
monitor.record_latency("peer2".to_string(), 600);
monitor.record_latency("peer3".to_string(), 80);
let healthy = monitor.get_healthy_peers();
assert!(healthy.contains(&"peer1".to_string()));
assert!(healthy.contains(&"peer3".to_string()));
}
#[test]
fn test_get_excellent_peers() {
let mut monitor = NetworkMonitor::new();
monitor.record_latency("peer1".to_string(), 30);
monitor.record_latency("peer2".to_string(), 100);
let excellent = monitor.get_excellent_peers();
assert_eq!(excellent.len(), 1);
assert_eq!(excellent[0], "peer1");
}
#[test]
fn test_average_health() {
let mut monitor = NetworkMonitor::new();
monitor.record_latency("peer1".to_string(), 40);
monitor.record_latency("peer2".to_string(), 150);
let avg = monitor.average_health();
assert!(avg > 0.5 && avg < 1.0);
}
#[test]
fn test_cleanup_old_peers() {
let mut monitor = NetworkMonitor::new();
monitor.record_latency("peer1".to_string(), 50);
assert_eq!(monitor.peer_count(), 1);
monitor.cleanup_old_peers(0);
assert_eq!(monitor.peer_count(), 0);
}
#[test]
fn test_network_stats_stability() {
let mut monitor = NetworkMonitor::new();
for _ in 0..20 {
monitor.record_latency("peer1".to_string(), 50);
}
let stats = monitor.get_stats("peer1").unwrap();
assert!(stats.is_stable());
for i in 0..20 {
monitor.record_latency("peer2".to_string(), 50 + (i * 20));
}
let stats2 = monitor.get_stats("peer2").unwrap();
assert!(!stats2.is_stable());
}
}