Skip to main content

stygian_proxy/strategy/
mod.rs

1//! Proxy rotation strategy trait and built-in implementations.
2//!
3//! ## Why Thompson Sampling
4//!
5//! Round-robin assumes proxy health is stationary; in practice a proxy that
6//! worked a minute ago may now be banned, and a dead one may recover.
7//! Round-robin therefore wastes a fixed fraction of every batch on dead
8//! proxies. Thompson sampling — see
9//! `docs/dev/project/scraping-guide-2026-llm-context.md`
10//! §"Stop Round-Robining Dead Proxies: Bayesian Selection" (L3018-3021)
11//! — models each proxy's success probability as a `Beta(α, β)` posterior
12//! and draws one sample per candidate on each acquisition. The 2026
13//! `ProxyOps` benchmark cited in the guide reports **76 % success with
14//! Thompson sampling vs 36 % with round-robin** on identical proxies and
15//! targets (549 114 requests in 7 days) — more than double the success
16//! rate.
17//!
18//! [`ThompsonStrategy`] is the Stygian implementation; it lives behind the
19//! `bayesian-rotation` cargo feature (off by default to preserve
20//! deterministic round-robin in existing deployments). The hot path is
21//! two atomic loads + one xorshift64 draw per candidate, which stays well
22//! under the 1 µs acquisition budget documented in
23//! `crates/stygian-proxy/AGENTS.md`.
24
25use std::sync::Arc;
26
27use async_trait::async_trait;
28use uuid::Uuid;
29
30use crate::error::ProxyResult;
31use crate::types::{CapabilityRequirement, ProxyCapabilities, ProxyMetrics};
32
33mod least_used;
34mod random;
35mod round_robin;
36#[cfg(feature = "bayesian-rotation")]
37pub mod thompson;
38mod weighted;
39
40pub use least_used::LeastUsedStrategy;
41pub use random::RandomStrategy;
42pub use round_robin::RoundRobinStrategy;
43#[cfg(feature = "bayesian-rotation")]
44pub use thompson::ThompsonStrategy;
45pub use weighted::WeightedStrategy;
46
47// ─────────────────────────────────────────────────────────────────────────────
48// ProxyCandidate
49// ─────────────────────────────────────────────────────────────────────────────
50
51/// A lightweight view of a proxy considered for selection.
52///
53/// Strategies operate on slices of `ProxyCandidate` values built from the live
54/// proxy pool. The `metrics` field allows latency- or usage-aware selection
55/// without acquiring a write lock.
56#[derive(Debug, Clone)]
57pub struct ProxyCandidate {
58    /// Stable identifier matching the [`ProxyRecord`](crate::types::ProxyRecord).
59    pub id: Uuid,
60    /// Relative weight used by [`WeightedStrategy`].
61    pub weight: u32,
62    /// Shared atomics updated by every request through this proxy.
63    pub metrics: Arc<ProxyMetrics>,
64    /// Whether the proxy currently passes health checks.
65    pub healthy: bool,
66    /// Protocol-level capabilities this proxy exposes.
67    pub capabilities: ProxyCapabilities,
68}
69
70// ─────────────────────────────────────────────────────────────────────────────
71// RotationStrategy trait
72// ─────────────────────────────────────────────────────────────────────────────
73
74/// Selects a proxy from a slice of candidates on each request.
75///
76/// Implementations receive **all** candidates (healthy and unhealthy) so they
77/// can distinguish between an empty pool and a pool where every proxy is
78/// temporarily down. Call [`healthy_candidates`] to filter the slice.
79///
80/// # Example
81/// ```rust,no_run
82/// use stygian_proxy::strategy::{ProxyCandidate, RotationStrategy, RoundRobinStrategy};
83///
84/// async fn pick(candidates: &[ProxyCandidate]) {
85///     let strategy = RoundRobinStrategy::default();
86///     let chosen = strategy.select(candidates).await.unwrap();
87///     println!("selected: {}", chosen.id);
88/// }
89/// ```
90#[async_trait]
91pub trait RotationStrategy: Send + Sync + 'static {
92    /// Select one candidate from `candidates`.
93    ///
94    /// Returns [`crate::error::ProxyError::AllProxiesUnhealthy`] when every candidate has
95    /// `healthy == false`.
96    async fn select<'a>(&self, candidates: &'a [ProxyCandidate])
97    -> ProxyResult<&'a ProxyCandidate>;
98}
99
100/// Shared-ownership type alias for a [`RotationStrategy`] implementation.
101pub type BoxedRotationStrategy = Arc<dyn RotationStrategy>;
102
103// ─────────────────────────────────────────────────────────────────────────────
104// BayesianObserver
105// ─────────────────────────────────────────────────────────────────────────────
106
107/// Callback that receives outcome observations for a specific proxy.
108///
109/// `ProxyManager` calls `observe(success = true)` on
110/// [`ProxyHandle::mark_success`](crate::manager::ProxyHandle::mark_success)
111/// and `observe(success = false)` on drop without `mark_success`. The
112/// default implementation is [`NoopBayesianObserver`], which discards
113/// observations; the `bayesian-rotation` feature wires
114/// [`ThompsonStrategy`] in via
115/// [`crate::manager::ProxyManagerBuilder::with_thompson_sampling`].
116pub trait BayesianObserver: Send + Sync + 'static {
117    /// Record one outcome for `proxy_id`. Implementations must be safe to
118    /// call from any thread and must never block on locks the hot path
119    /// already holds.
120    fn observe(&self, proxy_id: Uuid, success: bool);
121}
122
123/// Discarding [`BayesianObserver`] used when no strategy is configured.
124///
125/// The default for `ProxyManager` so the observer field has a
126/// zero-overhead `None`-equivalent without changing the manager's type
127/// signature.
128#[derive(Debug, Default, Clone, Copy)]
129pub struct NoopBayesianObserver;
130
131impl BayesianObserver for NoopBayesianObserver {
132    #[inline]
133    fn observe(&self, _proxy_id: Uuid, _success: bool) {}
134}
135
136/// Shared-ownership type alias for a [`BayesianObserver`] implementation.
137pub type BoxedBayesianObserver = Arc<dyn BayesianObserver>;
138
139// ─────────────────────────────────────────────────────────────────────────────
140// Shared helper
141// ─────────────────────────────────────────────────────────────────────────────
142
143/// Filter `all` to only the candidates that are currently healthy.
144///
145/// Returns references into the original slice, so no allocation is needed
146/// beyond the returned `Vec`.
147#[must_use]
148pub fn healthy_candidates(all: &[ProxyCandidate]) -> Vec<&ProxyCandidate> {
149    all.iter().filter(|c| c.healthy).collect()
150}
151
152/// Filter `all` to candidates that are healthy **and** satisfy `req`.
153///
154/// An empty [`CapabilityRequirement`] (all flags `false`, no geo filter)
155/// behaves identically to [`healthy_candidates`].
156///
157/// # Example
158/// ```
159/// use std::sync::Arc;
160/// use stygian_proxy::strategy::{ProxyCandidate, capable_healthy_candidates};
161/// use stygian_proxy::types::{CapabilityRequirement, ProxyCapabilities, ProxyMetrics};
162/// use uuid::Uuid;
163///
164/// let caps = ProxyCapabilities { supports_https_connect: true, ..Default::default() };
165/// let candidate = ProxyCandidate {
166///     id: Uuid::new_v4(),
167///     weight: 1,
168///     metrics: Arc::new(ProxyMetrics::default()),
169///     healthy: true,
170///     capabilities: caps,
171/// };
172/// let req = CapabilityRequirement { require_https_connect: true, ..Default::default() };
173/// let candidates = [candidate];
174/// let result = capable_healthy_candidates(&candidates, &req);
175/// assert_eq!(result.len(), 1);
176/// ```
177#[must_use]
178pub fn capable_healthy_candidates<'a>(
179    all: &'a [ProxyCandidate],
180    req: &CapabilityRequirement,
181) -> Vec<&'a ProxyCandidate> {
182    all.iter()
183        .filter(|c| c.healthy && c.capabilities.satisfies(req))
184        .collect()
185}
186
187// ─────────────────────────────────────────────────────────────────────────────
188// Tests
189// ─────────────────────────────────────────────────────────────────────────────
190
191#[cfg(test)]
192pub(crate) mod tests {
193    use super::*;
194    use crate::error::ProxyError;
195    use crate::types::{CapabilityRequirement, ProxyCapabilities};
196    use std::sync::atomic::Ordering;
197
198    /// Build a `ProxyCandidate` with sensible test defaults.
199    pub fn candidate(id: u128, healthy: bool, weight: u32, requests: u64) -> ProxyCandidate {
200        let metrics = Arc::new(ProxyMetrics::default());
201        metrics.requests_total.store(requests, Ordering::Relaxed);
202        ProxyCandidate {
203            id: Uuid::from_u128(id),
204            weight,
205            metrics,
206            healthy,
207            capabilities: ProxyCapabilities::default(),
208        }
209    }
210
211    /// Build a `ProxyCandidate` with explicit capabilities.
212    pub fn candidate_with_caps(
213        id: u128,
214        healthy: bool,
215        weight: u32,
216        caps: ProxyCapabilities,
217    ) -> ProxyCandidate {
218        let metrics = Arc::new(ProxyMetrics::default());
219        ProxyCandidate {
220            id: Uuid::from_u128(id),
221            weight,
222            metrics,
223            healthy,
224            capabilities: caps,
225        }
226    }
227
228    #[tokio::test]
229    async fn healthy_candidates_filters() {
230        let c = vec![
231            candidate(1, true, 1, 0),
232            candidate(2, false, 1, 0),
233            candidate(3, true, 1, 0),
234        ];
235        let healthy = healthy_candidates(&c);
236        assert_eq!(healthy.len(), 2);
237        assert!(healthy.iter().all(|c| c.healthy));
238    }
239
240    #[tokio::test]
241    async fn all_unhealthy_returns_error() {
242        let c = vec![candidate(1, false, 1, 0), candidate(2, false, 1, 0)];
243        assert!(matches!(
244            RoundRobinStrategy::default().select(&c).await,
245            Err(ProxyError::AllProxiesUnhealthy)
246        ));
247    }
248
249    #[test]
250    fn capable_healthy_candidates_filters_by_capability() {
251        let c = vec![
252            candidate_with_caps(
253                1,
254                true,
255                1,
256                ProxyCapabilities {
257                    supports_https_connect: true,
258                    ..Default::default()
259                },
260            ),
261            candidate_with_caps(2, true, 1, ProxyCapabilities::default()),
262            candidate_with_caps(
263                3,
264                false,
265                1,
266                ProxyCapabilities {
267                    supports_https_connect: true,
268                    ..Default::default()
269                },
270            ),
271        ];
272        let req = CapabilityRequirement {
273            require_https_connect: true,
274            ..Default::default()
275        };
276        let result = capable_healthy_candidates(&c, &req);
277        // Only candidate 1: healthy AND supports_https_connect
278        assert_eq!(result.len(), 1);
279        assert_eq!(
280            result.first().map(|candidate| candidate.id),
281            Some(Uuid::from_u128(1))
282        );
283    }
284
285    #[test]
286    fn capable_healthy_candidates_empty_req_behaves_like_healthy() {
287        let c = vec![
288            candidate(1, true, 1, 0),
289            candidate(2, false, 1, 0),
290            candidate(3, true, 1, 0),
291        ];
292        let req = CapabilityRequirement::default();
293        let result = capable_healthy_candidates(&c, &req);
294        assert_eq!(result.len(), 2);
295    }
296
297    #[test]
298    fn capable_healthy_candidates_returns_empty_when_none_match() {
299        let c = vec![candidate(1, true, 1, 0), candidate(2, true, 1, 0)];
300        let req = CapabilityRequirement {
301            require_socks5_udp: true,
302            ..Default::default()
303        };
304        let result = capable_healthy_candidates(&c, &req);
305        assert!(result.is_empty());
306    }
307
308    #[test]
309    fn geo_country_filter_matches_exact_country() {
310        let gb_proxy_caps = ProxyCapabilities {
311            geo_country: Some("GB".into()),
312            ..Default::default()
313        };
314        let us_proxy_caps = ProxyCapabilities {
315            geo_country: Some("US".into()),
316            ..Default::default()
317        };
318        let c = vec![
319            candidate_with_caps(1, true, 1, gb_proxy_caps),
320            candidate_with_caps(2, true, 1, us_proxy_caps),
321            candidate_with_caps(3, true, 1, ProxyCapabilities::default()),
322        ];
323        let req = CapabilityRequirement {
324            require_geo_country: Some("GB".into()),
325            ..Default::default()
326        };
327        let result = capable_healthy_candidates(&c, &req);
328        assert_eq!(result.len(), 1);
329        assert_eq!(
330            result.first().map(|candidate| candidate.id),
331            Some(Uuid::from_u128(1))
332        );
333    }
334}