use crate::geo_routing::{GeoLocation, GeoRouter, GeoRouterConfig};
use crate::quality_predictor::{QualityPredictor, QualityPredictorConfig};
use dashmap::DashMap;
use libp2p::PeerId;
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::time::{Duration, Instant};
#[derive(Debug, Clone)]
pub struct PeerSelectorConfig {
pub distance_weight: f64,
pub quality_weight: f64,
pub latency_weight: f64,
pub bandwidth_weight: f64,
pub enable_caching: bool,
pub cache_ttl_secs: u64,
pub max_cache_entries: usize,
}
impl Default for PeerSelectorConfig {
fn default() -> Self {
Self {
distance_weight: 0.3,
quality_weight: 0.3,
latency_weight: 0.2,
bandwidth_weight: 0.2,
enable_caching: true,
cache_ttl_secs: 300, max_cache_entries: 1000,
}
}
}
impl PeerSelectorConfig {
pub fn low_latency() -> Self {
Self {
distance_weight: 0.4,
quality_weight: 0.1,
latency_weight: 0.4,
bandwidth_weight: 0.1,
enable_caching: true,
cache_ttl_secs: 180,
max_cache_entries: 500,
}
}
pub fn high_bandwidth() -> Self {
Self {
distance_weight: 0.1,
quality_weight: 0.2,
latency_weight: 0.1,
bandwidth_weight: 0.6,
enable_caching: true,
cache_ttl_secs: 300,
max_cache_entries: 1000,
}
}
pub fn balanced() -> Self {
Self::default()
}
pub fn mobile() -> Self {
Self {
distance_weight: 0.5, quality_weight: 0.3,
latency_weight: 0.1,
bandwidth_weight: 0.1,
enable_caching: true,
cache_ttl_secs: 600, max_cache_entries: 200,
}
}
}
#[derive(Debug, Clone)]
pub struct SelectionCriteria {
pub reference_location: Option<GeoLocation>,
pub min_quality_score: f64,
pub max_distance_km: Option<f64>,
pub max_results: usize,
}
impl Default for SelectionCriteria {
fn default() -> Self {
Self {
reference_location: None,
min_quality_score: 0.0,
max_distance_km: None,
max_results: 10,
}
}
}
#[derive(Debug, Clone)]
pub struct SelectedPeer {
pub peer_id: PeerId,
pub score: f64,
pub distance_score: f64,
pub quality_score: f64,
pub latency_score: f64,
pub bandwidth_score: f64,
pub location: Option<GeoLocation>,
pub distance_km: Option<f64>,
}
#[derive(Debug, Clone)]
struct CachedSelection {
peers: Vec<SelectedPeer>,
cached_at: Instant,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct PeerSelectorStats {
pub total_selections: u64,
pub cache_hits: u64,
pub cache_misses: u64,
pub total_peers_evaluated: u64,
pub avg_selection_time_us: f64,
}
impl PeerSelectorStats {
pub fn cache_hit_rate(&self) -> f64 {
let total = self.cache_hits + self.cache_misses;
if total > 0 {
self.cache_hits as f64 / total as f64
} else {
0.0
}
}
}
pub struct PeerSelector {
config: PeerSelectorConfig,
geo_router: Arc<GeoRouter>,
quality_predictor: Arc<QualityPredictor>,
cache: Arc<DashMap<String, CachedSelection>>,
stats: Arc<RwLock<PeerSelectorStats>>,
}
impl PeerSelector {
pub fn new(config: PeerSelectorConfig) -> Self {
let geo_config = GeoRouterConfig::default();
let quality_config = QualityPredictorConfig::default();
Self {
config,
geo_router: Arc::new(GeoRouter::new(geo_config)),
quality_predictor: Arc::new(
QualityPredictor::new(quality_config).expect("Default config should be valid"),
),
cache: Arc::new(DashMap::new()),
stats: Arc::new(RwLock::new(PeerSelectorStats::default())),
}
}
pub fn with_configs(
config: PeerSelectorConfig,
geo_config: GeoRouterConfig,
quality_config: QualityPredictorConfig,
) -> Self {
Self {
config,
geo_router: Arc::new(GeoRouter::new(geo_config)),
quality_predictor: Arc::new(
QualityPredictor::new(quality_config).expect("Config should be valid"),
),
cache: Arc::new(DashMap::new()),
stats: Arc::new(RwLock::new(PeerSelectorStats::default())),
}
}
pub fn add_peer_location(&self, peer_id: PeerId, location: GeoLocation) {
self.geo_router.update_peer_location(peer_id, location);
self.invalidate_cache();
}
pub fn remove_peer(&self, peer_id: &PeerId) {
self.geo_router.remove_peer(peer_id);
self.quality_predictor.remove_peer(peer_id);
self.invalidate_cache();
}
pub fn update_peer_quality(
&self,
peer_id: PeerId,
latency_ms: f64,
bandwidth_mbps: f64,
success: bool,
) {
self.quality_predictor
.record_latency(peer_id, latency_ms as u64);
let bytes_per_sec = (bandwidth_mbps * 1_000_000.0 / 8.0) as u64;
self.quality_predictor
.record_bandwidth(peer_id, bytes_per_sec);
if !success {
self.quality_predictor.record_failure(peer_id);
}
self.invalidate_cache();
}
pub fn select_peers(&self, criteria: &SelectionCriteria) -> Vec<SelectedPeer> {
let start = Instant::now();
if self.config.enable_caching {
let cache_key = self.make_cache_key(criteria);
if let Some(cached) = self.cache.get(&cache_key) {
let age = start.duration_since(cached.cached_at);
if age.as_secs() < self.config.cache_ttl_secs {
let mut stats = self.stats.write();
stats.total_selections += 1;
stats.cache_hits += 1;
return cached.peers.clone();
}
}
}
let selected = self.select_peers_impl(criteria);
let elapsed = start.elapsed();
let mut stats = self.stats.write();
stats.total_selections += 1;
stats.cache_misses += 1;
stats.total_peers_evaluated += selected.len() as u64;
let new_avg = if stats.total_selections > 1 {
(stats.avg_selection_time_us * (stats.total_selections - 1) as f64
+ elapsed.as_micros() as f64)
/ stats.total_selections as f64
} else {
elapsed.as_micros() as f64
};
stats.avg_selection_time_us = new_avg;
drop(stats);
if self.config.enable_caching {
let cache_key = self.make_cache_key(criteria);
self.cache.insert(
cache_key,
CachedSelection {
peers: selected.clone(),
cached_at: start,
},
);
if self.cache.len() > self.config.max_cache_entries {
self.evict_old_cache_entries();
}
}
selected
}
fn select_peers_impl(&self, criteria: &SelectionCriteria) -> Vec<SelectedPeer> {
let mut scored_peers = Vec::new();
let geo_peers = if let Some(ref_location) = &criteria.reference_location {
self.geo_router.rank_peers_by_proximity(ref_location)
} else {
vec![]
};
for geo_peer in geo_peers {
if let Some(max_dist) = criteria.max_distance_km {
if let Some(dist) = geo_peer.distance_km {
if dist > max_dist {
continue;
}
}
}
let distance_score = self.calculate_distance_score(&geo_peer.distance_km);
let quality_prediction = self.quality_predictor.predict_quality(&geo_peer.peer_id);
let quality_score = quality_prediction
.as_ref()
.map(|p| p.overall_score)
.unwrap_or(0.5);
let latency_score = quality_prediction
.as_ref()
.map(|p| p.latency_score)
.unwrap_or(0.5);
let bandwidth_score = quality_prediction
.as_ref()
.map(|p| p.bandwidth_score)
.unwrap_or(0.5);
if quality_score < criteria.min_quality_score {
continue;
}
let overall_score = self.calculate_overall_score(
distance_score,
quality_score,
latency_score,
bandwidth_score,
);
scored_peers.push(SelectedPeer {
peer_id: geo_peer.peer_id,
score: overall_score,
distance_score,
quality_score,
latency_score,
bandwidth_score,
location: Some(geo_peer.location),
distance_km: geo_peer.distance_km,
});
}
scored_peers.sort_by(|a, b| {
b.score
.partial_cmp(&a.score)
.unwrap_or(std::cmp::Ordering::Equal)
});
scored_peers.truncate(criteria.max_results);
scored_peers
}
fn calculate_distance_score(&self, distance_km: &Option<f64>) -> f64 {
match distance_km {
Some(dist) => {
(-dist / 1000.0).exp()
}
None => 0.5, }
}
fn calculate_overall_score(
&self,
distance: f64,
quality: f64,
latency: f64,
bandwidth: f64,
) -> f64 {
distance * self.config.distance_weight
+ quality * self.config.quality_weight
+ latency * self.config.latency_weight
+ bandwidth * self.config.bandwidth_weight
}
fn make_cache_key(&self, criteria: &SelectionCriteria) -> String {
format!(
"loc:{:?}_qual:{}_dist:{:?}_max:{}",
criteria.reference_location,
criteria.min_quality_score,
criteria.max_distance_km,
criteria.max_results
)
}
fn invalidate_cache(&self) {
self.cache.clear();
}
fn evict_old_cache_entries(&self) {
let now = Instant::now();
let ttl = Duration::from_secs(self.config.cache_ttl_secs);
self.cache
.retain(|_, entry| now.duration_since(entry.cached_at) < ttl);
while self.cache.len() > self.config.max_cache_entries {
if let Some(entry) = self.cache.iter().next() {
let key = entry.key().clone();
drop(entry);
self.cache.remove(&key);
} else {
break;
}
}
}
pub fn stats(&self) -> PeerSelectorStats {
self.stats.read().clone()
}
#[allow(dead_code)]
pub fn reset(&self) {
self.cache.clear();
let mut stats = self.stats.write();
*stats = PeerSelectorStats::default();
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_peer_selector_config() {
let config = PeerSelectorConfig::default();
assert_eq!(config.distance_weight, 0.3);
assert_eq!(config.quality_weight, 0.3);
}
#[test]
fn test_config_presets() {
let low_latency = PeerSelectorConfig::low_latency();
assert!(low_latency.latency_weight > 0.3);
let high_bandwidth = PeerSelectorConfig::high_bandwidth();
assert!(high_bandwidth.bandwidth_weight > 0.5);
let mobile = PeerSelectorConfig::mobile();
assert!(mobile.distance_weight >= 0.5);
}
#[test]
fn test_peer_selector_creation() {
let config = PeerSelectorConfig::default();
let selector = PeerSelector::new(config);
let stats = selector.stats();
assert_eq!(stats.total_selections, 0);
}
#[test]
fn test_add_peer_location() {
let config = PeerSelectorConfig::default();
let selector = PeerSelector::new(config);
let peer = PeerId::random();
let location = GeoLocation::new(40.7128, -74.0060);
selector.add_peer_location(peer, location);
let loc = selector.geo_router.get_peer_location(&peer);
assert!(loc.is_some());
}
#[test]
fn test_remove_peer() {
let config = PeerSelectorConfig::default();
let selector = PeerSelector::new(config);
let peer = PeerId::random();
let location = GeoLocation::new(40.7128, -74.0060);
selector.add_peer_location(peer, location);
selector.remove_peer(&peer);
let loc = selector.geo_router.get_peer_location(&peer);
assert!(loc.is_none());
}
#[test]
fn test_selection_criteria() {
let criteria = SelectionCriteria {
reference_location: Some(GeoLocation::new(37.7749, -122.4194)),
min_quality_score: 0.5,
max_distance_km: Some(1000.0),
max_results: 5,
};
assert!(criteria.reference_location.is_some());
assert_eq!(criteria.max_results, 5);
}
#[test]
fn test_select_peers_empty() {
let config = PeerSelectorConfig::default();
let selector = PeerSelector::new(config);
let criteria = SelectionCriteria::default();
let selected = selector.select_peers(&criteria);
assert_eq!(selected.len(), 0);
}
#[test]
fn test_select_peers_with_location() {
let config = PeerSelectorConfig::default();
let selector = PeerSelector::new(config);
let peer1 = PeerId::random();
let peer2 = PeerId::random();
selector.add_peer_location(peer1, GeoLocation::new(40.7128, -74.0060)); selector.add_peer_location(peer2, GeoLocation::new(34.0522, -118.2437));
let criteria = SelectionCriteria {
reference_location: Some(GeoLocation::new(37.7749, -122.4194)),
min_quality_score: 0.0,
max_distance_km: None,
max_results: 10,
};
let selected = selector.select_peers(&criteria);
assert!(!selected.is_empty());
if selected.len() == 2 {
assert_eq!(selected[0].peer_id, peer2);
}
}
#[test]
fn test_distance_score_calculation() {
let config = PeerSelectorConfig::default();
let selector = PeerSelector::new(config);
let score_close = selector.calculate_distance_score(&Some(100.0));
let score_far = selector.calculate_distance_score(&Some(5000.0));
let score_none = selector.calculate_distance_score(&None);
assert!(score_close > score_far);
assert_eq!(score_none, 0.5);
}
#[test]
fn test_overall_score_calculation() {
let config = PeerSelectorConfig::default();
let selector = PeerSelector::new(config);
let score = selector.calculate_overall_score(1.0, 1.0, 1.0, 1.0);
assert!(score > 0.0 && score <= 1.0);
}
#[test]
fn test_cache_functionality() {
let config = PeerSelectorConfig {
enable_caching: true,
..Default::default()
};
let selector = PeerSelector::new(config);
let peer = PeerId::random();
selector.add_peer_location(peer, GeoLocation::new(40.7128, -74.0060));
let criteria = SelectionCriteria {
reference_location: Some(GeoLocation::new(37.7749, -122.4194)),
min_quality_score: 0.0,
max_distance_km: None,
max_results: 10,
};
selector.select_peers(&criteria);
let stats1 = selector.stats();
assert_eq!(stats1.cache_misses, 1);
selector.select_peers(&criteria);
let stats2 = selector.stats();
assert_eq!(stats2.cache_hits, 1);
}
#[test]
fn test_stats_cache_hit_rate() {
let stats = PeerSelectorStats {
cache_hits: 7,
cache_misses: 3,
..Default::default()
};
assert!((stats.cache_hit_rate() - 0.7).abs() < 0.01);
}
#[test]
fn test_max_distance_filtering() {
let config = PeerSelectorConfig::default();
let selector = PeerSelector::new(config);
let peer_la = PeerId::random();
let peer_london = PeerId::random();
selector.add_peer_location(peer_la, GeoLocation::new(34.0522, -118.2437)); selector.add_peer_location(peer_london, GeoLocation::new(51.5074, -0.1278));
let criteria = SelectionCriteria {
reference_location: Some(GeoLocation::new(37.7749, -122.4194)), min_quality_score: 0.0,
max_distance_km: Some(1000.0),
max_results: 10,
};
let selected = selector.select_peers(&criteria);
assert_eq!(selected.len(), 1); assert_eq!(selected[0].peer_id, peer_la);
}
#[test]
fn test_min_quality_filtering() {
let config = PeerSelectorConfig::default();
let selector = PeerSelector::new(config);
let peer = PeerId::random();
selector.add_peer_location(peer, GeoLocation::new(40.7128, -74.0060));
selector.update_peer_quality(peer, 1000.0, 0.1, false);
let criteria = SelectionCriteria {
reference_location: Some(GeoLocation::new(37.7749, -122.4194)),
min_quality_score: 0.9, max_distance_km: None,
max_results: 10,
};
let selected = selector.select_peers(&criteria);
assert!(selected.is_empty() || selected[0].quality_score >= criteria.min_quality_score);
}
}