Skip to main content

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}