use dashmap::DashMap;
use libp2p::PeerId;
use serde::{Deserialize, Serialize};
use std::collections::VecDeque;
use std::sync::Arc;
use std::time::Instant;
use thiserror::Error;
#[derive(Debug, Error)]
pub enum QualityPredictorError {
#[error("No historical data available for peer")]
NoHistoricalData,
#[error("Invalid configuration: {0}")]
InvalidConfig(String),
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QualityPredictorConfig {
pub max_samples: usize,
pub latency_weight: f64,
pub bandwidth_weight: f64,
pub reliability_weight: f64,
pub uptime_weight: f64,
pub smoothing_factor: f64,
pub min_acceptable_quality: f64,
pub switch_threshold: f64,
pub enable_predictions: bool,
}
impl Default for QualityPredictorConfig {
fn default() -> Self {
Self {
max_samples: 100,
latency_weight: 0.3,
bandwidth_weight: 0.3,
reliability_weight: 0.25,
uptime_weight: 0.15,
smoothing_factor: 0.2,
min_acceptable_quality: 0.5,
switch_threshold: 0.6,
enable_predictions: true,
}
}
}
impl QualityPredictorConfig {
pub fn low_latency() -> Self {
Self {
latency_weight: 0.5,
bandwidth_weight: 0.2,
reliability_weight: 0.2,
uptime_weight: 0.1,
..Default::default()
}
}
pub fn high_bandwidth() -> Self {
Self {
latency_weight: 0.15,
bandwidth_weight: 0.5,
reliability_weight: 0.25,
uptime_weight: 0.1,
..Default::default()
}
}
pub fn high_reliability() -> Self {
Self {
latency_weight: 0.2,
bandwidth_weight: 0.2,
reliability_weight: 0.4,
uptime_weight: 0.2,
..Default::default()
}
}
pub fn validate(&self) -> Result<(), QualityPredictorError> {
if self.max_samples == 0 {
return Err(QualityPredictorError::InvalidConfig(
"max_samples must be > 0".to_string(),
));
}
let total_weight = self.latency_weight
+ self.bandwidth_weight
+ self.reliability_weight
+ self.uptime_weight;
if (total_weight - 1.0).abs() > 0.01 {
return Err(QualityPredictorError::InvalidConfig(format!(
"weights must sum to 1.0, got {}",
total_weight
)));
}
if !(0.0..=1.0).contains(&self.smoothing_factor) {
return Err(QualityPredictorError::InvalidConfig(
"smoothing_factor must be between 0.0 and 1.0".to_string(),
));
}
Ok(())
}
}
#[derive(Debug, Clone)]
struct ConnectionHistory {
latency_samples: VecDeque<u64>,
bandwidth_samples: VecDeque<u64>,
success_count: u64,
failure_count: u64,
first_seen: Instant,
last_seen: Instant,
quality_ema: Option<f64>,
}
impl ConnectionHistory {
fn new() -> Self {
let now = Instant::now();
Self {
latency_samples: VecDeque::new(),
bandwidth_samples: VecDeque::new(),
success_count: 0,
failure_count: 0,
first_seen: now,
last_seen: now,
quality_ema: None,
}
}
fn record_latency(&mut self, latency_ms: u64, max_samples: usize) {
if self.latency_samples.len() >= max_samples {
self.latency_samples.pop_front();
}
self.latency_samples.push_back(latency_ms);
self.last_seen = Instant::now();
}
fn record_bandwidth(&mut self, bytes_per_sec: u64, max_samples: usize) {
if self.bandwidth_samples.len() >= max_samples {
self.bandwidth_samples.pop_front();
}
self.bandwidth_samples.push_back(bytes_per_sec);
self.last_seen = Instant::now();
}
fn record_success(&mut self) {
self.success_count += 1;
self.last_seen = Instant::now();
}
fn record_failure(&mut self) {
self.failure_count += 1;
self.last_seen = Instant::now();
}
fn avg_latency(&self) -> Option<f64> {
if self.latency_samples.is_empty() {
None
} else {
let sum: u64 = self.latency_samples.iter().sum();
Some(sum as f64 / self.latency_samples.len() as f64)
}
}
fn avg_bandwidth(&self) -> Option<f64> {
if self.bandwidth_samples.is_empty() {
None
} else {
let sum: u64 = self.bandwidth_samples.iter().sum();
Some(sum as f64 / self.bandwidth_samples.len() as f64)
}
}
fn reliability_score(&self) -> f64 {
let total = self.success_count + self.failure_count;
if total == 0 {
0.5 } else {
self.success_count as f64 / total as f64
}
}
fn uptime_score(&self) -> f64 {
let total_duration = self.first_seen.elapsed().as_secs_f64();
if total_duration < 1.0 {
1.0 } else {
let active_duration = self.last_seen.duration_since(self.first_seen).as_secs_f64();
(active_duration / total_duration).min(1.0)
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QualityPrediction {
pub overall_score: f64,
pub latency_score: f64,
pub bandwidth_score: f64,
pub reliability_score: f64,
pub uptime_score: f64,
pub avg_latency_ms: Option<f64>,
pub avg_bandwidth_bps: Option<f64>,
pub is_acceptable: bool,
pub should_switch: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QualityPredictorStats {
pub tracked_peers: usize,
pub predictions_made: u64,
pub switch_recommendations: u64,
pub avg_quality: f64,
}
pub struct QualityPredictor {
config: QualityPredictorConfig,
history: Arc<DashMap<PeerId, ConnectionHistory>>,
stats: Arc<parking_lot::RwLock<QualityPredictorStats>>,
}
impl QualityPredictor {
pub fn new(config: QualityPredictorConfig) -> Result<Self, QualityPredictorError> {
config.validate()?;
Ok(Self {
config,
history: Arc::new(DashMap::new()),
stats: Arc::new(parking_lot::RwLock::new(QualityPredictorStats {
tracked_peers: 0,
predictions_made: 0,
switch_recommendations: 0,
avg_quality: 0.0,
})),
})
}
pub fn record_latency(&self, peer: PeerId, latency_ms: u64) {
let mut entry = self
.history
.entry(peer)
.or_insert_with(ConnectionHistory::new);
entry.record_latency(latency_ms, self.config.max_samples);
}
pub fn record_bandwidth(&self, peer: PeerId, bytes_per_sec: u64) {
let mut entry = self
.history
.entry(peer)
.or_insert_with(ConnectionHistory::new);
entry.record_bandwidth(bytes_per_sec, self.config.max_samples);
}
pub fn record_success(&self, peer: PeerId) {
let mut entry = self
.history
.entry(peer)
.or_insert_with(ConnectionHistory::new);
entry.record_success();
}
pub fn record_failure(&self, peer: PeerId) {
let mut entry = self
.history
.entry(peer)
.or_insert_with(ConnectionHistory::new);
entry.record_failure();
}
pub fn predict_quality(&self, peer: &PeerId) -> Option<QualityPrediction> {
let history = self.history.get(peer)?;
let latency_score = self.calculate_latency_score(history.avg_latency());
let bandwidth_score = self.calculate_bandwidth_score(history.avg_bandwidth());
let reliability_score = history.reliability_score();
let uptime_score = history.uptime_score();
let overall_score = latency_score * self.config.latency_weight
+ bandwidth_score * self.config.bandwidth_weight
+ reliability_score * self.config.reliability_weight
+ uptime_score * self.config.uptime_weight;
drop(history);
if let Some(mut history) = self.history.get_mut(peer) {
if let Some(prev_ema) = history.quality_ema {
history.quality_ema = Some(
self.config.smoothing_factor * overall_score
+ (1.0 - self.config.smoothing_factor) * prev_ema,
);
} else {
history.quality_ema = Some(overall_score);
}
}
let is_acceptable = overall_score >= self.config.min_acceptable_quality;
let should_switch =
self.config.enable_predictions && overall_score < self.config.switch_threshold;
let mut stats = self.stats.write();
stats.predictions_made += 1;
if should_switch {
stats.switch_recommendations += 1;
}
Some(QualityPrediction {
overall_score,
latency_score,
bandwidth_score,
reliability_score,
uptime_score,
avg_latency_ms: self.history.get(peer).and_then(|h| h.avg_latency()),
avg_bandwidth_bps: self.history.get(peer).and_then(|h| h.avg_bandwidth()),
is_acceptable,
should_switch,
})
}
pub fn should_switch_connection(&self, peer: &PeerId) -> bool {
self.predict_quality(peer)
.map(|p| p.should_switch)
.unwrap_or(false)
}
pub fn get_best_peer(&self, peers: &[PeerId]) -> Option<(PeerId, QualityPrediction)> {
peers
.iter()
.filter_map(|peer| {
self.predict_quality(peer)
.map(|prediction| (*peer, prediction))
})
.max_by(|a, b| {
a.1.overall_score
.partial_cmp(&b.1.overall_score)
.unwrap_or(std::cmp::Ordering::Equal)
})
}
pub fn rank_peers(&self, peers: &[PeerId]) -> Vec<(PeerId, QualityPrediction)> {
let mut ranked: Vec<_> = peers
.iter()
.filter_map(|peer| {
self.predict_quality(peer)
.map(|prediction| (*peer, prediction))
})
.collect();
ranked.sort_by(|a, b| {
b.1.overall_score
.partial_cmp(&a.1.overall_score)
.unwrap_or(std::cmp::Ordering::Equal)
});
ranked
}
pub fn remove_peer(&self, peer: &PeerId) {
self.history.remove(peer);
}
pub fn clear(&self) {
self.history.clear();
let mut stats = self.stats.write();
stats.tracked_peers = 0;
stats.predictions_made = 0;
stats.switch_recommendations = 0;
stats.avg_quality = 0.0;
}
pub fn stats(&self) -> QualityPredictorStats {
let mut stats = self.stats.read().clone();
stats.tracked_peers = self.history.len();
if stats.tracked_peers > 0 {
let total_quality: f64 = self
.history
.iter()
.filter_map(|entry| entry.quality_ema)
.sum();
stats.avg_quality = total_quality / stats.tracked_peers as f64;
}
stats
}
fn calculate_latency_score(&self, avg_latency: Option<f64>) -> f64 {
match avg_latency {
None => 0.5, Some(latency) => {
if latency <= 0.0 {
1.0
} else if latency >= 1000.0 {
0.0
} else {
1.0 - (latency / 1000.0)
}
}
}
}
fn calculate_bandwidth_score(&self, avg_bandwidth: Option<f64>) -> f64 {
match avg_bandwidth {
None => 0.5, Some(bandwidth) => {
let mb_per_sec = bandwidth / 1_000_000.0;
if mb_per_sec >= 100.0 {
1.0
} else if mb_per_sec <= 0.0 {
0.0
} else {
(mb_per_sec / 100.0).min(1.0)
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_config_default() {
let config = QualityPredictorConfig::default();
assert!(config.validate().is_ok());
}
#[test]
fn test_config_validation_weights() {
let config = QualityPredictorConfig {
latency_weight: 0.5,
bandwidth_weight: 0.3,
reliability_weight: 0.1,
uptime_weight: 0.05, ..Default::default()
};
assert!(config.validate().is_err());
}
#[test]
fn test_config_presets() {
assert!(QualityPredictorConfig::low_latency().validate().is_ok());
assert!(QualityPredictorConfig::high_bandwidth().validate().is_ok());
assert!(QualityPredictorConfig::high_reliability()
.validate()
.is_ok());
}
#[test]
fn test_record_metrics() {
let config = QualityPredictorConfig::default();
let predictor = QualityPredictor::new(config)
.expect("test: QualityPredictor::new should succeed with default config");
let peer = PeerId::random();
predictor.record_latency(peer, 50);
predictor.record_bandwidth(peer, 1_000_000);
predictor.record_success(peer);
let prediction = predictor
.predict_quality(&peer)
.expect("test: predict_quality should return Some after recording metrics");
assert!(prediction.avg_latency_ms.is_some());
assert!(prediction.avg_bandwidth_bps.is_some());
assert!(prediction.overall_score > 0.0);
}
#[test]
fn test_latency_score() {
let predictor = QualityPredictor::new(QualityPredictorConfig::default())
.expect("test: QualityPredictor::new should succeed with default config");
assert_eq!(predictor.calculate_latency_score(Some(0.0)), 1.0);
assert!(predictor.calculate_latency_score(Some(100.0)) > 0.7);
assert!(predictor.calculate_latency_score(Some(500.0)) < 0.6);
assert_eq!(predictor.calculate_latency_score(Some(1000.0)), 0.0);
}
#[test]
fn test_bandwidth_score() {
let predictor = QualityPredictor::new(QualityPredictorConfig::default())
.expect("test: QualityPredictor::new should succeed with default config");
assert_eq!(predictor.calculate_bandwidth_score(Some(0.0)), 0.0);
assert!(predictor.calculate_bandwidth_score(Some(1_000_000.0)) > 0.0);
assert!(predictor.calculate_bandwidth_score(Some(10_000_000.0)) > 0.05);
assert_eq!(
predictor.calculate_bandwidth_score(Some(100_000_000.0)),
1.0
);
}
#[test]
fn test_reliability_tracking() {
let config = QualityPredictorConfig::default();
let predictor = QualityPredictor::new(config)
.expect("test: QualityPredictor::new should succeed with default config");
let peer = PeerId::random();
predictor.record_success(peer);
predictor.record_success(peer);
predictor.record_failure(peer);
let prediction = predictor
.predict_quality(&peer)
.expect("test: predict_quality should return Some after recording success/failure");
assert!((prediction.reliability_score - 0.666).abs() < 0.01);
}
#[test]
fn test_get_best_peer() {
let config = QualityPredictorConfig::default();
let predictor = QualityPredictor::new(config)
.expect("test: QualityPredictor::new should succeed with default config");
let peer1 = PeerId::random();
let peer2 = PeerId::random();
let peer3 = PeerId::random();
predictor.record_latency(peer1, 10);
predictor.record_bandwidth(peer1, 10_000_000);
predictor.record_success(peer1);
predictor.record_latency(peer2, 500);
predictor.record_bandwidth(peer2, 100_000);
predictor.record_failure(peer2);
predictor.record_latency(peer3, 50);
predictor.record_bandwidth(peer3, 5_000_000);
predictor.record_success(peer3);
let peers = vec![peer1, peer2, peer3];
let (best, _) = predictor
.get_best_peer(&peers)
.expect("test: get_best_peer should return Some for non-empty peer list");
assert_eq!(best, peer1);
}
#[test]
fn test_rank_peers() {
let config = QualityPredictorConfig::default();
let predictor = QualityPredictor::new(config)
.expect("test: QualityPredictor::new should succeed with default config");
let peer1 = PeerId::random();
let peer2 = PeerId::random();
let peer3 = PeerId::random();
predictor.record_latency(peer1, 10);
predictor.record_latency(peer2, 100);
predictor.record_latency(peer3, 50);
let peers = vec![peer1, peer2, peer3];
let ranked = predictor.rank_peers(&peers);
assert_eq!(ranked.len(), 3);
assert_eq!(ranked[0].0, peer1); assert_eq!(ranked[2].0, peer2); }
#[test]
fn test_should_switch() {
let config = QualityPredictorConfig {
switch_threshold: 0.7,
enable_predictions: true,
..Default::default()
};
let predictor = QualityPredictor::new(config)
.expect("test: QualityPredictor::new should succeed with default config");
let peer = PeerId::random();
predictor.record_latency(peer, 800);
predictor.record_bandwidth(peer, 50_000);
predictor.record_failure(peer);
predictor.record_failure(peer);
predictor.record_success(peer);
assert!(predictor.should_switch_connection(&peer));
}
#[test]
fn test_ema_smoothing() {
let config = QualityPredictorConfig {
smoothing_factor: 0.5,
..Default::default()
};
let predictor = QualityPredictor::new(config)
.expect("test: QualityPredictor::new should succeed with default config");
let peer = PeerId::random();
predictor.record_latency(peer, 100);
let pred1 = predictor
.predict_quality(&peer)
.expect("test: predict_quality should return Some after recording latency");
predictor.record_latency(peer, 50);
let pred2 = predictor
.predict_quality(&peer)
.expect("test: predict_quality should return Some after recording second latency");
assert!(pred2.overall_score > pred1.overall_score);
}
#[test]
fn test_stats() {
let config = QualityPredictorConfig::default();
let predictor = QualityPredictor::new(config)
.expect("test: QualityPredictor::new should succeed with default config");
let peer1 = PeerId::random();
let peer2 = PeerId::random();
predictor.record_latency(peer1, 50);
predictor.record_latency(peer2, 100);
predictor.predict_quality(&peer1);
predictor.predict_quality(&peer2);
let stats = predictor.stats();
assert_eq!(stats.tracked_peers, 2);
assert_eq!(stats.predictions_made, 2);
}
#[test]
fn test_remove_peer() {
let config = QualityPredictorConfig::default();
let predictor = QualityPredictor::new(config)
.expect("test: QualityPredictor::new should succeed with default config");
let peer = PeerId::random();
predictor.record_latency(peer, 50);
assert!(predictor.predict_quality(&peer).is_some());
predictor.remove_peer(&peer);
assert!(predictor.predict_quality(&peer).is_none());
}
#[test]
fn test_clear() {
let config = QualityPredictorConfig::default();
let predictor = QualityPredictor::new(config)
.expect("test: QualityPredictor::new should succeed with default config");
predictor.record_latency(PeerId::random(), 50);
predictor.record_latency(PeerId::random(), 100);
predictor.clear();
let stats = predictor.stats();
assert_eq!(stats.tracked_peers, 0);
}
}