stygian_proxy/strategy/
mod.rs

1//! Proxy rotation strategy trait and built-in implementations.
2
3use std::sync::Arc;
4
5use async_trait::async_trait;
6use uuid::Uuid;
7
8use crate::error::ProxyResult;
9use crate::types::{CapabilityRequirement, ProxyCapabilities, ProxyMetrics};
10
11mod least_used;
12mod random;
13mod round_robin;
14mod weighted;
15
16pub use least_used::LeastUsedStrategy;
17pub use random::RandomStrategy;
18pub use round_robin::RoundRobinStrategy;
19pub use weighted::WeightedStrategy;
20
21// ─────────────────────────────────────────────────────────────────────────────
22// ProxyCandidate
23// ─────────────────────────────────────────────────────────────────────────────
24
25/// A lightweight view of a proxy considered for selection.
26///
27/// Strategies operate on slices of `ProxyCandidate` values built from the live
28/// proxy pool. The `metrics` field allows latency- or usage-aware selection
29/// without acquiring a write lock.
30#[derive(Debug, Clone)]
31pub struct ProxyCandidate {
32    /// Stable identifier matching the [`ProxyRecord`](crate::types::ProxyRecord).
33    pub id: Uuid,
34    /// Relative weight used by [`WeightedStrategy`].
35    pub weight: u32,
36    /// Shared atomics updated by every request through this proxy.
37    pub metrics: Arc<ProxyMetrics>,
38    /// Whether the proxy currently passes health checks.
39    pub healthy: bool,
40    /// Protocol-level capabilities this proxy exposes.
41    pub capabilities: ProxyCapabilities,
42}
43
44// ─────────────────────────────────────────────────────────────────────────────
45// RotationStrategy trait
46// ─────────────────────────────────────────────────────────────────────────────
47
48/// Selects a proxy from a slice of candidates on each request.
49///
50/// Implementations receive **all** candidates (healthy and unhealthy) so they
51/// can distinguish between an empty pool and a pool where every proxy is
52/// temporarily down. Call [`healthy_candidates`] to filter the slice.
53///
54/// # Example
55/// ```rust,no_run
56/// use stygian_proxy::strategy::{ProxyCandidate, RotationStrategy, RoundRobinStrategy};
57///
58/// async fn pick(candidates: &[ProxyCandidate]) {
59///     let strategy = RoundRobinStrategy::default();
60///     let chosen = strategy.select(candidates).await.unwrap();
61///     println!("selected: {}", chosen.id);
62/// }
63/// ```
64#[async_trait]
65pub trait RotationStrategy: Send + Sync + 'static {
66    /// Select one candidate from `candidates`.
67    ///
68    /// Returns [`crate::error::ProxyError::AllProxiesUnhealthy`] when every candidate has
69    /// `healthy == false`.
70    async fn select<'a>(&self, candidates: &'a [ProxyCandidate])
71    -> ProxyResult<&'a ProxyCandidate>;
72}
73
74/// Shared-ownership type alias for a [`RotationStrategy`] implementation.
75pub type BoxedRotationStrategy = Arc<dyn RotationStrategy>;
76
77// ─────────────────────────────────────────────────────────────────────────────
78// Shared helper
79// ─────────────────────────────────────────────────────────────────────────────
80
81/// Filter `all` to only the candidates that are currently healthy.
82///
83/// Returns references into the original slice, so no allocation is needed
84/// beyond the returned `Vec`.
85pub fn healthy_candidates(all: &[ProxyCandidate]) -> Vec<&ProxyCandidate> {
86    all.iter().filter(|c| c.healthy).collect()
87}
88
89/// Filter `all` to candidates that are healthy **and** satisfy `req`.
90///
91/// An empty [`CapabilityRequirement`] (all flags `false`, no geo filter)
92/// behaves identically to [`healthy_candidates`].
93///
94/// # Example
95/// ```
96/// use std::sync::Arc;
97/// use stygian_proxy::strategy::{ProxyCandidate, capable_healthy_candidates};
98/// use stygian_proxy::types::{CapabilityRequirement, ProxyCapabilities, ProxyMetrics};
99/// use uuid::Uuid;
100///
101/// let caps = ProxyCapabilities { supports_https_connect: true, ..Default::default() };
102/// let candidate = ProxyCandidate {
103///     id: Uuid::new_v4(),
104///     weight: 1,
105///     metrics: Arc::new(ProxyMetrics::default()),
106///     healthy: true,
107///     capabilities: caps,
108/// };
109/// let req = CapabilityRequirement { require_https_connect: true, ..Default::default() };
110/// let candidates = [candidate];
111/// let result = capable_healthy_candidates(&candidates, &req);
112/// assert_eq!(result.len(), 1);
113/// ```
114pub fn capable_healthy_candidates<'a>(
115    all: &'a [ProxyCandidate],
116    req: &CapabilityRequirement,
117) -> Vec<&'a ProxyCandidate> {
118    all.iter()
119        .filter(|c| c.healthy && c.capabilities.satisfies(req))
120        .collect()
121}
122
123// ─────────────────────────────────────────────────────────────────────────────
124// Tests
125// ─────────────────────────────────────────────────────────────────────────────
126
127#[cfg(test)]
128pub(crate) mod tests {
129    use super::*;
130    use crate::error::ProxyError;
131    use crate::types::{CapabilityRequirement, ProxyCapabilities};
132    use std::sync::atomic::Ordering;
133
134    /// Build a `ProxyCandidate` with sensible test defaults.
135    pub fn candidate(id: u128, healthy: bool, weight: u32, requests: u64) -> ProxyCandidate {
136        let metrics = Arc::new(ProxyMetrics::default());
137        metrics.requests_total.store(requests, Ordering::Relaxed);
138        ProxyCandidate {
139            id: Uuid::from_u128(id),
140            weight,
141            metrics,
142            healthy,
143            capabilities: ProxyCapabilities::default(),
144        }
145    }
146
147    /// Build a `ProxyCandidate` with explicit capabilities.
148    pub fn candidate_with_caps(
149        id: u128,
150        healthy: bool,
151        weight: u32,
152        caps: ProxyCapabilities,
153    ) -> ProxyCandidate {
154        let metrics = Arc::new(ProxyMetrics::default());
155        ProxyCandidate {
156            id: Uuid::from_u128(id),
157            weight,
158            metrics,
159            healthy,
160            capabilities: caps,
161        }
162    }
163
164    #[tokio::test]
165    async fn healthy_candidates_filters() {
166        let c = vec![
167            candidate(1, true, 1, 0),
168            candidate(2, false, 1, 0),
169            candidate(3, true, 1, 0),
170        ];
171        let healthy = healthy_candidates(&c);
172        assert_eq!(healthy.len(), 2);
173        assert!(healthy.iter().all(|c| c.healthy));
174    }
175
176    #[tokio::test]
177    async fn all_unhealthy_returns_error() {
178        let c = vec![candidate(1, false, 1, 0), candidate(2, false, 1, 0)];
179        assert!(matches!(
180            RoundRobinStrategy::default().select(&c).await,
181            Err(ProxyError::AllProxiesUnhealthy)
182        ));
183    }
184
185    #[test]
186    fn capable_healthy_candidates_filters_by_capability() {
187        let c = vec![
188            candidate_with_caps(
189                1,
190                true,
191                1,
192                ProxyCapabilities {
193                    supports_https_connect: true,
194                    ..Default::default()
195                },
196            ),
197            candidate_with_caps(2, true, 1, ProxyCapabilities::default()),
198            candidate_with_caps(
199                3,
200                false,
201                1,
202                ProxyCapabilities {
203                    supports_https_connect: true,
204                    ..Default::default()
205                },
206            ),
207        ];
208        let req = CapabilityRequirement {
209            require_https_connect: true,
210            ..Default::default()
211        };
212        let result = capable_healthy_candidates(&c, &req);
213        // Only candidate 1: healthy AND supports_https_connect
214        assert_eq!(result.len(), 1);
215        assert_eq!(
216            result.first().map(|candidate| candidate.id),
217            Some(Uuid::from_u128(1))
218        );
219    }
220
221    #[test]
222    fn capable_healthy_candidates_empty_req_behaves_like_healthy() {
223        let c = vec![
224            candidate(1, true, 1, 0),
225            candidate(2, false, 1, 0),
226            candidate(3, true, 1, 0),
227        ];
228        let req = CapabilityRequirement::default();
229        let result = capable_healthy_candidates(&c, &req);
230        assert_eq!(result.len(), 2);
231    }
232
233    #[test]
234    fn capable_healthy_candidates_returns_empty_when_none_match() {
235        let c = vec![candidate(1, true, 1, 0), candidate(2, true, 1, 0)];
236        let req = CapabilityRequirement {
237            require_socks5_udp: true,
238            ..Default::default()
239        };
240        let result = capable_healthy_candidates(&c, &req);
241        assert!(result.is_empty());
242    }
243
244    #[test]
245    fn geo_country_filter_matches_exact_country() {
246        let gb_proxy_caps = ProxyCapabilities {
247            geo_country: Some("GB".into()),
248            ..Default::default()
249        };
250        let us_proxy_caps = ProxyCapabilities {
251            geo_country: Some("US".into()),
252            ..Default::default()
253        };
254        let c = vec![
255            candidate_with_caps(1, true, 1, gb_proxy_caps),
256            candidate_with_caps(2, true, 1, us_proxy_caps),
257            candidate_with_caps(3, true, 1, ProxyCapabilities::default()),
258        ];
259        let req = CapabilityRequirement {
260            require_geo_country: Some("GB".into()),
261            ..Default::default()
262        };
263        let result = capable_healthy_candidates(&c, &req);
264        assert_eq!(result.len(), 1);
265        assert_eq!(
266            result.first().map(|candidate| candidate.id),
267            Some(Uuid::from_u128(1))
268        );
269    }
270}
stygian_proxy/strategy/mod.rs

stygian_proxy/strategy/
mod.rs
