stygian_proxy/strategy/weighted.rs
1//! Weighted random proxy rotation strategy.
2
3use async_trait::async_trait;
4use rand::Rng as _;
5
6use crate::error::{ProxyError, ProxyResult};
7use crate::strategy::{ProxyCandidate, RotationStrategy, healthy_candidates};
8
9/// Selects a healthy proxy with probability proportional to its `weight`.
10///
11/// Proxies with `weight == 0` are never selected.
12///
13/// # Example
14/// ```
15/// # tokio_test::block_on(async {
16/// use stygian_proxy::strategy::{WeightedStrategy, RotationStrategy, ProxyCandidate};
17/// use stygian_proxy::types::ProxyMetrics;
18/// use std::sync::Arc;
19/// use uuid::Uuid;
20///
21/// let strategy = WeightedStrategy;
22/// let candidates = vec![
23/// ProxyCandidate { id: Uuid::new_v4(), weight: 10, metrics: Arc::new(ProxyMetrics::default()), healthy: true },
24/// ProxyCandidate { id: Uuid::new_v4(), weight: 1, metrics: Arc::new(ProxyMetrics::default()), healthy: true },
25/// ];
26/// strategy.select(&candidates).await.unwrap();
27/// # })
28/// ```
29#[derive(Debug, Default, Clone, Copy)]
30pub struct WeightedStrategy;
31
32#[async_trait]
33impl RotationStrategy for WeightedStrategy {
34 async fn select<'a>(
35 &self,
36 candidates: &'a [ProxyCandidate],
37 ) -> ProxyResult<&'a ProxyCandidate> {
38 let healthy: Vec<&ProxyCandidate> = healthy_candidates(candidates)
39 .into_iter()
40 .filter(|c| c.weight > 0)
41 .collect();
42
43 if healthy.is_empty() {
44 return Err(ProxyError::AllProxiesUnhealthy);
45 }
46
47 let total: u64 = healthy.iter().map(|c| c.weight as u64).sum();
48 let mut cursor: u64 = rand::rng().random_range(0..total);
49
50 for candidate in &healthy {
51 if cursor < candidate.weight as u64 {
52 return Ok(candidate);
53 }
54 cursor -= candidate.weight as u64;
55 }
56
57 // Unreachable: cursor always exhausts within the loop.
58 Ok(healthy[healthy.len() - 1])
59 }
60}