Skip to main content

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::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}
41
42// ─────────────────────────────────────────────────────────────────────────────
43// RotationStrategy trait
44// ─────────────────────────────────────────────────────────────────────────────
45
46/// Selects a proxy from a slice of candidates on each request.
47///
48/// Implementations receive **all** candidates (healthy and unhealthy) so they
49/// can distinguish between an empty pool and a pool where every proxy is
50/// temporarily down. Call [`healthy_candidates`] to filter the slice.
51///
52/// # Example
53/// ```rust,no_run
54/// use stygian_proxy::strategy::{ProxyCandidate, RotationStrategy, RoundRobinStrategy};
55///
56/// async fn pick(candidates: &[ProxyCandidate]) {
57///     let strategy = RoundRobinStrategy::default();
58///     let chosen = strategy.select(candidates).await.unwrap();
59///     println!("selected: {}", chosen.id);
60/// }
61/// ```
62#[async_trait]
63pub trait RotationStrategy: Send + Sync + 'static {
64    /// Select one candidate from `candidates`.
65    ///
66    /// Returns [`crate::error::ProxyError::AllProxiesUnhealthy`] when every candidate has
67    /// `healthy == false`.
68    async fn select<'a>(&self, candidates: &'a [ProxyCandidate])
69    -> ProxyResult<&'a ProxyCandidate>;
70}
71
72/// Shared-ownership type alias for a [`RotationStrategy`] implementation.
73pub type BoxedRotationStrategy = Arc<dyn RotationStrategy>;
74
75// ─────────────────────────────────────────────────────────────────────────────
76// Shared helper
77// ─────────────────────────────────────────────────────────────────────────────
78
79/// Filter `all` to only the candidates that are currently healthy.
80///
81/// Returns references into the original slice, so no allocation is needed
82/// beyond the returned `Vec`.
83pub fn healthy_candidates(all: &[ProxyCandidate]) -> Vec<&ProxyCandidate> {
84    all.iter().filter(|c| c.healthy).collect()
85}
86
87// ─────────────────────────────────────────────────────────────────────────────
88// Tests
89// ─────────────────────────────────────────────────────────────────────────────
90
91#[cfg(test)]
92pub(crate) mod tests {
93    use super::*;
94    use crate::error::ProxyError;
95    use std::sync::atomic::Ordering;
96
97    /// Build a `ProxyCandidate` with sensible test defaults.
98    pub fn candidate(id: u128, healthy: bool, weight: u32, requests: u64) -> ProxyCandidate {
99        let metrics = Arc::new(ProxyMetrics::default());
100        metrics.requests_total.store(requests, Ordering::Relaxed);
101        ProxyCandidate {
102            id: Uuid::from_u128(id),
103            weight,
104            metrics,
105            healthy,
106        }
107    }
108
109    #[tokio::test]
110    async fn healthy_candidates_filters() {
111        let c = vec![
112            candidate(1, true, 1, 0),
113            candidate(2, false, 1, 0),
114            candidate(3, true, 1, 0),
115        ];
116        let healthy = healthy_candidates(&c);
117        assert_eq!(healthy.len(), 2);
118        assert!(healthy.iter().all(|c| c.healthy));
119    }
120
121    #[tokio::test]
122    async fn all_unhealthy_returns_error() {
123        let c = vec![candidate(1, false, 1, 0), candidate(2, false, 1, 0)];
124        let err = RoundRobinStrategy::default().select(&c).await.unwrap_err();
125        assert!(matches!(err, ProxyError::AllProxiesUnhealthy));
126    }
127}