millipede_core/
proxy.rs

1//! Proxy configuration and rotation.
2//!
3//! Tiered configurations implement a simplified Crawlee-compatible policy: blocking escalates a
4//! domain, periodic requests probe the next lower tier, and a `None` slot explicitly means a direct
5//! request without a proxy.
6
7use std::{
8    collections::HashMap,
9    fmt,
10    sync::{
11        Arc, Mutex,
12        atomic::{AtomicUsize, Ordering},
13    },
14};
15
16use url::Url;
17
18use crate::{errors::CrawlError, request::Request, session::SessionId};
19
20/// Selection policy for a static proxy list.
21///
22/// ```
23/// use millipede_core::proxy::RotationStrategy;
24/// assert_eq!(RotationStrategy::default(), RotationStrategy::RoundRobin);
25/// ```
26#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
27#[non_exhaustive]
28pub enum RotationStrategy {
29    /// Selects proxies in stable cyclic order.
30    #[default]
31    RoundRobin,
32    /// Selects a proxy using the crate's lightweight process-local generator.
33    Random,
34}
35
36/// Borrowed inputs available to proxy resolution.
37///
38/// ```
39/// use millipede_core::proxy::ProxyResolveContext;
40/// let context = ProxyResolveContext::new().attempt(2);
41/// assert_eq!(context.attempt, 2);
42/// ```
43#[derive(Debug, Clone, Copy, Default)]
44pub struct ProxyResolveContext<'a> {
45    /// Request being routed, when available.
46    pub request: Option<&'a Request>,
47    /// Checked-out session identifier, when available.
48    pub session_id: Option<&'a SessionId>,
49    /// Zero-based request attempt.
50    pub attempt: u32,
51}
52
53impl<'a> ProxyResolveContext<'a> {
54    /// Creates an empty resolution context.
55    pub fn new() -> Self {
56        Self::default()
57    }
58    /// Sets the current request.
59    pub fn request(mut self, value: &'a Request) -> Self {
60        self.request = Some(value);
61        self
62    }
63    /// Sets the current session identifier.
64    pub fn session_id(mut self, value: &'a SessionId) -> Self {
65        self.session_id = Some(value);
66        self
67    }
68    /// Sets the current attempt.
69    pub fn attempt(mut self, value: u32) -> Self {
70        self.attempt = value;
71        self
72    }
73}
74
75/// Asynchronous custom proxy URL resolver.
76///
77/// ```
78/// # use millipede_core::{errors::CrawlError, proxy::{ProxyResolveContext, ProxyResolver}};
79/// # struct Direct;
80/// # #[async_trait::async_trait]
81/// # impl ProxyResolver for Direct {
82/// #   async fn resolve(&self, _: ProxyResolveContext<'_>) -> Result<Option<url::Url>, CrawlError> { Ok(None) }
83/// # }
84/// ```
85#[async_trait::async_trait]
86pub trait ProxyResolver: Send + Sync + 'static {
87    /// Resolves a proxy URL or chooses a direct request with `None`.
88    async fn resolve(&self, ctx: ProxyResolveContext<'_>) -> Result<Option<Url>, CrawlError>;
89}
90
91/// Parsed connection details for a selected proxy.
92///
93/// ```
94/// use millipede_core::proxy::ProxyInfo;
95/// let info = ProxyInfo::from_url(url::Url::parse("http://proxy.example:8080")?);
96/// assert_eq!(info.port, 8080);
97/// # Ok::<(), url::ParseError>(())
98/// ```
99#[derive(Debug, Clone, PartialEq)]
100#[non_exhaustive]
101pub struct ProxyInfo {
102    /// Complete proxy URL.
103    pub url: Url,
104    /// Proxy host name.
105    pub hostname: String,
106    /// Explicit or scheme-default port.
107    pub port: u16,
108    /// Optional non-empty username.
109    pub username: Option<String>,
110    /// Optional password.
111    pub password: Option<String>,
112    /// Tier that supplied this proxy.
113    pub tier: Option<u8>,
114    /// Session associated with the resolution request.
115    pub session_id: Option<SessionId>,
116}
117
118impl ProxyInfo {
119    /// Parses connection metadata from a URL.
120    pub fn from_url(url: Url) -> Self {
121        let hostname = url.host_str().unwrap_or_default().to_owned();
122        let port = url.port_or_known_default().unwrap_or(80);
123        let username = (!url.username().is_empty()).then(|| url.username().to_owned());
124        let password = url.password().map(str::to_owned);
125        Self {
126            url,
127            hostname,
128            port,
129            username,
130            password,
131            tier: None,
132            session_id: None,
133        }
134    }
135    /// Sets the originating tier.
136    pub fn with_tier(mut self, value: u8) -> Self {
137        self.tier = Some(value);
138        self
139    }
140    /// Sets the associated session.
141    pub fn with_session_id(mut self, value: SessionId) -> Self {
142        self.session_id = Some(value);
143        self
144    }
145}
146
147enum ProxyInner {
148    Static {
149        urls: Vec<Url>,
150        rotation: RotationStrategy,
151        cursor: AtomicUsize,
152    },
153    Custom(Arc<dyn ProxyResolver>),
154    Tiered(TieredState),
155}
156
157struct TieredState {
158    tiers: Vec<Vec<Option<Url>>>,
159    probe_interval: u32,
160    domains: Mutex<HashMap<String, DomainTier>>,
161}
162
163#[derive(Default)]
164struct DomainTier {
165    tier: usize,
166    requests: u32,
167    probing: Option<usize>,
168}
169
170/// Static, custom, or per-domain tiered proxy selection.
171///
172/// ```
173/// use millipede_core::proxy::ProxyConfiguration;
174/// let config = ProxyConfiguration::round_robin([url::Url::parse("http://proxy.example")?]);
175/// assert!(format!("{config:?}").contains("Static"));
176/// # Ok::<(), url::ParseError>(())
177/// ```
178pub struct ProxyConfiguration {
179    inner: ProxyInner,
180}
181
182impl ProxyConfiguration {
183    /// Creates a round-robin static list.
184    pub fn round_robin(urls: impl IntoIterator<Item = Url>) -> Self {
185        Self::rotating(urls, RotationStrategy::RoundRobin)
186    }
187    /// Creates a static list using `rotation`.
188    pub fn rotating(urls: impl IntoIterator<Item = Url>, rotation: RotationStrategy) -> Self {
189        Self {
190            inner: ProxyInner::Static {
191                urls: urls.into_iter().collect(),
192                rotation,
193                cursor: AtomicUsize::new(0),
194            },
195        }
196    }
197    /// Creates a configuration backed by a custom resolver.
198    pub fn custom<R: ProxyResolver>(resolver: R) -> Self {
199        Self {
200            inner: ProxyInner::Custom(Arc::new(resolver)),
201        }
202    }
203    /// Creates tiered rotation with a probe every 20 requests.
204    pub fn tiered(tiers: Vec<Vec<Option<Url>>>) -> Self {
205        Self::tiered_with_probe_interval(tiers, 20)
206    }
207    /// Creates tiered rotation with the requested probe interval.
208    pub fn tiered_with_probe_interval(tiers: Vec<Vec<Option<Url>>>, probe_interval: u32) -> Self {
209        Self {
210            inner: ProxyInner::Tiered(TieredState {
211                tiers,
212                probe_interval: probe_interval.max(1),
213                domains: Mutex::new(HashMap::new()),
214            }),
215        }
216    }
217
218    fn tiered_url(state: &TieredState, ctx: ProxyResolveContext<'_>) -> (Option<Url>, Option<u8>) {
219        if state.tiers.is_empty() {
220            return (None, None);
221        }
222        let key = ctx
223            .request
224            .and_then(|request| request.url.host_str())
225            .unwrap_or_default()
226            .to_owned();
227        let mut domains = state.domains.lock().unwrap_or_else(|e| e.into_inner());
228        let domain = domains.entry(key).or_default();
229        domain.tier = domain.tier.min(state.tiers.len() - 1);
230        domain.requests = domain.requests.saturating_add(1);
231        let serving = if domain.tier > 0 && domain.requests % state.probe_interval == 0 {
232            let probe = domain.tier - 1;
233            domain.probing = Some(probe);
234            probe
235        } else {
236            domain.tier
237        };
238        let tier = &state.tiers[serving];
239        if tier.is_empty() {
240            return (None, Some(serving as u8));
241        }
242        (
243            tier[domain.requests as usize % tier.len()].clone(),
244            Some(serving as u8),
245        )
246    }
247
248    async fn resolve(
249        &self,
250        ctx: ProxyResolveContext<'_>,
251    ) -> Result<(Option<Url>, Option<u8>), CrawlError> {
252        match &self.inner {
253            ProxyInner::Static {
254                urls,
255                rotation,
256                cursor,
257            } => {
258                if urls.is_empty() {
259                    return Ok((None, None));
260                }
261                let index = match rotation {
262                    RotationStrategy::RoundRobin => cursor.fetch_add(1, Ordering::Relaxed),
263                    RotationStrategy::Random => crate::util::rand_u64() as usize,
264                } % urls.len();
265                Ok((Some(urls[index].clone()), None))
266            }
267            ProxyInner::Custom(resolver) => resolver.resolve(ctx).await.map(|url| (url, None)),
268            ProxyInner::Tiered(state) => Ok(Self::tiered_url(state, ctx)),
269        }
270    }
271
272    /// Selects the next proxy URL, or `None` for a direct request.
273    pub async fn new_url(&self, ctx: ProxyResolveContext<'_>) -> Result<Option<Url>, CrawlError> {
274        self.resolve(ctx).await.map(|(url, _)| url)
275    }
276    /// Selects and parses the next proxy, retaining tier and session metadata.
277    pub async fn new_proxy_info(
278        &self,
279        ctx: ProxyResolveContext<'_>,
280    ) -> Result<Option<ProxyInfo>, CrawlError> {
281        let session_id = ctx.session_id.cloned();
282        let (url, tier) = self.resolve(ctx).await?;
283        Ok(url.map(|url| {
284            let mut info = ProxyInfo::from_url(url);
285            info.tier = tier;
286            info.session_id = session_id;
287            info
288        }))
289    }
290    /// Reports that `target` was blocked, escalating its domain unless a recovery probe failed.
291    pub fn report_blocked(&self, target: &Url) {
292        let ProxyInner::Tiered(state) = &self.inner else {
293            return;
294        };
295        if state.tiers.is_empty() {
296            return;
297        }
298        let key = target.host_str().unwrap_or_default().to_owned();
299        let mut domains = state.domains.lock().unwrap_or_else(|e| e.into_inner());
300        let domain = domains.entry(key).or_default();
301        if domain.probing.take().is_none() {
302            domain.tier = (domain.tier + 1).min(state.tiers.len() - 1);
303            domain.requests = 0;
304        }
305    }
306    /// Reports success for `target`, accepting a pending lower-tier recovery probe.
307    pub fn report_success(&self, target: &Url) {
308        let ProxyInner::Tiered(state) = &self.inner else {
309            return;
310        };
311        let key = target.host_str().unwrap_or_default().to_owned();
312        let mut domains = state.domains.lock().unwrap_or_else(|e| e.into_inner());
313        let domain = domains.entry(key).or_default();
314        if let Some(probe) = domain.probing.take() {
315            domain.tier = probe;
316            domain.requests = 0;
317        }
318    }
319}
320
321impl fmt::Debug for ProxyConfiguration {
322    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
323        let variant = match self.inner {
324            ProxyInner::Static { .. } => "Static",
325            ProxyInner::Custom(_) => "Custom",
326            ProxyInner::Tiered(_) => "Tiered",
327        };
328        formatter
329            .debug_struct("ProxyConfiguration")
330            .field("variant", &variant)
331            .finish()
332    }
333}
334
335/// Synchronous policy selecting a proxy bucket for a request.
336///
337/// ```
338/// # use millipede_core::proxy::{ProxyKind, ProxyRouteContext, ProxyStrategy};
339/// # struct DefaultRoute;
340/// # impl ProxyStrategy for DefaultRoute { fn route(&self, _: &ProxyRouteContext<'_>) -> ProxyKind { ProxyKind::Default } }
341/// ```
342pub trait ProxyStrategy: Send + Sync + 'static {
343    /// Selects the logical proxy bucket.
344    fn route(&self, ctx: &ProxyRouteContext<'_>) -> ProxyKind;
345}
346
347/// Borrowed inputs available to a [`ProxyStrategy`].
348///
349/// ```
350/// # use millipede_core::{proxy::ProxyRouteContext, request::Request};
351/// # let request = Request::get("https://example.com").build()?;
352/// let context = ProxyRouteContext::new(&request, 0).previous_profile_key("old");
353/// assert_eq!(context.previous_profile_key, Some("old"));
354/// # Ok::<(), millipede_core::request::RequestBuildError>(())
355/// ```
356pub struct ProxyRouteContext<'a> {
357    /// Request being routed.
358    pub request: &'a Request,
359    /// Zero-based attempt.
360    pub attempt: u32,
361    /// Profile selected on the previous attempt, when any.
362    pub previous_profile_key: Option<&'a str>,
363}
364
365impl<'a> ProxyRouteContext<'a> {
366    /// Creates a route context.
367    pub fn new(request: &'a Request, attempt: u32) -> Self {
368        Self {
369            request,
370            attempt,
371            previous_profile_key: None,
372        }
373    }
374    /// Sets the prior profile key.
375    pub fn previous_profile_key(mut self, value: &'a str) -> Self {
376        self.previous_profile_key = Some(value);
377        self
378    }
379}
380
381/// Logical proxy configuration bucket.
382///
383/// ```
384/// use millipede_core::proxy::ProxyKind;
385/// assert_eq!(ProxyKind::default(), ProxyKind::Default);
386/// ```
387#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
388#[non_exhaustive]
389pub enum ProxyKind {
390    /// General-purpose proxy configuration.
391    #[default]
392    Default,
393    /// Proxy configuration specialized for media assets.
394    MediaAsset,
395    /// User-named proxy configuration.
396    Custom(String),
397}
398
399/// Named proxy configurations with deterministic fallbacks.
400///
401/// ```
402/// use millipede_core::proxy::{ProxyBuckets, ProxyKind};
403/// assert!(ProxyBuckets::new().for_kind(&ProxyKind::Default).is_none());
404/// ```
405#[derive(Default)]
406#[must_use = "proxy buckets do nothing unless installed on a crawler"]
407pub struct ProxyBuckets {
408    default_bucket: Option<ProxyConfiguration>,
409    media: Option<ProxyConfiguration>,
410    custom: HashMap<String, ProxyConfiguration>,
411}
412
413impl ProxyBuckets {
414    /// Creates empty buckets.
415    pub fn new() -> Self {
416        Self::default()
417    }
418    /// Sets the default bucket.
419    pub fn with_default(mut self, value: ProxyConfiguration) -> Self {
420        self.default_bucket = Some(value);
421        self
422    }
423    /// Sets the media bucket.
424    pub fn with_media(mut self, value: ProxyConfiguration) -> Self {
425        self.media = Some(value);
426        self
427    }
428    /// Inserts or replaces a named custom bucket.
429    pub fn with_custom(mut self, name: impl Into<String>, value: ProxyConfiguration) -> Self {
430        self.custom.insert(name.into(), value);
431        self
432    }
433    /// Finds a bucket, falling media and unknown custom names back to the default.
434    pub fn for_kind(&self, kind: &ProxyKind) -> Option<&ProxyConfiguration> {
435        match kind {
436            ProxyKind::Default => self.default_bucket.as_ref(),
437            ProxyKind::MediaAsset => self.media.as_ref().or(self.default_bucket.as_ref()),
438            ProxyKind::Custom(name) => self.custom.get(name).or(self.default_bucket.as_ref()),
439        }
440    }
441}
442
443impl fmt::Debug for ProxyBuckets {
444    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
445        formatter
446            .debug_struct("ProxyBuckets")
447            .field("default_bucket", &self.default_bucket)
448            .field("media", &self.media)
449            .field("custom", &self.custom)
450            .finish()
451    }
452}
millipede_core/proxy.rs

millipede_core/
proxy.rs
