stealthscraper_rs/
proxy_pool.rs

1//! Pure upstream-proxy pool and rotation strategy.
2//!
3//! This is domain logic — no I/O — so it lives outside the `browser` feature
4//! and is fully unit-tested. The pool tracks a set of upstream proxy URLs,
5//! which one is currently selected, and per-endpoint health. When the
6//! orchestration layer detects that the current egress is blocked it calls
7//! [`ProxyPool::rotate`], which retires the current proxy and selects the next
8//! healthy one according to the configured [`RotationStrategy`].
9
10use rand::RngExt;
11
12use crate::geo::CountryCode;
13
14/// How [`ProxyPool::rotate`] picks the next healthy endpoint.
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
16pub enum RotationStrategy {
17    /// Cycle through endpoints in declaration order.
18    #[default]
19    RoundRobin,
20    /// Pick uniformly at random among the healthy endpoints.
21    Random,
22}
23
24/// A single upstream proxy and its current health flag.
25#[derive(Debug, Clone)]
26struct ProxyEndpoint {
27    url: String,
28    country: Option<CountryCode>,
29    healthy: bool,
30}
31
32/// A rotatable pool of upstream proxy URLs.
33#[derive(Debug, Clone)]
34pub struct ProxyPool {
35    endpoints: Vec<ProxyEndpoint>,
36    strategy: RotationStrategy,
37    current: Option<usize>,
38}
39
40impl ProxyPool {
41    /// Build a pool from upstream proxy URLs (untagged). Duplicate and empty URLs
42    /// are dropped.
43    ///
44    /// The initially selected endpoint is the first one for
45    /// [`RotationStrategy::RoundRobin`], or a random one for
46    /// [`RotationStrategy::Random`].
47    pub fn new(urls: impl IntoIterator<Item = String>, strategy: RotationStrategy) -> Self {
48        Self::with_endpoints(urls.into_iter().map(|url| (url, None)), strategy)
49    }
50
51    /// Build a pool from `(url, country)` pairs, where `country` tags the proxy's
52    /// exit country for proxy-led locale derivation. Duplicate/empty URLs dropped.
53    pub fn with_endpoints(
54        endpoints: impl IntoIterator<Item = (String, Option<CountryCode>)>,
55        strategy: RotationStrategy,
56    ) -> Self {
57        let mut collected: Vec<ProxyEndpoint> = Vec::new();
58        for (url, country) in endpoints {
59            let url = url.trim().to_string();
60            if url.is_empty() || collected.iter().any(|e| e.url == url) {
61                continue;
62            }
63            collected.push(ProxyEndpoint {
64                url,
65                country,
66                healthy: true,
67            });
68        }
69
70        let current = if collected.is_empty() {
71            None
72        } else {
73            match strategy {
74                RotationStrategy::RoundRobin => Some(0),
75                RotationStrategy::Random => Some(rand::rng().random_range(0..collected.len())),
76            }
77        };
78
79        Self {
80            endpoints: collected,
81            strategy,
82            current,
83        }
84    }
85
86    /// Total number of endpoints in the pool (healthy or not).
87    pub fn len(&self) -> usize {
88        self.endpoints.len()
89    }
90
91    /// Returns `true` when the pool holds no endpoints.
92    pub fn is_empty(&self) -> bool {
93        self.endpoints.is_empty()
94    }
95
96    /// Number of endpoints currently marked healthy.
97    pub fn healthy_count(&self) -> usize {
98        self.endpoints.iter().filter(|e| e.healthy).count()
99    }
100
101    /// The currently selected upstream proxy URL, if any.
102    pub fn selected(&self) -> Option<&str> {
103        self.current.map(|i| self.endpoints[i].url.as_str())
104    }
105
106    /// The tagged exit country of the currently selected proxy, if any.
107    pub fn selected_country(&self) -> Option<CountryCode> {
108        self.current.and_then(|i| self.endpoints[i].country)
109    }
110
111    /// Retire the current endpoint and select the next healthy one.
112    ///
113    /// Marks the current endpoint unhealthy (so rotation never returns to a
114    /// known-blocked proxy), then selects the next healthy endpoint per the
115    /// [`RotationStrategy`]. Returns the newly selected URL, or `None` when no
116    /// healthy endpoints remain.
117    pub fn rotate(&mut self) -> Option<String> {
118        if let Some(cur) = self.current {
119            self.endpoints[cur].healthy = false;
120        }
121
122        let next = match self.strategy {
123            RotationStrategy::RoundRobin => self.next_round_robin(),
124            RotationStrategy::Random => self.next_random(),
125        };
126
127        self.current = next;
128        next.map(|i| self.endpoints[i].url.clone())
129    }
130
131    fn next_round_robin(&self) -> Option<usize> {
132        let n = self.endpoints.len();
133        if n == 0 {
134            return None;
135        }
136        let start = self.current.unwrap_or(0);
137        (1..=n)
138            .map(|step| (start + step) % n)
139            .find(|&idx| self.endpoints[idx].healthy)
140    }
141
142    fn next_random(&self) -> Option<usize> {
143        let healthy: Vec<usize> = self
144            .endpoints
145            .iter()
146            .enumerate()
147            .filter(|(_, e)| e.healthy)
148            .map(|(i, _)| i)
149            .collect();
150        if healthy.is_empty() {
151            return None;
152        }
153        Some(healthy[rand::rng().random_range(0..healthy.len())])
154    }
155}
156
157#[cfg(test)]
158mod tests {
159    use super::*;
160
161    fn pool(urls: &[&str], strategy: RotationStrategy) -> ProxyPool {
162        ProxyPool::new(urls.iter().map(|s| s.to_string()), strategy)
163    }
164
165    #[test]
166    fn new_empty_pool_has_no_selection() {
167        let p = ProxyPool::new(Vec::<String>::new(), RotationStrategy::RoundRobin);
168        assert!(p.is_empty());
169        assert_eq!(p.selected(), None);
170        assert_eq!(p.healthy_count(), 0);
171    }
172
173    #[test]
174    fn new_dedups_and_trims() {
175        let p = pool(
176            &["http://a", " http://a ", "http://b", ""],
177            RotationStrategy::RoundRobin,
178        );
179        assert_eq!(p.len(), 2);
180    }
181
182    #[test]
183    fn round_robin_selects_first_then_advances() {
184        let mut p = pool(
185            &["http://a", "http://b", "http://c"],
186            RotationStrategy::RoundRobin,
187        );
188        assert_eq!(p.selected(), Some("http://a"));
189        assert_eq!(p.rotate().as_deref(), Some("http://b"));
190        assert_eq!(p.rotate().as_deref(), Some("http://c"));
191    }
192
193    #[test]
194    fn rotate_retires_current_and_skips_it_when_wrapping() {
195        let mut p = pool(&["http://a", "http://b"], RotationStrategy::RoundRobin);
196        // a -> b (a retired)
197        assert_eq!(p.rotate().as_deref(), Some("http://b"));
198        // b retired, a already retired -> none healthy left
199        assert_eq!(p.rotate(), None);
200        assert_eq!(p.healthy_count(), 0);
201        assert_eq!(p.selected(), None);
202    }
203
204    #[test]
205    fn rotate_on_empty_pool_is_none() {
206        let mut p = ProxyPool::new(Vec::<String>::new(), RotationStrategy::RoundRobin);
207        assert_eq!(p.rotate(), None);
208    }
209
210    #[test]
211    fn healthy_count_decreases_on_rotate() {
212        let mut p = pool(
213            &["http://a", "http://b", "http://c"],
214            RotationStrategy::RoundRobin,
215        );
216        assert_eq!(p.healthy_count(), 3);
217        p.rotate();
218        assert_eq!(p.healthy_count(), 2);
219    }
220
221    #[test]
222    fn endpoints_carry_country_tags_through_rotation() {
223        let de = CountryCode::new("DE");
224        let fr = CountryCode::new("FR");
225        let mut p = ProxyPool::with_endpoints(
226            [("http://a".to_string(), de), ("http://b".to_string(), fr)],
227            RotationStrategy::RoundRobin,
228        );
229        assert_eq!(p.selected(), Some("http://a"));
230        assert_eq!(p.selected_country(), de);
231        assert_eq!(p.rotate().as_deref(), Some("http://b"));
232        assert_eq!(p.selected_country(), fr);
233    }
234
235    #[test]
236    fn untagged_pool_has_no_country() {
237        let p = pool(&["http://a"], RotationStrategy::RoundRobin);
238        assert_eq!(p.selected_country(), None);
239    }
240
241    #[test]
242    fn random_strategy_selects_a_member_and_exhausts() {
243        let urls = ["http://a", "http://b", "http://c"];
244        let mut p = pool(&urls, RotationStrategy::Random);
245        assert!(urls.contains(&p.selected().unwrap()));
246        // Rotate through all remaining; every result is a valid member, then None.
247        let mut seen = 0;
248        while let Some(url) = p.rotate() {
249            assert!(urls.contains(&url.as_str()));
250            seen += 1;
251        }
252        assert_eq!(seen, 2); // 3 endpoints, initial selection consumed one
253        assert_eq!(p.healthy_count(), 0);
254    }
255}
stealthscraper_rs/proxy_pool.rs

stealthscraper_rs/
proxy_pool.rs
