1use std::collections::HashMap;
38use std::time::{Duration, SystemTime};
39
40#[derive(Debug, Clone)]
42pub struct PeerCandidate {
43 pub peer_id: String,
45 pub reputation_score: f64,
47 pub network_health: f64,
49 pub current_load: f64,
51 pub latency_ms: f64,
53 pub bandwidth_mbps: f64,
55 pub distance_km: Option<f64>,
57 pub last_seen: SystemTime,
59}
60
61#[derive(Debug, Clone, Copy, PartialEq, Eq)]
63pub enum SelectionStrategy {
64 Best,
66 WeightedRandom,
68 RoundRobin,
70 LeastLoaded,
72 LowestLatency,
74}
75
76#[derive(Debug, Clone)]
78pub struct SelectionWeights {
79 pub reputation: f64,
81 pub network_health: f64,
83 pub load: f64,
85 pub latency: f64,
87 pub bandwidth: f64,
89 pub distance: f64,
91}
92
93impl Default for SelectionWeights {
94 fn default() -> Self {
95 Self {
96 reputation: 0.3,
97 network_health: 0.25,
98 load: 0.2,
99 latency: 0.15,
100 bandwidth: 0.1,
101 distance: 0.0,
102 }
103 }
104}
105
106pub struct PeerSelector {
108 candidates: Vec<PeerCandidate>,
109 weights: SelectionWeights,
110 strategy: SelectionStrategy,
111 round_robin_index: usize,
112 peer_request_counts: HashMap<String, u64>,
113 stale_threshold: Duration,
114}
115
116impl PeerSelector {
117 #[must_use]
119 #[inline]
120 pub fn new() -> Self {
121 Self {
122 candidates: Vec::new(),
123 weights: SelectionWeights::default(),
124 strategy: SelectionStrategy::Best,
125 round_robin_index: 0,
126 peer_request_counts: HashMap::new(),
127 stale_threshold: Duration::from_secs(300), }
129 }
130
131 #[must_use]
133 #[inline]
134 pub fn with_weights(weights: SelectionWeights) -> Self {
135 Self {
136 weights,
137 ..Self::new()
138 }
139 }
140
141 #[inline]
143 pub fn set_strategy(&mut self, strategy: SelectionStrategy) {
144 self.strategy = strategy;
145 }
146
147 #[inline]
149 pub fn set_stale_threshold(&mut self, threshold: Duration) {
150 self.stale_threshold = threshold;
151 }
152
153 pub fn add_candidate(&mut self, candidate: PeerCandidate) {
155 self.candidates.retain(|c| c.peer_id != candidate.peer_id);
157 self.candidates.push(candidate);
158 }
159
160 #[inline]
162 pub fn remove_candidate(&mut self, peer_id: &str) {
163 self.candidates.retain(|c| c.peer_id != peer_id);
164 self.peer_request_counts.remove(peer_id);
165 }
166
167 pub fn remove_stale_peers(&mut self) -> usize {
169 let now = SystemTime::now();
170 let initial_count = self.candidates.len();
171
172 self.candidates.retain(|c| {
173 if let Ok(duration) = now.duration_since(c.last_seen) {
174 duration < self.stale_threshold
175 } else {
176 true }
178 });
179
180 initial_count - self.candidates.len()
181 }
182
183 #[inline]
185 fn calculate_score(&self, peer: &PeerCandidate) -> f64 {
186 let mut score = 0.0;
187
188 score += peer.reputation_score * self.weights.reputation;
190
191 score += peer.network_health * self.weights.network_health;
193
194 score += (1.0 - peer.current_load) * self.weights.load;
196
197 let normalized_latency = 1.0 - (peer.latency_ms.min(500.0) / 500.0);
200 score += normalized_latency * self.weights.latency;
201
202 let normalized_bandwidth = peer.bandwidth_mbps.min(1000.0) / 1000.0;
205 score += normalized_bandwidth * self.weights.bandwidth;
206
207 if let Some(distance) = peer.distance_km {
209 let normalized_distance = 1.0 - (distance.min(10000.0) / 10000.0);
211 score += normalized_distance * self.weights.distance;
212 }
213
214 score
215 }
216
217 #[must_use]
219 pub fn select_best(&mut self) -> Option<PeerCandidate> {
220 if self.candidates.is_empty() {
221 return None;
222 }
223
224 match self.strategy {
225 SelectionStrategy::Best => self.select_highest_score(),
226 SelectionStrategy::WeightedRandom => self.select_weighted_random(),
227 SelectionStrategy::RoundRobin => self.select_round_robin(),
228 SelectionStrategy::LeastLoaded => self.select_least_loaded(),
229 SelectionStrategy::LowestLatency => self.select_lowest_latency(),
230 }
231 }
232
233 fn select_highest_score(&mut self) -> Option<PeerCandidate> {
235 let mut scored: Vec<_> = self
236 .candidates
237 .iter()
238 .map(|c| (c.clone(), self.calculate_score(c)))
239 .collect();
240
241 scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
242
243 scored.first().map(|(peer, _)| {
244 *self
245 .peer_request_counts
246 .entry(peer.peer_id.clone())
247 .or_insert(0) += 1;
248 peer.clone()
249 })
250 }
251
252 fn select_weighted_random(&mut self) -> Option<PeerCandidate> {
254 use rand::Rng;
255
256 let scores: Vec<_> = self
257 .candidates
258 .iter()
259 .map(|c| self.calculate_score(c))
260 .collect();
261
262 let total_score: f64 = scores.iter().sum();
263 if total_score == 0.0 {
264 return self.candidates.first().cloned();
265 }
266
267 let mut rng = rand::thread_rng();
268 let mut random_value = rng.gen_range(0.0..total_score);
269
270 for (i, score) in scores.iter().enumerate() {
271 random_value -= score;
272 if random_value <= 0.0 {
273 let peer = self.candidates[i].clone();
274 *self
275 .peer_request_counts
276 .entry(peer.peer_id.clone())
277 .or_insert(0) += 1;
278 return Some(peer);
279 }
280 }
281
282 self.candidates.last().cloned()
284 }
285
286 fn select_round_robin(&mut self) -> Option<PeerCandidate> {
288 if self.candidates.is_empty() {
289 return None;
290 }
291
292 let peer = self.candidates[self.round_robin_index % self.candidates.len()].clone();
293 self.round_robin_index = (self.round_robin_index + 1) % self.candidates.len();
294 *self
295 .peer_request_counts
296 .entry(peer.peer_id.clone())
297 .or_insert(0) += 1;
298 Some(peer)
299 }
300
301 fn select_least_loaded(&mut self) -> Option<PeerCandidate> {
303 let mut sorted = self.candidates.clone();
304 sorted.sort_by(|a, b| {
305 a.current_load
306 .partial_cmp(&b.current_load)
307 .unwrap_or(std::cmp::Ordering::Equal)
308 });
309
310 sorted.first().map(|peer| {
311 *self
312 .peer_request_counts
313 .entry(peer.peer_id.clone())
314 .or_insert(0) += 1;
315 peer.clone()
316 })
317 }
318
319 fn select_lowest_latency(&mut self) -> Option<PeerCandidate> {
321 let mut sorted = self.candidates.clone();
322 sorted.sort_by(|a, b| {
323 a.latency_ms
324 .partial_cmp(&b.latency_ms)
325 .unwrap_or(std::cmp::Ordering::Equal)
326 });
327
328 sorted.first().map(|peer| {
329 *self
330 .peer_request_counts
331 .entry(peer.peer_id.clone())
332 .or_insert(0) += 1;
333 peer.clone()
334 })
335 }
336
337 #[must_use]
339 #[inline]
340 pub fn select_top_n(&self, n: usize) -> Vec<PeerCandidate> {
341 let mut scored: Vec<_> = self
342 .candidates
343 .iter()
344 .map(|c| (c.clone(), self.calculate_score(c)))
345 .collect();
346
347 scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
348
349 scored.into_iter().take(n).map(|(peer, _)| peer).collect()
350 }
351
352 #[must_use]
354 #[inline]
355 pub fn get_qualified_peers(&self, min_score: f64) -> Vec<PeerCandidate> {
356 self.candidates
357 .iter()
358 .filter(|c| self.calculate_score(c) >= min_score)
359 .cloned()
360 .collect()
361 }
362
363 #[must_use]
365 #[inline]
366 pub fn candidate_count(&self) -> usize {
367 self.candidates.len()
368 }
369
370 #[must_use]
372 #[inline]
373 pub fn candidates(&self) -> &[PeerCandidate] {
374 &self.candidates
375 }
376
377 #[inline]
379 pub fn clear(&mut self) {
380 self.candidates.clear();
381 self.peer_request_counts.clear();
382 self.round_robin_index = 0;
383 }
384
385 #[must_use]
387 #[inline]
388 pub fn get_request_count(&self, peer_id: &str) -> u64 {
389 self.peer_request_counts.get(peer_id).copied().unwrap_or(0)
390 }
391
392 #[must_use]
394 #[inline]
395 pub fn get_statistics(&self) -> PeerSelectionStats {
396 if self.candidates.is_empty() {
397 return PeerSelectionStats::default();
398 }
399
400 let scores: Vec<f64> = self
401 .candidates
402 .iter()
403 .map(|c| self.calculate_score(c))
404 .collect();
405
406 let total_score: f64 = scores.iter().sum();
407 let avg_score = total_score / scores.len() as f64;
408 let max_score = scores.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
409 let min_score = scores.iter().cloned().fold(f64::INFINITY, f64::min);
410
411 let total_requests: u64 = self.peer_request_counts.values().sum();
412
413 PeerSelectionStats {
414 total_candidates: self.candidates.len(),
415 average_score: avg_score,
416 max_score,
417 min_score,
418 total_requests,
419 }
420 }
421}
422
423impl Default for PeerSelector {
424 fn default() -> Self {
425 Self::new()
426 }
427}
428
429#[derive(Debug, Clone, Default)]
431pub struct PeerSelectionStats {
432 pub total_candidates: usize,
434 pub average_score: f64,
436 pub max_score: f64,
438 pub min_score: f64,
440 pub total_requests: u64,
442}
443
444#[cfg(test)]
445mod tests {
446 use super::*;
447
448 fn create_test_peer(peer_id: &str, reputation: f64, health: f64, load: f64) -> PeerCandidate {
449 PeerCandidate {
450 peer_id: peer_id.to_string(),
451 reputation_score: reputation,
452 network_health: health,
453 current_load: load,
454 latency_ms: 100.0,
455 bandwidth_mbps: 100.0,
456 distance_km: None,
457 last_seen: SystemTime::now(),
458 }
459 }
460
461 #[test]
462 fn test_peer_selection_best_strategy() {
463 let mut selector = PeerSelector::new();
464
465 selector.add_candidate(create_test_peer("peer1", 0.5, 0.5, 0.8));
466 selector.add_candidate(create_test_peer("peer2", 0.9, 0.9, 0.2));
467 selector.add_candidate(create_test_peer("peer3", 0.7, 0.7, 0.5));
468
469 selector.set_strategy(SelectionStrategy::Best);
470 let best = selector.select_best().unwrap();
471 assert_eq!(best.peer_id, "peer2");
472 }
473
474 #[test]
475 fn test_peer_selection_least_loaded() {
476 let mut selector = PeerSelector::new();
477
478 selector.add_candidate(create_test_peer("peer1", 0.9, 0.9, 0.8));
479 selector.add_candidate(create_test_peer("peer2", 0.5, 0.5, 0.1));
480 selector.add_candidate(create_test_peer("peer3", 0.7, 0.7, 0.5));
481
482 selector.set_strategy(SelectionStrategy::LeastLoaded);
483 let best = selector.select_best().unwrap();
484 assert_eq!(best.peer_id, "peer2");
485 }
486
487 #[test]
488 fn test_peer_selection_top_n() {
489 let mut selector = PeerSelector::new();
490
491 selector.add_candidate(create_test_peer("peer1", 0.5, 0.5, 0.8));
492 selector.add_candidate(create_test_peer("peer2", 0.9, 0.9, 0.2));
493 selector.add_candidate(create_test_peer("peer3", 0.7, 0.7, 0.5));
494
495 let top_2 = selector.select_top_n(2);
496 assert_eq!(top_2.len(), 2);
497 assert_eq!(top_2[0].peer_id, "peer2");
498 assert_eq!(top_2[1].peer_id, "peer3");
499 }
500
501 #[test]
502 fn test_peer_selection_round_robin() {
503 let mut selector = PeerSelector::new();
504
505 selector.add_candidate(create_test_peer("peer1", 0.5, 0.5, 0.5));
506 selector.add_candidate(create_test_peer("peer2", 0.5, 0.5, 0.5));
507 selector.add_candidate(create_test_peer("peer3", 0.5, 0.5, 0.5));
508
509 selector.set_strategy(SelectionStrategy::RoundRobin);
510
511 assert_eq!(selector.select_best().unwrap().peer_id, "peer1");
512 assert_eq!(selector.select_best().unwrap().peer_id, "peer2");
513 assert_eq!(selector.select_best().unwrap().peer_id, "peer3");
514 assert_eq!(selector.select_best().unwrap().peer_id, "peer1");
515 }
516
517 #[test]
518 fn test_remove_candidate() {
519 let mut selector = PeerSelector::new();
520
521 selector.add_candidate(create_test_peer("peer1", 0.5, 0.5, 0.5));
522 selector.add_candidate(create_test_peer("peer2", 0.5, 0.5, 0.5));
523
524 assert_eq!(selector.candidate_count(), 2);
525
526 selector.remove_candidate("peer1");
527 assert_eq!(selector.candidate_count(), 1);
528 assert_eq!(selector.candidates()[0].peer_id, "peer2");
529 }
530
531 #[test]
532 fn test_custom_weights() {
533 let weights = SelectionWeights {
534 reputation: 1.0,
535 network_health: 0.0,
536 load: 0.0,
537 latency: 0.0,
538 bandwidth: 0.0,
539 distance: 0.0,
540 };
541
542 let mut selector = PeerSelector::with_weights(weights);
543
544 selector.add_candidate(create_test_peer("peer1", 0.5, 1.0, 0.0));
545 selector.add_candidate(create_test_peer("peer2", 1.0, 0.0, 1.0));
546
547 selector.set_strategy(SelectionStrategy::Best);
548 let best = selector.select_best().unwrap();
549 assert_eq!(best.peer_id, "peer2"); }
551
552 #[test]
553 fn test_qualified_peers() {
554 let mut selector = PeerSelector::new();
555
556 selector.add_candidate(create_test_peer("peer1", 0.3, 0.3, 0.9));
557 selector.add_candidate(create_test_peer("peer2", 0.9, 0.9, 0.1));
558 selector.add_candidate(create_test_peer("peer3", 0.7, 0.7, 0.5));
559
560 let qualified = selector.get_qualified_peers(0.5);
561 assert!(qualified.len() >= 2);
562 }
563
564 #[test]
565 fn test_statistics() {
566 let mut selector = PeerSelector::new();
567
568 selector.add_candidate(create_test_peer("peer1", 0.5, 0.5, 0.5));
569 selector.add_candidate(create_test_peer("peer2", 0.9, 0.9, 0.2));
570 selector.add_candidate(create_test_peer("peer3", 0.7, 0.7, 0.5));
571
572 let _ = selector.select_best();
573 let _ = selector.select_best();
574
575 let stats = selector.get_statistics();
576 assert_eq!(stats.total_candidates, 3);
577 assert!(stats.average_score > 0.0);
578 assert_eq!(stats.total_requests, 2);
579 }
580
581 #[test]
582 fn test_stale_peer_removal() {
583 let mut selector = PeerSelector::new();
584 selector.set_stale_threshold(Duration::from_secs(1));
585
586 let mut old_peer = create_test_peer("peer1", 0.5, 0.5, 0.5);
587 old_peer.last_seen = SystemTime::now() - Duration::from_secs(5);
588
589 selector.add_candidate(old_peer);
590 selector.add_candidate(create_test_peer("peer2", 0.5, 0.5, 0.5));
591
592 assert_eq!(selector.candidate_count(), 2);
593
594 let removed = selector.remove_stale_peers();
595 assert_eq!(removed, 1);
596 assert_eq!(selector.candidate_count(), 1);
597 }
598
599 #[test]
600 fn test_request_counting() {
601 let mut selector = PeerSelector::new();
602
603 selector.add_candidate(create_test_peer("peer1", 0.9, 0.9, 0.1));
604
605 selector.set_strategy(SelectionStrategy::Best);
606 let _ = selector.select_best();
607 let _ = selector.select_best();
608 let _ = selector.select_best();
609
610 assert_eq!(selector.get_request_count("peer1"), 3);
611 assert_eq!(selector.get_request_count("peer2"), 0);
612 }
613}