use std::collections::HashMap;
use std::time::{Duration, SystemTime};
#[derive(Debug, Clone)]
pub struct PeerCandidate {
pub peer_id: String,
pub reputation_score: f64,
pub network_health: f64,
pub current_load: f64,
pub latency_ms: f64,
pub bandwidth_mbps: f64,
pub distance_km: Option<f64>,
pub last_seen: SystemTime,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SelectionStrategy {
Best,
WeightedRandom,
RoundRobin,
LeastLoaded,
LowestLatency,
}
#[derive(Debug, Clone)]
pub struct SelectionWeights {
pub reputation: f64,
pub network_health: f64,
pub load: f64,
pub latency: f64,
pub bandwidth: f64,
pub distance: f64,
}
impl Default for SelectionWeights {
fn default() -> Self {
Self {
reputation: 0.3,
network_health: 0.25,
load: 0.2,
latency: 0.15,
bandwidth: 0.1,
distance: 0.0,
}
}
}
pub struct PeerSelector {
candidates: Vec<PeerCandidate>,
weights: SelectionWeights,
strategy: SelectionStrategy,
round_robin_index: usize,
peer_request_counts: HashMap<String, u64>,
stale_threshold: Duration,
}
impl PeerSelector {
#[must_use]
#[inline]
pub fn new() -> Self {
Self {
candidates: Vec::new(),
weights: SelectionWeights::default(),
strategy: SelectionStrategy::Best,
round_robin_index: 0,
peer_request_counts: HashMap::new(),
stale_threshold: Duration::from_secs(300), }
}
#[must_use]
#[inline]
pub fn with_weights(weights: SelectionWeights) -> Self {
Self {
weights,
..Self::new()
}
}
#[inline]
pub fn set_strategy(&mut self, strategy: SelectionStrategy) {
self.strategy = strategy;
}
#[inline]
pub fn set_stale_threshold(&mut self, threshold: Duration) {
self.stale_threshold = threshold;
}
pub fn add_candidate(&mut self, candidate: PeerCandidate) {
self.candidates.retain(|c| c.peer_id != candidate.peer_id);
self.candidates.push(candidate);
}
#[inline]
pub fn remove_candidate(&mut self, peer_id: &str) {
self.candidates.retain(|c| c.peer_id != peer_id);
self.peer_request_counts.remove(peer_id);
}
pub fn remove_stale_peers(&mut self) -> usize {
let now = SystemTime::now();
let initial_count = self.candidates.len();
self.candidates.retain(|c| {
if let Ok(duration) = now.duration_since(c.last_seen) {
duration < self.stale_threshold
} else {
true }
});
initial_count - self.candidates.len()
}
#[inline]
fn calculate_score(&self, peer: &PeerCandidate) -> f64 {
let mut score = 0.0;
score += peer.reputation_score * self.weights.reputation;
score += peer.network_health * self.weights.network_health;
score += (1.0 - peer.current_load) * self.weights.load;
let normalized_latency = 1.0 - (peer.latency_ms.min(500.0) / 500.0);
score += normalized_latency * self.weights.latency;
let normalized_bandwidth = peer.bandwidth_mbps.min(1000.0) / 1000.0;
score += normalized_bandwidth * self.weights.bandwidth;
if let Some(distance) = peer.distance_km {
let normalized_distance = 1.0 - (distance.min(10000.0) / 10000.0);
score += normalized_distance * self.weights.distance;
}
score
}
#[must_use]
pub fn select_best(&mut self) -> Option<PeerCandidate> {
if self.candidates.is_empty() {
return None;
}
match self.strategy {
SelectionStrategy::Best => self.select_highest_score(),
SelectionStrategy::WeightedRandom => self.select_weighted_random(),
SelectionStrategy::RoundRobin => self.select_round_robin(),
SelectionStrategy::LeastLoaded => self.select_least_loaded(),
SelectionStrategy::LowestLatency => self.select_lowest_latency(),
}
}
fn select_highest_score(&mut self) -> Option<PeerCandidate> {
let mut scored: Vec<_> = self
.candidates
.iter()
.map(|c| (c.clone(), self.calculate_score(c)))
.collect();
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
scored.first().map(|(peer, _)| {
*self
.peer_request_counts
.entry(peer.peer_id.clone())
.or_insert(0) += 1;
peer.clone()
})
}
fn select_weighted_random(&mut self) -> Option<PeerCandidate> {
use rand::RngExt as _;
let scores: Vec<_> = self
.candidates
.iter()
.map(|c| self.calculate_score(c))
.collect();
let total_score: f64 = scores.iter().sum();
if total_score == 0.0 {
return self.candidates.first().cloned();
}
let mut rng = rand::rng();
let mut random_value = rng.random_range(0.0..total_score);
for (i, score) in scores.iter().enumerate() {
random_value -= score;
if random_value <= 0.0 {
let peer = self.candidates[i].clone();
*self
.peer_request_counts
.entry(peer.peer_id.clone())
.or_insert(0) += 1;
return Some(peer);
}
}
self.candidates.last().cloned()
}
fn select_round_robin(&mut self) -> Option<PeerCandidate> {
if self.candidates.is_empty() {
return None;
}
let peer = self.candidates[self.round_robin_index % self.candidates.len()].clone();
self.round_robin_index = (self.round_robin_index + 1) % self.candidates.len();
*self
.peer_request_counts
.entry(peer.peer_id.clone())
.or_insert(0) += 1;
Some(peer)
}
fn select_least_loaded(&mut self) -> Option<PeerCandidate> {
let mut sorted = self.candidates.clone();
sorted.sort_by(|a, b| {
a.current_load
.partial_cmp(&b.current_load)
.unwrap_or(std::cmp::Ordering::Equal)
});
sorted.first().map(|peer| {
*self
.peer_request_counts
.entry(peer.peer_id.clone())
.or_insert(0) += 1;
peer.clone()
})
}
fn select_lowest_latency(&mut self) -> Option<PeerCandidate> {
let mut sorted = self.candidates.clone();
sorted.sort_by(|a, b| {
a.latency_ms
.partial_cmp(&b.latency_ms)
.unwrap_or(std::cmp::Ordering::Equal)
});
sorted.first().map(|peer| {
*self
.peer_request_counts
.entry(peer.peer_id.clone())
.or_insert(0) += 1;
peer.clone()
})
}
#[must_use]
#[inline]
pub fn select_top_n(&self, n: usize) -> Vec<PeerCandidate> {
let mut scored: Vec<_> = self
.candidates
.iter()
.map(|c| (c.clone(), self.calculate_score(c)))
.collect();
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
scored.into_iter().take(n).map(|(peer, _)| peer).collect()
}
#[must_use]
#[inline]
pub fn get_qualified_peers(&self, min_score: f64) -> Vec<PeerCandidate> {
self.candidates
.iter()
.filter(|c| self.calculate_score(c) >= min_score)
.cloned()
.collect()
}
#[must_use]
#[inline]
pub fn candidate_count(&self) -> usize {
self.candidates.len()
}
#[must_use]
#[inline]
pub fn candidates(&self) -> &[PeerCandidate] {
&self.candidates
}
#[inline]
pub fn clear(&mut self) {
self.candidates.clear();
self.peer_request_counts.clear();
self.round_robin_index = 0;
}
#[must_use]
#[inline]
pub fn get_request_count(&self, peer_id: &str) -> u64 {
self.peer_request_counts.get(peer_id).copied().unwrap_or(0)
}
#[must_use]
#[inline]
pub fn get_statistics(&self) -> PeerSelectionStats {
if self.candidates.is_empty() {
return PeerSelectionStats::default();
}
let scores: Vec<f64> = self
.candidates
.iter()
.map(|c| self.calculate_score(c))
.collect();
let total_score: f64 = scores.iter().sum();
let avg_score = total_score / scores.len() as f64;
let max_score = scores.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
let min_score = scores.iter().cloned().fold(f64::INFINITY, f64::min);
let total_requests: u64 = self.peer_request_counts.values().sum();
PeerSelectionStats {
total_candidates: self.candidates.len(),
average_score: avg_score,
max_score,
min_score,
total_requests,
}
}
}
impl Default for PeerSelector {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Default)]
pub struct PeerSelectionStats {
pub total_candidates: usize,
pub average_score: f64,
pub max_score: f64,
pub min_score: f64,
pub total_requests: u64,
}
#[cfg(test)]
mod tests {
use super::*;
fn create_test_peer(peer_id: &str, reputation: f64, health: f64, load: f64) -> PeerCandidate {
PeerCandidate {
peer_id: peer_id.to_string(),
reputation_score: reputation,
network_health: health,
current_load: load,
latency_ms: 100.0,
bandwidth_mbps: 100.0,
distance_km: None,
last_seen: SystemTime::now(),
}
}
#[test]
fn test_peer_selection_best_strategy() {
let mut selector = PeerSelector::new();
selector.add_candidate(create_test_peer("peer1", 0.5, 0.5, 0.8));
selector.add_candidate(create_test_peer("peer2", 0.9, 0.9, 0.2));
selector.add_candidate(create_test_peer("peer3", 0.7, 0.7, 0.5));
selector.set_strategy(SelectionStrategy::Best);
let best = selector.select_best().unwrap();
assert_eq!(best.peer_id, "peer2");
}
#[test]
fn test_peer_selection_least_loaded() {
let mut selector = PeerSelector::new();
selector.add_candidate(create_test_peer("peer1", 0.9, 0.9, 0.8));
selector.add_candidate(create_test_peer("peer2", 0.5, 0.5, 0.1));
selector.add_candidate(create_test_peer("peer3", 0.7, 0.7, 0.5));
selector.set_strategy(SelectionStrategy::LeastLoaded);
let best = selector.select_best().unwrap();
assert_eq!(best.peer_id, "peer2");
}
#[test]
fn test_peer_selection_top_n() {
let mut selector = PeerSelector::new();
selector.add_candidate(create_test_peer("peer1", 0.5, 0.5, 0.8));
selector.add_candidate(create_test_peer("peer2", 0.9, 0.9, 0.2));
selector.add_candidate(create_test_peer("peer3", 0.7, 0.7, 0.5));
let top_2 = selector.select_top_n(2);
assert_eq!(top_2.len(), 2);
assert_eq!(top_2[0].peer_id, "peer2");
assert_eq!(top_2[1].peer_id, "peer3");
}
#[test]
fn test_peer_selection_round_robin() {
let mut selector = PeerSelector::new();
selector.add_candidate(create_test_peer("peer1", 0.5, 0.5, 0.5));
selector.add_candidate(create_test_peer("peer2", 0.5, 0.5, 0.5));
selector.add_candidate(create_test_peer("peer3", 0.5, 0.5, 0.5));
selector.set_strategy(SelectionStrategy::RoundRobin);
assert_eq!(selector.select_best().unwrap().peer_id, "peer1");
assert_eq!(selector.select_best().unwrap().peer_id, "peer2");
assert_eq!(selector.select_best().unwrap().peer_id, "peer3");
assert_eq!(selector.select_best().unwrap().peer_id, "peer1");
}
#[test]
fn test_remove_candidate() {
let mut selector = PeerSelector::new();
selector.add_candidate(create_test_peer("peer1", 0.5, 0.5, 0.5));
selector.add_candidate(create_test_peer("peer2", 0.5, 0.5, 0.5));
assert_eq!(selector.candidate_count(), 2);
selector.remove_candidate("peer1");
assert_eq!(selector.candidate_count(), 1);
assert_eq!(selector.candidates()[0].peer_id, "peer2");
}
#[test]
fn test_custom_weights() {
let weights = SelectionWeights {
reputation: 1.0,
network_health: 0.0,
load: 0.0,
latency: 0.0,
bandwidth: 0.0,
distance: 0.0,
};
let mut selector = PeerSelector::with_weights(weights);
selector.add_candidate(create_test_peer("peer1", 0.5, 1.0, 0.0));
selector.add_candidate(create_test_peer("peer2", 1.0, 0.0, 1.0));
selector.set_strategy(SelectionStrategy::Best);
let best = selector.select_best().unwrap();
assert_eq!(best.peer_id, "peer2"); }
#[test]
fn test_qualified_peers() {
let mut selector = PeerSelector::new();
selector.add_candidate(create_test_peer("peer1", 0.3, 0.3, 0.9));
selector.add_candidate(create_test_peer("peer2", 0.9, 0.9, 0.1));
selector.add_candidate(create_test_peer("peer3", 0.7, 0.7, 0.5));
let qualified = selector.get_qualified_peers(0.5);
assert!(qualified.len() >= 2);
}
#[test]
fn test_statistics() {
let mut selector = PeerSelector::new();
selector.add_candidate(create_test_peer("peer1", 0.5, 0.5, 0.5));
selector.add_candidate(create_test_peer("peer2", 0.9, 0.9, 0.2));
selector.add_candidate(create_test_peer("peer3", 0.7, 0.7, 0.5));
let _ = selector.select_best();
let _ = selector.select_best();
let stats = selector.get_statistics();
assert_eq!(stats.total_candidates, 3);
assert!(stats.average_score > 0.0);
assert_eq!(stats.total_requests, 2);
}
#[test]
fn test_stale_peer_removal() {
let mut selector = PeerSelector::new();
selector.set_stale_threshold(Duration::from_secs(1));
let mut old_peer = create_test_peer("peer1", 0.5, 0.5, 0.5);
old_peer.last_seen = SystemTime::now() - Duration::from_secs(5);
selector.add_candidate(old_peer);
selector.add_candidate(create_test_peer("peer2", 0.5, 0.5, 0.5));
assert_eq!(selector.candidate_count(), 2);
let removed = selector.remove_stale_peers();
assert_eq!(removed, 1);
assert_eq!(selector.candidate_count(), 1);
}
#[test]
fn test_request_counting() {
let mut selector = PeerSelector::new();
selector.add_candidate(create_test_peer("peer1", 0.9, 0.9, 0.1));
selector.set_strategy(SelectionStrategy::Best);
let _ = selector.select_best();
let _ = selector.select_best();
let _ = selector.select_best();
assert_eq!(selector.get_request_count("peer1"), 3);
assert_eq!(selector.get_request_count("peer2"), 0);
}
}