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}