1use std::collections::HashMap;
7use std::sync::Arc;
8use std::sync::atomic::{AtomicBool, Ordering};
9#[cfg(feature = "vendor-stickiness")]
10use std::time::Duration;
11
12use serde::Serialize;
13use tokio::sync::RwLock;
14use tokio::task::JoinHandle;
15use tokio_util::sync::CancellationToken;
16use uuid::Uuid;
17
18use crate::circuit_breaker::CircuitBreaker;
19use crate::error::{ProxyError, ProxyResult};
20use crate::health::{HealthChecker, HealthMap};
21#[cfg(feature = "coherence-validation")]
22use crate::ports::coherence::{
23 BoxedCoherencePort, CoherenceContext, CoherencePolicy, CoherenceVerdict,
24};
25#[cfg(feature = "vendor-stickiness")]
26use crate::session::SessionDecision;
27use crate::session::{SessionMap, StickyPolicy};
28#[cfg(feature = "vendor-stickiness")]
29use crate::stickiness::VendorStickinessMap;
30use crate::storage::ProxyStoragePort;
31use crate::strategy::{
32 BoxedBayesianObserver, BoxedRotationStrategy, LeastUsedStrategy, NoopBayesianObserver,
33 ProxyCandidate, RandomStrategy, RoundRobinStrategy, WeightedStrategy,
34 capable_healthy_candidates,
35};
36#[cfg(feature = "vendor-stickiness")]
37use crate::types::VendorId;
38use crate::types::{CapabilityRequirement, Proxy, ProxyConfig};
39
40#[derive(Debug, Serialize)]
46pub struct PoolStats {
47 pub total: usize,
49 pub healthy: usize,
51 pub open: usize,
53 pub active_sessions: usize,
55}
56
57pub struct ProxyHandle {
67 pub proxy_url: String,
69 circuit_breaker: Arc<CircuitBreaker>,
70 succeeded: AtomicBool,
71 session_key: Option<String>,
73 sessions: Option<SessionMap>,
74 proxy_id: Uuid,
77 observer: BoxedBayesianObserver,
81}
82
83impl ProxyHandle {
84 const fn new(
85 proxy_url: String,
86 circuit_breaker: Arc<CircuitBreaker>,
87 proxy_id: Uuid,
88 observer: BoxedBayesianObserver,
89 ) -> Self {
90 Self {
91 proxy_url,
92 circuit_breaker,
93 succeeded: AtomicBool::new(false),
94 session_key: None,
95 sessions: None,
96 proxy_id,
97 observer,
98 }
99 }
100
101 const fn new_sticky(
102 proxy_url: String,
103 circuit_breaker: Arc<CircuitBreaker>,
104 session_key: String,
105 sessions: SessionMap,
106 proxy_id: Uuid,
107 observer: BoxedBayesianObserver,
108 ) -> Self {
109 Self {
110 proxy_url,
111 circuit_breaker,
112 succeeded: AtomicBool::new(false),
113 session_key: Some(session_key),
114 sessions: Some(sessions),
115 proxy_id,
116 observer,
117 }
118 }
119
120 #[must_use]
125 pub fn direct() -> Self {
126 let noop_cb = Arc::new(CircuitBreaker::new(u32::MAX, u64::MAX));
127 let noop_observer: BoxedBayesianObserver = Arc::new(NoopBayesianObserver);
128 Self {
129 proxy_url: String::new(),
130 circuit_breaker: noop_cb,
131 succeeded: AtomicBool::new(true),
132 session_key: None,
133 sessions: None,
134 proxy_id: Uuid::nil(),
135 observer: noop_observer,
136 }
137 }
138
139 pub fn mark_success(&self) {
141 self.succeeded.store(true, Ordering::Release);
142 self.observer.observe(self.proxy_id, true);
145 }
146}
147
148impl std::fmt::Debug for ProxyHandle {
149 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
150 f.debug_struct("ProxyHandle")
151 .field("proxy_url", &self.proxy_url)
152 .finish_non_exhaustive()
153 }
154}
155
156impl Drop for ProxyHandle {
157 fn drop(&mut self) {
158 if self.succeeded.load(Ordering::Acquire) {
159 self.circuit_breaker.record_success();
160 } else {
161 self.circuit_breaker.record_failure();
162 if let (Some(key), Some(sessions)) = (&self.session_key, &self.sessions) {
164 sessions.unbind(key);
165 }
166 self.observer.observe(self.proxy_id, false);
167 }
168 }
169}
170
171pub struct ProxyManager {
212 storage: Arc<dyn ProxyStoragePort>,
213 strategy: BoxedRotationStrategy,
214 health_checker: HealthChecker,
215 circuit_breakers: Arc<RwLock<HashMap<Uuid, Arc<CircuitBreaker>>>>,
216 config: ProxyConfig,
217 sessions: SessionMap,
219 observer: BoxedBayesianObserver,
224 #[cfg(feature = "coherence-validation")]
232 coherence_validator: Option<BoxedCoherencePort>,
233 #[cfg(feature = "vendor-stickiness")]
242 stickiness_map: VendorStickinessMap,
243}
244
245impl ProxyManager {
246 #[must_use]
248 pub fn builder() -> ProxyManagerBuilder {
249 ProxyManagerBuilder::default()
250 }
251
252 pub fn with_round_robin(
259 storage: Arc<dyn ProxyStoragePort>,
260 config: ProxyConfig,
261 ) -> ProxyResult<Self> {
262 Self::builder()
263 .storage(storage)
264 .strategy(Arc::new(RoundRobinStrategy::default()))
265 .config(config)
266 .build()
267 }
268
269 pub fn with_random(
276 storage: Arc<dyn ProxyStoragePort>,
277 config: ProxyConfig,
278 ) -> ProxyResult<Self> {
279 Self::builder()
280 .storage(storage)
281 .strategy(Arc::new(RandomStrategy))
282 .config(config)
283 .build()
284 }
285
286 pub fn with_weighted(
293 storage: Arc<dyn ProxyStoragePort>,
294 config: ProxyConfig,
295 ) -> ProxyResult<Self> {
296 Self::builder()
297 .storage(storage)
298 .strategy(Arc::new(WeightedStrategy))
299 .config(config)
300 .build()
301 }
302
303 pub fn with_least_used(
310 storage: Arc<dyn ProxyStoragePort>,
311 config: ProxyConfig,
312 ) -> ProxyResult<Self> {
313 Self::builder()
314 .storage(storage)
315 .strategy(Arc::new(LeastUsedStrategy))
316 .config(config)
317 .build()
318 }
319
320 #[cfg(feature = "bayesian-rotation")]
336 pub fn with_thompson_sampling(
337 storage: Arc<dyn ProxyStoragePort>,
338 config: ProxyConfig,
339 decay_interval: std::time::Duration,
340 ) -> ProxyResult<Self> {
341 Self::builder()
342 .storage(storage)
343 .config(config)
344 .with_thompson_sampling(decay_interval)
345 .build()
346 }
347
348 #[allow(clippy::significant_drop_tightening)]
367 pub async fn add_proxy(&self, proxy: Proxy) -> ProxyResult<Uuid> {
368 let mut cb_map = self.circuit_breakers.write().await;
369 let record = self.storage.add(proxy).await?;
370 cb_map.insert(
371 record.id,
372 Arc::new(CircuitBreaker::new(
373 self.config.circuit_open_threshold,
374 u64::try_from(self.config.circuit_half_open_after.as_millis()).unwrap_or(u64::MAX),
375 )),
376 );
377 Ok(record.id)
378 }
379
380 pub async fn remove_proxy(&self, id: Uuid) -> ProxyResult<()> {
387 self.storage.remove(id).await?;
388 self.circuit_breakers.write().await.remove(&id);
389 Ok(())
390 }
391
392 #[allow(clippy::significant_drop_tightening)]
442 pub async fn add_proxy_with_metadata(
443 &self,
444 url: &str,
445 asn: Option<u32>,
446 city: Option<&str>,
447 postal_code: Option<&str>,
448 ) -> ProxyResult<Uuid> {
449 let capabilities = crate::types::ProxyCapabilities {
450 asn,
451 city: city.map(str::to_owned),
452 postal_code: postal_code.map(str::to_owned),
453 ..Default::default()
454 };
455 let proxy = Proxy {
456 url: url.to_owned(),
457 proxy_type: crate::types::ProxyType::Http,
458 username: None,
459 password: None,
460 weight: 1,
461 tags: Vec::new(),
462 capabilities,
463 ip_class: crate::types::IpClass::Unknown,
464 target_compatibility: crate::types::TargetVendorCompatibility::default(),
465 };
466 self.add_proxy(proxy).await
467 }
468
469 #[must_use]
476 pub fn start(&self) -> (CancellationToken, JoinHandle<()>) {
477 let token = CancellationToken::new();
478 let health_handle = self.health_checker.clone().spawn(token.clone());
479
480 let sessions = self.sessions.clone();
481 let purge_token = token.clone();
482 let purge_handle = tokio::spawn(async move {
483 let mut interval = tokio::time::interval(std::time::Duration::from_mins(1));
484 loop {
485 tokio::select! {
486 _ = interval.tick() => { let _ = sessions.purge_expired(); }
487 () = purge_token.cancelled() => break,
488 }
489 }
490 });
491
492 let combined = tokio::spawn(async move {
493 let _ = tokio::join!(health_handle, purge_handle);
494 });
495
496 (token, combined)
497 }
498
499 pub fn strategy_warmup_observe(&self, proxy_id: Uuid, success: bool) {
507 self.observer.observe(proxy_id, success);
508 }
509
510 #[must_use]
514 pub fn storage(&self) -> &Arc<dyn ProxyStoragePort> {
515 &self.storage
516 }
517
518 #[allow(clippy::significant_drop_tightening)]
524 async fn select_proxy_inner(&self) -> ProxyResult<(String, Arc<CircuitBreaker>, Uuid)> {
525 let with_metrics = self.storage.list_with_metrics().await?;
526 if with_metrics.is_empty() {
527 return Err(ProxyError::PoolExhausted);
528 }
529
530 let candidates = {
533 let health_map_ref = Arc::clone(self.health_checker.health_map());
534 let health_map = health_map_ref.read().await;
535 let cb_map_ref = Arc::clone(&self.circuit_breakers);
536 let cb_map = cb_map_ref.read().await;
537 let candidates: Vec<ProxyCandidate> = with_metrics
538 .iter()
539 .map(|(record, metrics)| {
540 let healthy = health_map.get(&record.id).copied().unwrap_or(true);
541 let available = cb_map.get(&record.id).is_none_or(|cb| cb.is_available());
542 ProxyCandidate {
543 id: record.id,
544 weight: record.proxy.weight,
545 metrics: Arc::clone(metrics),
546 healthy: healthy && available,
547 capabilities: record.proxy.capabilities.clone(),
548 }
549 })
550 .collect();
551 candidates
552 };
554
555 let selected = self.strategy.select(&candidates).await?;
556 let id = selected.id;
557
558 let cb = self
560 .circuit_breakers
561 .read()
562 .await
563 .get(&id)
564 .cloned()
565 .ok_or(ProxyError::PoolExhausted)?;
566 let url = with_metrics
567 .iter()
568 .find(|(r, _)| r.id == id)
569 .map(|(r, _)| r.proxy.url.clone())
570 .unwrap_or_default();
571
572 Ok((url, cb, id))
573 }
574
575 pub async fn acquire_proxy(&self) -> ProxyResult<ProxyHandle> {
587 let (url, cb, id) = self.select_proxy_inner().await?;
588 Ok(ProxyHandle::new(url, cb, id, Arc::clone(&self.observer)))
589 }
590
591 pub async fn acquire_with_capabilities(
617 &self,
618 req: &CapabilityRequirement,
619 ) -> ProxyResult<ProxyHandle> {
620 let with_metrics = self.storage.list_with_metrics().await?;
621
622 if with_metrics.is_empty() {
623 return Err(ProxyError::PoolExhausted);
624 }
625
626 let candidates = {
627 let health_map_ref = Arc::clone(self.health_checker.health_map());
628 let health_map = health_map_ref.read().await;
629 let cb_map_ref = Arc::clone(&self.circuit_breakers);
630 let cb_map = cb_map_ref.read().await;
631 let candidates: Vec<ProxyCandidate> = with_metrics
632 .iter()
633 .map(|(record, metrics)| {
634 let healthy = health_map.get(&record.id).copied().unwrap_or(true);
635 let available = cb_map.get(&record.id).is_none_or(|cb| cb.is_available());
636 ProxyCandidate {
637 id: record.id,
638 weight: record.proxy.weight,
639 metrics: Arc::clone(metrics),
640 healthy: healthy && available,
641 capabilities: record.proxy.capabilities.clone(),
642 }
643 })
644 .collect();
645 candidates
646 };
647
648 let compatible: Vec<ProxyCandidate> = capable_healthy_candidates(&candidates, req)
650 .into_iter()
651 .cloned()
652 .collect();
653 if compatible.is_empty() {
654 return Err(ProxyError::NoCompatibleProxy);
655 }
656
657 let selected = self.strategy.select(&compatible).await?;
658 let id = selected.id;
659
660 let cb = self
661 .circuit_breakers
662 .read()
663 .await
664 .get(&id)
665 .cloned()
666 .ok_or(ProxyError::PoolExhausted)?;
667 let url = with_metrics
668 .iter()
669 .find(|(r, _)| r.id == id)
670 .map(|(r, _)| r.proxy.url.clone())
671 .unwrap_or_default();
672
673 Ok(ProxyHandle::new(url, cb, id, Arc::clone(&self.observer)))
674 }
675
676 pub async fn acquire_for_domain(&self, domain: &str) -> ProxyResult<ProxyHandle> {
697 let ttl = match &self.config.sticky_policy {
698 StickyPolicy::Disabled => return self.acquire_proxy().await,
699 StickyPolicy::Domain { ttl } => *ttl,
700 };
701
702 if let Some(proxy_id) = self.sessions.lookup(domain) {
704 let cb_map = self.circuit_breakers.read().await;
705 if let Some(cb) = cb_map.get(&proxy_id).cloned()
706 && cb.is_available()
707 {
708 let with_metrics = self.storage.list_with_metrics().await?;
710 if let Some((record, _)) = with_metrics.iter().find(|(r, _)| r.id == proxy_id) {
711 let url = record.proxy.url.clone();
712 drop(cb_map);
713 return Ok(ProxyHandle::new_sticky(
714 url,
715 cb,
716 domain.to_string(),
717 self.sessions.clone(),
718 proxy_id,
719 Arc::clone(&self.observer),
720 ));
721 }
722 }
723 drop(cb_map);
725 self.sessions.unbind(domain);
726 }
727
728 let (url, cb, proxy_id) = self.select_proxy_inner().await?;
730 self.sessions.bind(domain, proxy_id, ttl);
731 Ok(ProxyHandle::new_sticky(
732 url,
733 cb,
734 domain.to_string(),
735 self.sessions.clone(),
736 proxy_id,
737 Arc::clone(&self.observer),
738 ))
739 }
740
741 #[cfg(feature = "vendor-stickiness")]
772 pub async fn acquire_for_domain_with_vendor(
773 &self,
774 domain: &str,
775 vendor: VendorId,
776 ) -> ProxyResult<ProxyHandle> {
777 let decision = self
778 .sessions
779 .acquire_session(domain, vendor, &self.stickiness_map);
780 match decision {
781 SessionDecision::UseSticky(proxy_id) => {
782 let cb = self.lookup_validated_cb(proxy_id).await?;
787 if let Some(cb) = cb {
788 let url = self.lookup_url(proxy_id).await?;
789 if let Some(url) = url {
790 return Ok(ProxyHandle::new_sticky(
791 url,
792 cb,
793 domain.to_string(),
794 self.sessions.clone(),
795 proxy_id,
796 Arc::clone(&self.observer),
797 ));
798 }
799 }
800 self.sessions.unbind(domain);
802 let (url, cb, proxy_id) = self.select_proxy_inner().await?;
803 let ttl = self.stickiness_ttl(vendor);
808 self.sessions.bind(domain, proxy_id, ttl);
809 Ok(ProxyHandle::new_sticky(
810 url,
811 cb,
812 domain.to_string(),
813 self.sessions.clone(),
814 proxy_id,
815 Arc::clone(&self.observer),
816 ))
817 }
818 SessionDecision::AcquireFresh => self.acquire_proxy().await,
819 SessionDecision::AcquireAndBind(ttl) => {
820 let (url, cb, proxy_id) = self.select_proxy_inner().await?;
821 self.sessions.bind(domain, proxy_id, ttl);
822 Ok(ProxyHandle::new_sticky(
823 url,
824 cb,
825 domain.to_string(),
826 self.sessions.clone(),
827 proxy_id,
828 Arc::clone(&self.observer),
829 ))
830 }
831 }
832 }
833
834 #[cfg(feature = "vendor-stickiness")]
840 fn stickiness_ttl(&self, vendor: VendorId) -> Duration {
841 use crate::stickiness::StickinessPolicy;
842 match self.stickiness_map.for_vendor(vendor) {
843 StickinessPolicy::StickyForever => Duration::MAX,
844 StickinessPolicy::StickyForTtl { ttl } => ttl,
845 _ => Duration::from_mins(30),
847 }
848 }
849
850 #[cfg(feature = "vendor-stickiness")]
854 #[allow(clippy::significant_drop_tightening)]
855 async fn lookup_validated_cb(
856 &self,
857 proxy_id: Uuid,
858 ) -> ProxyResult<Option<Arc<CircuitBreaker>>> {
859 let cb_map = self.circuit_breakers.read().await;
860 let Some(cb) = cb_map.get(&proxy_id).cloned() else {
861 return Ok(None);
862 };
863 if !cb.is_available() {
864 return Ok(None);
865 }
866 Ok(Some(cb))
867 }
868
869 #[cfg(feature = "vendor-stickiness")]
872 async fn lookup_url(&self, proxy_id: Uuid) -> ProxyResult<Option<String>> {
873 let with_metrics = self.storage.list_with_metrics().await?;
874 Ok(with_metrics
875 .iter()
876 .find(|(r, _)| r.id == proxy_id)
877 .map(|(r, _)| r.proxy.url.clone()))
878 }
879
880 pub async fn pool_stats(&self) -> ProxyResult<PoolStats> {
889 let records = self.storage.list().await?;
890 let total = records.len();
891 let health_map = self.health_checker.health_map().read().await;
892 let cb_map = self.circuit_breakers.read().await;
893
894 let mut healthy = 0usize;
895 let mut open = 0usize;
896 for r in &records {
897 if health_map.get(&r.id).copied().unwrap_or(true) {
898 healthy += 1;
899 }
900 if cb_map.get(&r.id).is_some_and(|cb| !cb.is_available()) {
901 open += 1;
902 }
903 }
904 drop(health_map);
905 drop(cb_map);
906 Ok(PoolStats {
907 total,
908 healthy,
909 open,
910 active_sessions: self.sessions.active_count(),
911 })
912 }
913
914 #[cfg(feature = "coherence-validation")]
948 pub async fn acquire_proxy_with_coherence(
949 &self,
950 ctx: &CoherenceContext,
951 policy: &CoherencePolicy,
952 ) -> ProxyResult<ProxyHandle> {
953 let validator = self.coherence_validator.as_ref().ok_or_else(|| {
954 ProxyError::ConfigError(
955 "ProxyManager::acquire_proxy_with_coherence: no coherence_validator wired in; enable the `coherence-validation` cargo feature or call ProxyManagerBuilder::coherence_validator(...)"
956 .into(),
957 )
958 })?;
959
960 let handle = self.acquire_proxy().await?;
961 let verdict = validator.evaluate(ctx);
962 match verdict {
963 CoherenceVerdict::Coherent => Ok(handle),
964 CoherenceVerdict::Mismatch { field, severity } => {
965 if policy.is_hard_fail(field) && severity.is_hard() {
966 Err(ProxyError::CoherenceMismatch { field, severity })
967 } else {
968 tracing::warn!(
969 target: "stygian_proxy::coherence",
970 field = %field,
971 severity = %severity,
972 "coherence mismatch (advisory) — proceeding with the selected proxy"
973 );
974 Ok(handle)
975 }
976 }
977 CoherenceVerdict::Unknown(reason) => {
978 tracing::debug!(
979 target: "stygian_proxy::coherence",
980 reason = %reason,
981 "coherence verdict unknown — proceeding with the selected proxy"
982 );
983 Ok(handle)
984 }
985 }
986 }
987}
988
989#[derive(Default)]
995pub struct ProxyManagerBuilder {
996 storage: Option<Arc<dyn ProxyStoragePort>>,
997 strategy: Option<BoxedRotationStrategy>,
998 config: Option<ProxyConfig>,
999 observer: Option<BoxedBayesianObserver>,
1002 #[cfg(feature = "coherence-validation")]
1007 coherence_validator: Option<BoxedCoherencePort>,
1008 #[cfg(feature = "vendor-stickiness")]
1013 stickiness_map: Option<VendorStickinessMap>,
1014}
1015
1016impl ProxyManagerBuilder {
1017 #[must_use]
1018 pub fn storage(mut self, s: Arc<dyn ProxyStoragePort>) -> Self {
1019 self.storage = Some(s);
1020 self
1021 }
1022
1023 #[must_use]
1024 pub fn strategy(mut self, s: BoxedRotationStrategy) -> Self {
1025 self.strategy = Some(s);
1026 self
1027 }
1028
1029 #[must_use]
1030 pub fn config(mut self, c: ProxyConfig) -> Self {
1031 self.config = Some(c);
1032 self
1033 }
1034
1035 #[must_use]
1039 pub fn observer(mut self, o: BoxedBayesianObserver) -> Self {
1040 self.observer = Some(o);
1041 self
1042 }
1043
1044 #[cfg(feature = "coherence-validation")]
1051 #[must_use]
1052 pub fn coherence_validator(mut self, validator: BoxedCoherencePort) -> Self {
1053 self.coherence_validator = Some(validator);
1054 self
1055 }
1056
1057 #[cfg(feature = "bayesian-rotation")]
1065 #[must_use]
1066 pub fn with_thompson_sampling(mut self, decay_interval: std::time::Duration) -> Self {
1067 let strategy = Arc::new(crate::strategy::ThompsonStrategy::with_decay(
1068 decay_interval,
1069 crate::strategy::thompson::DEFAULT_DECAY_FACTOR,
1070 ));
1071 let observer: BoxedBayesianObserver = Arc::clone(&strategy) as BoxedBayesianObserver;
1078 self.strategy = Some(strategy);
1079 self.observer = Some(observer);
1080 self
1081 }
1082
1083 #[cfg(feature = "vendor-stickiness")]
1093 #[must_use]
1094 pub fn stickiness_map(mut self, map: VendorStickinessMap) -> Self {
1095 self.stickiness_map = Some(map);
1096 self
1097 }
1098
1099 pub fn build(self) -> ProxyResult<ProxyManager> {
1111 let storage = self.storage.ok_or_else(|| {
1112 ProxyError::ConfigError("ProxyManagerBuilder: storage is required".into())
1113 })?;
1114 let strategy = self
1115 .strategy
1116 .unwrap_or_else(|| Arc::new(RoundRobinStrategy::default()));
1117 let observer = self
1118 .observer
1119 .unwrap_or_else(|| Arc::new(NoopBayesianObserver));
1120 let config = self.config.unwrap_or_default();
1121 let health_map: HealthMap = Arc::new(RwLock::new(HashMap::new()));
1122 let checker = HealthChecker::new(
1123 config.clone(),
1124 Arc::clone(&storage),
1125 Arc::clone(&health_map),
1126 );
1127
1128 #[cfg(feature = "tls-profiled")]
1129 let health_checker = if let Some(mode) = config.profiled_request_mode {
1130 checker.with_profiled_mode(mode)?
1131 } else {
1132 checker
1133 };
1134
1135 #[cfg(not(feature = "tls-profiled"))]
1136 let health_checker = checker;
1137
1138 Ok(ProxyManager {
1139 storage,
1140 strategy,
1141 health_checker,
1142 circuit_breakers: Arc::new(RwLock::new(HashMap::new())),
1143 config,
1144 sessions: SessionMap::new(),
1145 observer,
1146 #[cfg(feature = "coherence-validation")]
1147 coherence_validator: self.coherence_validator.or_else(|| {
1148 Some(std::sync::Arc::new(
1149 crate::adapters::coherence::DefaultCoherenceValidator,
1150 ))
1151 }),
1152 #[cfg(feature = "vendor-stickiness")]
1153 stickiness_map: self
1154 .stickiness_map
1155 .unwrap_or_else(VendorStickinessMap::with_builtin_defaults),
1156 })
1157 }
1158}
1159
1160#[cfg(test)]
1165#[allow(
1166 clippy::unwrap_used,
1167 clippy::expect_used,
1168 clippy::significant_drop_tightening,
1169 clippy::manual_let_else,
1170 clippy::panic,
1171 clippy::indexing_slicing
1172)]
1173mod tests {
1174 use std::collections::HashSet;
1175 use std::time::Duration;
1176
1177 use super::*;
1178 use crate::circuit_breaker::{STATE_CLOSED, STATE_OPEN};
1179 use crate::storage::MemoryProxyStore;
1180 use crate::types::ProxyType;
1181
1182 fn make_proxy(url: &str) -> Proxy {
1183 Proxy {
1184 url: url.into(),
1185 proxy_type: ProxyType::Http,
1186 username: None,
1187 password: None,
1188 weight: 1,
1189 tags: vec![],
1190 capabilities: crate::types::ProxyCapabilities::default(),
1191 ip_class: crate::types::IpClass::Unknown,
1192 target_compatibility: crate::types::TargetVendorCompatibility::default(),
1193 }
1194 }
1195
1196 fn storage() -> Arc<MemoryProxyStore> {
1197 Arc::new(MemoryProxyStore::default())
1198 }
1199
1200 #[tokio::test]
1202 async fn round_robin_distribution() {
1203 let store = storage();
1204 let mgr = ProxyManager::with_round_robin(store.clone(), ProxyConfig::default()).unwrap();
1205 mgr.add_proxy(make_proxy("http://a.test:8080"))
1206 .await
1207 .unwrap();
1208 mgr.add_proxy(make_proxy("http://b.test:8080"))
1209 .await
1210 .unwrap();
1211 mgr.add_proxy(make_proxy("http://c.test:8080"))
1212 .await
1213 .unwrap();
1214
1215 let mut seen = HashSet::new();
1216 for _ in 0..10 {
1217 let h = mgr.acquire_proxy().await.unwrap();
1218 h.mark_success();
1219 seen.insert(h.proxy_url.clone());
1220 }
1221 assert_eq!(seen.len(), 3, "all three proxies should have been selected");
1222 }
1223
1224 #[tokio::test]
1232 async fn acquire_proxy_hot_path_budget() {
1233 let store = storage();
1234 let mgr = ProxyManager::with_round_robin(store.clone(), ProxyConfig::default()).unwrap();
1235 mgr.add_proxy(make_proxy("http://a.test:8080"))
1236 .await
1237 .unwrap();
1238 mgr.add_proxy(make_proxy("http://b.test:8080"))
1239 .await
1240 .unwrap();
1241 mgr.add_proxy(make_proxy("http://c.test:8080"))
1242 .await
1243 .unwrap();
1244
1245 let start = std::time::Instant::now();
1246 for _ in 0..1_000 {
1247 let h = mgr.acquire_proxy().await.unwrap();
1248 h.mark_success();
1249 }
1250 let elapsed = start.elapsed();
1251 assert!(
1252 elapsed < std::time::Duration::from_secs(1),
1253 "1000 acquisitions took {elapsed:?}; hot-path budget violated"
1254 );
1255 }
1256
1257 #[tokio::test]
1264 async fn acquire_with_capabilities_hot_path_budget() {
1265 let store = storage();
1266 let mgr = ProxyManager::with_round_robin(store.clone(), ProxyConfig::default()).unwrap();
1267 mgr.add_proxy(make_proxy("http://a.test:8080"))
1268 .await
1269 .unwrap();
1270 mgr.add_proxy(make_proxy("http://b.test:8080"))
1271 .await
1272 .unwrap();
1273 mgr.add_proxy(make_proxy("http://c.test:8080"))
1274 .await
1275 .unwrap();
1276
1277 let req = crate::types::CapabilityRequirement::default();
1280 let start = std::time::Instant::now();
1281 for _ in 0..1_000 {
1282 let h = mgr.acquire_with_capabilities(&req).await.unwrap();
1283 h.mark_success();
1284 }
1285 let elapsed = start.elapsed();
1286 assert!(
1287 elapsed < std::time::Duration::from_secs(1),
1288 "1000 capability-aware acquisitions took {elapsed:?}; hot-path budget violated"
1289 );
1290 }
1291
1292 #[tokio::test]
1294 async fn all_open_returns_error() {
1295 let store = storage();
1296 let mgr = ProxyManager::with_round_robin(
1297 store.clone(),
1298 ProxyConfig {
1299 circuit_open_threshold: 1,
1300 ..ProxyConfig::default()
1301 },
1302 )
1303 .unwrap();
1304 let id = mgr
1305 .add_proxy(make_proxy("http://x.test:8080"))
1306 .await
1307 .unwrap();
1308
1309 {
1311 let map = mgr.circuit_breakers.read().await;
1312 let cb = map.get(&id).unwrap();
1313 cb.record_failure();
1314 }
1315
1316 let err = mgr.acquire_proxy().await.unwrap_err();
1317 assert!(
1318 matches!(err, ProxyError::AllProxiesUnhealthy),
1319 "expected AllProxiesUnhealthy, got {err:?}"
1320 );
1321 }
1322
1323 #[tokio::test]
1325 async fn handle_drop_records_failure() {
1326 let store = storage();
1327 let mgr = ProxyManager::with_round_robin(
1328 store.clone(),
1329 ProxyConfig {
1330 circuit_open_threshold: 1,
1331 ..ProxyConfig::default()
1332 },
1333 )
1334 .unwrap();
1335 let id = mgr
1336 .add_proxy(make_proxy("http://y.test:8080"))
1337 .await
1338 .unwrap();
1339
1340 {
1341 let _h = mgr.acquire_proxy().await.unwrap();
1342 }
1344
1345 let cb_map = mgr.circuit_breakers.read().await;
1346 let cb = cb_map.get(&id).unwrap();
1347 assert_eq!(cb.state(), STATE_OPEN);
1348 }
1349
1350 #[tokio::test]
1352 async fn handle_success_keeps_closed() {
1353 let store = storage();
1354 let mgr = ProxyManager::with_round_robin(store.clone(), ProxyConfig::default()).unwrap();
1355 let id = mgr
1356 .add_proxy(make_proxy("http://z.test:8080"))
1357 .await
1358 .unwrap();
1359
1360 let h = mgr.acquire_proxy().await.unwrap();
1361 h.mark_success();
1362 drop(h);
1363
1364 let cb_map = mgr.circuit_breakers.read().await;
1365 let cb = cb_map.get(&id).unwrap();
1366 assert_eq!(cb.state(), STATE_CLOSED);
1367 }
1368
1369 #[tokio::test]
1371 async fn start_and_graceful_shutdown() {
1372 let store = storage();
1373 let mgr = ProxyManager::with_round_robin(
1374 store,
1375 ProxyConfig {
1376 health_check_interval: Duration::from_hours(1),
1377 ..ProxyConfig::default()
1378 },
1379 )
1380 .unwrap();
1381 let (token, handle) = mgr.start();
1382 token.cancel();
1383 let result = tokio::time::timeout(Duration::from_secs(1), handle).await;
1384 assert!(result.is_ok(), "health checker task should exit within 1s");
1385 }
1386
1387 #[cfg(feature = "tls-profiled")]
1388 #[tokio::test]
1389 async fn builder_accepts_profiled_request_mode_preset() {
1390 let store = storage();
1391 let cfg = ProxyConfig {
1392 profiled_request_mode: Some(crate::types::ProfiledRequestMode::Preset),
1393 ..ProxyConfig::default()
1394 };
1395
1396 let result = ProxyManager::builder()
1397 .storage(store)
1398 .strategy(Arc::new(RoundRobinStrategy::default()))
1399 .config(cfg)
1400 .build();
1401
1402 assert!(
1403 result.is_ok(),
1404 "builder should accept profiled preset mode: {:?}",
1405 result.err()
1406 );
1407 }
1408
1409 #[cfg(feature = "tls-profiled")]
1410 #[tokio::test]
1411 async fn builder_rejects_profiled_request_mode_strict_all_for_chrome() {
1412 let store = storage();
1413 let cfg = ProxyConfig {
1414 profiled_request_mode: Some(crate::types::ProfiledRequestMode::StrictAll),
1415 ..ProxyConfig::default()
1416 };
1417
1418 let result = ProxyManager::builder()
1419 .storage(store)
1420 .strategy(Arc::new(RoundRobinStrategy::default()))
1421 .config(cfg)
1422 .build();
1423
1424 let Err(err) = result else {
1425 panic!("strict_all should fail for default Chrome baseline profile")
1426 };
1427
1428 assert!(
1429 matches!(err, ProxyError::ConfigError(_)),
1430 "expected ConfigError, got {err:?}"
1431 );
1432 }
1433
1434 fn sticky_config() -> ProxyConfig {
1437 use crate::session::StickyPolicy;
1438 ProxyConfig {
1439 sticky_policy: StickyPolicy::domain_default(),
1440 ..ProxyConfig::default()
1441 }
1442 }
1443
1444 #[tokio::test]
1446 async fn sticky_same_domain_returns_same_proxy() {
1447 let store = storage();
1448 let mgr = ProxyManager::with_round_robin(store, sticky_config()).unwrap();
1449 mgr.add_proxy(make_proxy("http://p1.test:8080"))
1450 .await
1451 .unwrap();
1452 mgr.add_proxy(make_proxy("http://p2.test:8080"))
1453 .await
1454 .unwrap();
1455
1456 let h1 = mgr.acquire_for_domain("example.com").await.unwrap();
1457 let url1 = h1.proxy_url.clone();
1458 h1.mark_success();
1459
1460 let h2 = mgr.acquire_for_domain("example.com").await.unwrap();
1461 let url2 = h2.proxy_url.clone();
1462 h2.mark_success();
1463
1464 assert_eq!(url1, url2, "same domain should return the same proxy");
1465 }
1466
1467 #[tokio::test]
1469 async fn sticky_different_domains_may_differ() {
1470 let store = storage();
1471 let mgr = ProxyManager::with_round_robin(store, sticky_config()).unwrap();
1472 mgr.add_proxy(make_proxy("http://pa.test:8080"))
1473 .await
1474 .unwrap();
1475 mgr.add_proxy(make_proxy("http://pb.test:8080"))
1476 .await
1477 .unwrap();
1478
1479 let ha = mgr.acquire_for_domain("a.com").await.unwrap();
1480 let url_a = ha.proxy_url.clone();
1481 ha.mark_success();
1482
1483 let hb = mgr.acquire_for_domain("b.com").await.unwrap();
1484 let url_b = hb.proxy_url.clone();
1485 hb.mark_success();
1486
1487 assert_ne!(
1489 url_a, url_b,
1490 "different domains should get different proxies"
1491 );
1492 }
1493
1494 #[tokio::test]
1497 async fn sticky_expired_session_re_acquires() {
1498 use crate::session::StickyPolicy;
1499 let store = storage();
1500 let mgr = ProxyManager::with_round_robin(
1501 store,
1502 ProxyConfig {
1503 sticky_policy: StickyPolicy::domain(Duration::from_millis(1)),
1504 ..ProxyConfig::default()
1505 },
1506 )
1507 .unwrap();
1508 mgr.add_proxy(make_proxy("http://x.test:8080"))
1509 .await
1510 .unwrap();
1511
1512 let h1 = mgr.acquire_for_domain("expired.com").await.unwrap();
1513 h1.mark_success();
1514
1515 tokio::time::sleep(Duration::from_millis(5)).await;
1517
1518 let h2 = mgr.acquire_for_domain("expired.com").await.unwrap();
1520 h2.mark_success();
1521 }
1522
1523 #[tokio::test]
1526 async fn sticky_cb_trip_invalidates_session() {
1527 let store = storage();
1528 let mgr = ProxyManager::with_round_robin(
1529 store,
1530 ProxyConfig {
1531 circuit_open_threshold: 1,
1532 sticky_policy: sticky_config().sticky_policy,
1533 ..ProxyConfig::default()
1534 },
1535 )
1536 .unwrap();
1537 mgr.add_proxy(make_proxy("http://q1.test:8080"))
1538 .await
1539 .unwrap();
1540 mgr.add_proxy(make_proxy("http://q2.test:8080"))
1541 .await
1542 .unwrap();
1543
1544 let h1 = mgr.acquire_for_domain("cb.com").await.unwrap();
1546 let url1 = h1.proxy_url.clone();
1547 drop(h1);
1549
1550 tokio::task::yield_now().await;
1552
1553 let _h2 = mgr.acquire_for_domain("cb.com").await;
1557 let _ = url1;
1559 }
1560
1561 #[tokio::test]
1563 async fn sticky_purge_expired() {
1564 use crate::session::StickyPolicy;
1565 let store = storage();
1566 let mgr = ProxyManager::with_round_robin(
1567 store,
1568 ProxyConfig {
1569 sticky_policy: StickyPolicy::domain(Duration::from_millis(1)),
1570 ..ProxyConfig::default()
1571 },
1572 )
1573 .unwrap();
1574 mgr.add_proxy(make_proxy("http://r.test:8080"))
1575 .await
1576 .unwrap();
1577
1578 let h = mgr.acquire_for_domain("purge.com").await.unwrap();
1579 h.mark_success();
1580
1581 assert_eq!(mgr.sessions.active_count(), 1);
1582
1583 tokio::time::sleep(Duration::from_millis(5)).await;
1585 let _ = mgr.sessions.purge_expired();
1586
1587 assert_eq!(mgr.sessions.active_count(), 0);
1588 }
1589
1590 #[tokio::test]
1592 async fn pool_stats_includes_sessions() {
1593 let store = storage();
1594 let mgr = ProxyManager::with_round_robin(store, sticky_config()).unwrap();
1595 mgr.add_proxy(make_proxy("http://s.test:8080"))
1596 .await
1597 .unwrap();
1598
1599 let stats = mgr.pool_stats().await.unwrap();
1600 assert_eq!(stats.active_sessions, 0);
1601
1602 let h = mgr.acquire_for_domain("stats.com").await.unwrap();
1603 h.mark_success();
1604
1605 let stats = mgr.pool_stats().await.unwrap();
1606 assert_eq!(stats.active_sessions, 1);
1607 }
1608
1609 #[cfg(feature = "bayesian-rotation")]
1616 #[tokio::test]
1617 async fn thompson_observer_records_outcomes_through_manager() {
1618 use crate::strategy::{BayesianObserver, ThompsonStrategy};
1619 use std::time::Duration;
1620
1621 let store = storage();
1622 let strategy = Arc::new(ThompsonStrategy::with_decay(Duration::from_hours(1), 0.99));
1623 let observer: BoxedBayesianObserver = Arc::clone(&strategy) as BoxedBayesianObserver;
1624
1625 let mgr = ProxyManager::builder()
1628 .storage(store.clone())
1629 .strategy(Arc::clone(&strategy) as BoxedRotationStrategy)
1630 .observer(observer)
1631 .config(ProxyConfig::default())
1632 .build()
1633 .unwrap();
1634 mgr.add_proxy(make_proxy("http://alpha.test:8080"))
1635 .await
1636 .unwrap();
1637 mgr.add_proxy(make_proxy("http://beta.test:8080"))
1638 .await
1639 .unwrap();
1640
1641 let records = store.list().await.unwrap();
1642 let mut by_url: std::collections::HashMap<String, Uuid> = std::collections::HashMap::new();
1643 for r in &records {
1644 by_url.insert(r.proxy.url.clone(), r.id);
1645 }
1646 let alpha_id = *by_url
1647 .get("http://alpha.test:8080")
1648 .expect("alpha proxy should be in storage");
1649 let beta_id = *by_url
1650 .get("http://beta.test:8080")
1651 .expect("beta proxy should be in storage");
1652
1653 for _ in 0..8 {
1655 strategy.observe(alpha_id, true);
1660 strategy.observe(beta_id, false);
1661 }
1662 let (alpha_succ, alpha_fail) = strategy.counts_for(alpha_id);
1663 let (beta_succ, beta_fail) = strategy.counts_for(beta_id);
1664 assert!(
1665 alpha_succ >= 8,
1666 "alpha should have many successes (got {alpha_succ})"
1667 );
1668 assert!(
1669 alpha_fail < 5,
1670 "alpha should have few failures (got {alpha_fail})"
1671 );
1672 assert!(
1673 beta_succ < 5,
1674 "beta should have few successes (got {beta_succ})"
1675 );
1676 assert!(
1677 beta_fail >= 8,
1678 "beta should have many failures (got {beta_fail})"
1679 );
1680 }
1681
1682 #[cfg(feature = "bayesian-rotation")]
1687 #[tokio::test]
1688 async fn thompson_manager_hot_path_budget() {
1689 use std::time::Duration;
1690 let store = storage();
1691 let mgr = ProxyManager::with_thompson_sampling(
1692 store.clone(),
1693 ProxyConfig::default(),
1694 Duration::from_hours(1),
1695 )
1696 .unwrap();
1697 mgr.add_proxy(make_proxy("http://p1.test:8080"))
1698 .await
1699 .unwrap();
1700 mgr.add_proxy(make_proxy("http://p2.test:8080"))
1701 .await
1702 .unwrap();
1703 mgr.add_proxy(make_proxy("http://p3.test:8080"))
1704 .await
1705 .unwrap();
1706
1707 for _ in 0..10 {
1709 let h = mgr.acquire_proxy().await.unwrap();
1710 h.mark_success();
1711 }
1712
1713 let start = std::time::Instant::now();
1714 for _ in 0..1_000 {
1715 let h = mgr.acquire_proxy().await.unwrap();
1716 h.mark_success();
1717 }
1718 let elapsed = start.elapsed();
1719 assert!(
1720 elapsed < std::time::Duration::from_secs(1),
1721 "1 000 Thompson manager round-trips took {elapsed:?}; hot-path budget violated"
1722 );
1723 }
1724
1725 #[cfg(feature = "bayesian-rotation")]
1731 #[tokio::test]
1732 async fn thompson_outperforms_round_robin_on_poisoned_pool() {
1733 use std::time::Duration;
1734
1735 let store_rr = storage();
1739 let store_th = storage();
1740 let mgr_rr = ProxyManager::with_round_robin(store_rr, ProxyConfig::default()).unwrap();
1741 let mgr_th = ProxyManager::with_thompson_sampling(
1742 store_th,
1743 ProxyConfig::default(),
1744 Duration::from_hours(1),
1745 )
1746 .unwrap();
1747
1748 let mut alive_urls: Vec<String> = Vec::new();
1750 let mut dead_urls: Vec<String> = Vec::new();
1751 for i in 0..5 {
1752 let url = format!("http://alive{i}.test:8080");
1753 mgr_rr.add_proxy(make_proxy(&url)).await.unwrap();
1754 mgr_th.add_proxy(make_proxy(&url)).await.unwrap();
1755 alive_urls.push(url);
1756 }
1757 for i in 0..5 {
1758 let url = format!("http://dead{i}.test:8080");
1759 mgr_rr.add_proxy(make_proxy(&url)).await.unwrap();
1760 mgr_th.add_proxy(make_proxy(&url)).await.unwrap();
1761 dead_urls.push(url);
1762 }
1763
1764 let records = mgr_th.storage().list().await.unwrap();
1768 for r in &records {
1769 if alive_urls.iter().any(|u| u == &r.proxy.url) {
1770 for _ in 0..5 {
1771 mgr_th.strategy_warmup_observe(r.id, true);
1772 }
1773 } else if dead_urls.iter().any(|u| u == &r.proxy.url) {
1774 for _ in 0..5 {
1775 mgr_th.strategy_warmup_observe(r.id, false);
1776 }
1777 }
1778 }
1779
1780 let mut rr_alive = 0_u64;
1785 let mut rr_dead = 0_u64;
1786 for _ in 0..200 {
1787 let h = mgr_rr.acquire_proxy().await.unwrap();
1788 let url = h.proxy_url.clone();
1789 if url.contains("alive") {
1790 h.mark_success();
1791 rr_alive += 1;
1792 } else {
1793 drop(h);
1794 rr_dead += 1;
1795 }
1796 }
1797 #[allow(clippy::cast_precision_loss)]
1801 let rr_dead_share = (rr_dead as f64) / ((rr_alive + rr_dead) as f64);
1802
1803 let mut th_alive = 0_u64;
1807 let mut th_dead = 0_u64;
1808 for _ in 0..200 {
1809 let h = mgr_th.acquire_proxy().await.unwrap();
1810 let url = h.proxy_url.clone();
1811 if url.contains("alive") {
1812 h.mark_success();
1813 th_alive += 1;
1814 } else {
1815 drop(h);
1816 th_dead += 1;
1817 }
1818 }
1819 #[allow(clippy::cast_precision_loss)]
1820 let th_dead_share = (th_dead as f64) / ((th_alive + th_dead) as f64);
1821
1822 assert!(
1826 th_dead_share < rr_dead_share,
1827 "Thompson dead-share ({th_dead_share:.3}) should be less than round-robin ({rr_dead_share:.3})"
1828 );
1829 let improvement = (rr_dead_share - th_dead_share) / rr_dead_share;
1831 assert!(
1832 improvement > 0.50,
1833 "expected >50% relative improvement in dead-share reduction (got {:.1}%)",
1834 improvement * 100.0
1835 );
1836 }
1837
1838 #[cfg(feature = "coherence-validation")]
1844 fn clean_us_context() -> crate::ports::coherence::CoherenceContext {
1845 use crate::ports::coherence::{AcceptLanguage, CoherenceContext, IsoCountry, Locale, Tz};
1846 use std::net::IpAddr;
1847 use std::str::FromStr;
1848 CoherenceContext {
1849 proxy_geo_country: Some(IsoCountry::new("US").unwrap()),
1850 dns_resolver_country: Some(IsoCountry::new("US").unwrap()),
1851 browser_locale: Locale::new("en-US").unwrap(),
1852 browser_timezone: Tz::new("America/New_York").unwrap(),
1853 accept_language: AcceptLanguage::new("en-US,en;q=0.9").unwrap(),
1854 webrtc_local_ip: None,
1855 webrtc_public_ip: Some(IpAddr::from_str("192.0.2.42").unwrap()),
1856 proxy_ip: Some(IpAddr::from_str("192.0.2.7").unwrap()),
1857 }
1858 }
1859
1860 #[cfg(feature = "coherence-validation")]
1864 #[tokio::test]
1865 async fn acquire_with_coherence_coherent_returns_proxy() {
1866 let store = storage();
1867 let mgr = ProxyManager::with_round_robin(store, ProxyConfig::default()).unwrap();
1868 mgr.add_proxy(make_proxy("http://a.test:8080"))
1869 .await
1870 .unwrap();
1871
1872 let ctx = clean_us_context();
1873 let policy = crate::ports::coherence::CoherencePolicy::advisory();
1874 let handle = mgr
1875 .acquire_proxy_with_coherence(&ctx, &policy)
1876 .await
1877 .unwrap();
1878 assert_eq!(handle.proxy_url, "http://a.test:8080");
1879 handle.mark_success();
1880 }
1881
1882 #[cfg(feature = "coherence-validation")]
1885 #[tokio::test]
1886 async fn acquire_with_coherence_hard_fail_returns_error() {
1887 use crate::ports::coherence::{
1888 AcceptLanguage, CoherenceContext, IsoCountry, Locale, MismatchField, MismatchSeverity,
1889 Tz,
1890 };
1891 let store = storage();
1892 let mgr = ProxyManager::with_round_robin(store, ProxyConfig::default()).unwrap();
1893 mgr.add_proxy(make_proxy("http://a.test:8080"))
1894 .await
1895 .unwrap();
1896
1897 let ctx = CoherenceContext {
1899 dns_resolver_country: Some(IsoCountry::new("PK").unwrap()),
1900 ..clean_us_context()
1901 };
1902 let policy =
1903 crate::ports::coherence::CoherencePolicy::hard_fail_on(MismatchField::ProxyGeoVsDns);
1904 let err = mgr
1905 .acquire_proxy_with_coherence(&ctx, &policy)
1906 .await
1907 .unwrap_err();
1908 match err {
1909 crate::error::ProxyError::CoherenceMismatch { field, severity } => {
1910 assert_eq!(field, MismatchField::ProxyGeoVsDns);
1911 assert_eq!(severity, MismatchSeverity::Hard);
1912 }
1913 other => panic!("expected CoherenceMismatch, got {other:?}"),
1914 }
1915 let _ = Locale::new("en-US").unwrap();
1918 let _ = AcceptLanguage::new("en-US").unwrap();
1919 let _ = Tz::new("America/New_York").unwrap();
1920 }
1921
1922 #[cfg(feature = "coherence-validation")]
1926 #[tokio::test]
1927 async fn acquire_with_coherence_advisory_mismatch_logs_and_returns_proxy() {
1928 use crate::ports::coherence::{CoherenceContext, Tz};
1929 let store = storage();
1930 let mgr = ProxyManager::with_round_robin(store, ProxyConfig::default()).unwrap();
1931 mgr.add_proxy(make_proxy("http://a.test:8080"))
1932 .await
1933 .unwrap();
1934
1935 let ctx = CoherenceContext {
1936 browser_timezone: Tz::new("Europe/London").unwrap(),
1937 ..clean_us_context()
1938 };
1939 let policy = crate::ports::coherence::CoherencePolicy::advisory();
1941 let handle = mgr
1942 .acquire_proxy_with_coherence(&ctx, &policy)
1943 .await
1944 .unwrap();
1945 assert_eq!(handle.proxy_url, "http://a.test:8080");
1946 handle.mark_success();
1947 }
1948
1949 #[cfg(feature = "coherence-validation")]
1955 #[tokio::test]
1956 async fn acquire_with_coherence_custom_validator_is_wired() {
1957 use crate::ports::coherence::{
1958 BoxedCoherencePort, CoherenceContext, CoherencePort, CoherenceVerdict,
1959 };
1960
1961 #[derive(Debug)]
1962 struct AlwaysCoherent;
1963 impl CoherencePort for AlwaysCoherent {
1964 fn evaluate(&self, _: &CoherenceContext) -> CoherenceVerdict {
1965 CoherenceVerdict::Coherent
1966 }
1967 }
1968
1969 let store = storage();
1970 let mgr = ProxyManager::builder()
1971 .storage(store)
1972 .coherence_validator(std::sync::Arc::new(AlwaysCoherent) as BoxedCoherencePort)
1973 .build()
1974 .unwrap();
1975 mgr.add_proxy(make_proxy("http://a.test:8080"))
1976 .await
1977 .unwrap();
1978
1979 let ctx = clean_us_context();
1981 let policy = crate::ports::coherence::CoherencePolicy::advisory();
1982 let handle = mgr
1983 .acquire_proxy_with_coherence(&ctx, &policy)
1984 .await
1985 .unwrap();
1986 assert_eq!(handle.proxy_url, "http://a.test:8080");
1987 handle.mark_success();
1988 }
1989
1990 #[cfg(feature = "coherence-validation")]
1995 #[tokio::test]
1996 async fn acquire_with_coherence_hot_path_budget() {
1997 let store = storage();
1998 let mgr = ProxyManager::with_round_robin(store, ProxyConfig::default()).unwrap();
1999 mgr.add_proxy(make_proxy("http://a.test:8080"))
2000 .await
2001 .unwrap();
2002 mgr.add_proxy(make_proxy("http://b.test:8080"))
2003 .await
2004 .unwrap();
2005 mgr.add_proxy(make_proxy("http://c.test:8080"))
2006 .await
2007 .unwrap();
2008
2009 let ctx = clean_us_context();
2010 let policy = crate::ports::coherence::CoherencePolicy::advisory();
2011 let start = std::time::Instant::now();
2012 for _ in 0..1_000 {
2013 let h = mgr
2014 .acquire_proxy_with_coherence(&ctx, &policy)
2015 .await
2016 .unwrap();
2017 h.mark_success();
2018 }
2019 let elapsed = start.elapsed();
2020 assert!(
2021 elapsed < std::time::Duration::from_secs(1),
2022 "1000 coherence-gated acquisitions took {elapsed:?}; hot-path budget violated"
2023 );
2024 }
2025
2026 #[cfg(feature = "vendor-stickiness")]
2032 #[tokio::test]
2033 async fn acquire_with_vendor_akamai_is_sticky() {
2034 let store = storage();
2035 let mgr = ProxyManager::with_round_robin(store, ProxyConfig::default()).unwrap();
2036 mgr.add_proxy(make_proxy("http://p1.test:8080"))
2037 .await
2038 .unwrap();
2039 mgr.add_proxy(make_proxy("http://p2.test:8080"))
2040 .await
2041 .unwrap();
2042
2043 let h1 = mgr
2044 .acquire_for_domain_with_vendor("example.com", crate::types::VendorId::Akamai)
2045 .await
2046 .unwrap();
2047 let url1 = h1.proxy_url.clone();
2048 h1.mark_success();
2049
2050 let h2 = mgr
2051 .acquire_for_domain_with_vendor("example.com", crate::types::VendorId::Akamai)
2052 .await
2053 .unwrap();
2054 let url2 = h2.proxy_url.clone();
2055 h2.mark_success();
2056
2057 assert_eq!(
2058 url1, url2,
2059 "Akamai sticky policy should return the same proxy across calls"
2060 );
2061 }
2062
2063 #[cfg(feature = "vendor-stickiness")]
2066 #[tokio::test]
2067 async fn acquire_with_vendor_data_dome_is_fresh_per_request() {
2068 let store = storage();
2069 let mgr = ProxyManager::with_round_robin(store, ProxyConfig::default()).unwrap();
2070 mgr.add_proxy(make_proxy("http://p1.test:8080"))
2071 .await
2072 .unwrap();
2073 mgr.add_proxy(make_proxy("http://p2.test:8080"))
2074 .await
2075 .unwrap();
2076
2077 let h1 = mgr
2078 .acquire_for_domain_with_vendor("example.com", crate::types::VendorId::DataDome)
2079 .await
2080 .unwrap();
2081 let url1 = h1.proxy_url.clone();
2082 h1.mark_success();
2083
2084 let h2 = mgr
2085 .acquire_for_domain_with_vendor("example.com", crate::types::VendorId::DataDome)
2086 .await
2087 .unwrap();
2088 let url2 = h2.proxy_url.clone();
2089 h2.mark_success();
2090
2091 assert_ne!(
2092 url1, url2,
2093 "DataDome fresh-per-request policy should yield different proxies"
2094 );
2095 }
2096
2097 #[cfg(feature = "vendor-stickiness")]
2100 #[tokio::test]
2101 async fn acquire_with_vendor_perimeter_x_is_fresh_per_domain() {
2102 let store = storage();
2103 let mgr = ProxyManager::with_round_robin(store, ProxyConfig::default()).unwrap();
2104 mgr.add_proxy(make_proxy("http://p1.test:8080"))
2105 .await
2106 .unwrap();
2107 mgr.add_proxy(make_proxy("http://p2.test:8080"))
2108 .await
2109 .unwrap();
2110
2111 let h1 = mgr
2112 .acquire_for_domain_with_vendor("example.com", crate::types::VendorId::PerimeterX)
2113 .await
2114 .unwrap();
2115 let url1 = h1.proxy_url.clone();
2116 h1.mark_success();
2117
2118 let h2 = mgr
2119 .acquire_for_domain_with_vendor("example.com", crate::types::VendorId::PerimeterX)
2120 .await
2121 .unwrap();
2122 let url2 = h2.proxy_url.clone();
2123 h2.mark_success();
2124
2125 assert_ne!(
2126 url1, url2,
2127 "PerimeterX fresh-per-domain policy should yield different proxies"
2128 );
2129 }
2130
2131 #[cfg(feature = "vendor-stickiness")]
2134 #[tokio::test]
2135 async fn acquire_with_vendor_unknown_defaults_to_fresh() {
2136 let store = storage();
2137 let mgr = ProxyManager::with_round_robin(store, ProxyConfig::default()).unwrap();
2138 mgr.add_proxy(make_proxy("http://p1.test:8080"))
2139 .await
2140 .unwrap();
2141 mgr.add_proxy(make_proxy("http://p2.test:8080"))
2142 .await
2143 .unwrap();
2144
2145 let h1 = mgr
2146 .acquire_for_domain_with_vendor("example.com", crate::types::VendorId::Unknown)
2147 .await
2148 .unwrap();
2149 let url1 = h1.proxy_url.clone();
2150 h1.mark_success();
2151
2152 let h2 = mgr
2153 .acquire_for_domain_with_vendor("example.com", crate::types::VendorId::Unknown)
2154 .await
2155 .unwrap();
2156 let url2 = h2.proxy_url.clone();
2157 h2.mark_success();
2158
2159 assert_ne!(url1, url2, "Unknown vendor should default to fresh");
2160 }
2161
2162 #[cfg(feature = "vendor-stickiness")]
2166 #[tokio::test]
2167 async fn acquire_with_vendor_sticky_binding_reacquires_after_failure() {
2168 let store = storage();
2169 let mgr = ProxyManager::with_round_robin(
2170 store,
2171 ProxyConfig {
2172 circuit_open_threshold: 1,
2173 ..ProxyConfig::default()
2174 },
2175 )
2176 .unwrap();
2177 mgr.add_proxy(make_proxy("http://p1.test:8080"))
2178 .await
2179 .unwrap();
2180 mgr.add_proxy(make_proxy("http://p2.test:8080"))
2181 .await
2182 .unwrap();
2183
2184 let h1 = mgr
2187 .acquire_for_domain_with_vendor("stale.com", crate::types::VendorId::Akamai)
2188 .await
2189 .unwrap();
2190 drop(h1);
2191 tokio::task::yield_now().await;
2192
2193 let result = mgr
2197 .acquire_for_domain_with_vendor("stale.com", crate::types::VendorId::Akamai)
2198 .await;
2199 match result {
2200 Ok(_h) => {} Err(crate::error::ProxyError::AllProxiesUnhealthy) => {} Err(e) => panic!("unexpected error after stale binding: {e:?}"),
2203 }
2204 }
2205
2206 #[cfg(feature = "vendor-stickiness")]
2209 #[tokio::test]
2210 async fn builder_stickiness_map_override_replaces_builtins() {
2211 use crate::stickiness::StickinessPolicy;
2212
2213 let store = storage();
2214 let custom = crate::stickiness::VendorStickinessMap::new()
2217 .with_override(
2218 crate::types::VendorId::Akamai,
2219 StickinessPolicy::StickyForever,
2220 )
2221 .with_override(
2222 crate::types::VendorId::DataDome,
2223 StickinessPolicy::StickyForTtl {
2224 ttl: Duration::from_mins(1),
2225 },
2226 );
2227 let mgr = ProxyManager::builder()
2228 .storage(store)
2229 .config(ProxyConfig::default())
2230 .stickiness_map(custom)
2231 .build()
2232 .unwrap();
2233 mgr.add_proxy(make_proxy("http://p1.test:8080"))
2234 .await
2235 .unwrap();
2236 mgr.add_proxy(make_proxy("http://p2.test:8080"))
2237 .await
2238 .unwrap();
2239
2240 let h1 = mgr
2242 .acquire_for_domain_with_vendor("example.com", crate::types::VendorId::Akamai)
2243 .await
2244 .unwrap();
2245 let url1 = h1.proxy_url.clone();
2246 h1.mark_success();
2247 let h2 = mgr
2248 .acquire_for_domain_with_vendor("example.com", crate::types::VendorId::Akamai)
2249 .await
2250 .unwrap();
2251 let url2 = h2.proxy_url.clone();
2252 h2.mark_success();
2253 assert_eq!(
2254 url1, url2,
2255 "custom StickyForever should keep the same proxy"
2256 );
2257
2258 let h1 = mgr
2261 .acquire_for_domain_with_vendor("dd.com", crate::types::VendorId::DataDome)
2262 .await
2263 .unwrap();
2264 let url1 = h1.proxy_url.clone();
2265 h1.mark_success();
2266 let h2 = mgr
2267 .acquire_for_domain_with_vendor("dd.com", crate::types::VendorId::DataDome)
2268 .await
2269 .unwrap();
2270 let url2 = h2.proxy_url.clone();
2271 h2.mark_success();
2272 assert_eq!(url1, url2, "custom StickyForTtl should keep the same proxy");
2273
2274 let h1 = mgr
2276 .acquire_for_domain_with_vendor("hc.com", crate::types::VendorId::Hcaptcha)
2277 .await
2278 .unwrap();
2279 let url1 = h1.proxy_url.clone();
2280 h1.mark_success();
2281 let h2 = mgr
2282 .acquire_for_domain_with_vendor("hc.com", crate::types::VendorId::Hcaptcha)
2283 .await
2284 .unwrap();
2285 let url2 = h2.proxy_url.clone();
2286 h2.mark_success();
2287 assert_ne!(
2288 url1, url2,
2289 "Hcaptcha should default to fresh when no override is installed"
2290 );
2291 }
2292
2293 #[cfg(feature = "vendor-stickiness")]
2298 #[tokio::test]
2299 async fn acquire_with_vendor_hot_path_budget() {
2300 let store = storage();
2301 let mgr = ProxyManager::with_round_robin(store, ProxyConfig::default()).unwrap();
2302 mgr.add_proxy(make_proxy("http://a.test:8080"))
2303 .await
2304 .unwrap();
2305 mgr.add_proxy(make_proxy("http://b.test:8080"))
2306 .await
2307 .unwrap();
2308 mgr.add_proxy(make_proxy("http://c.test:8080"))
2309 .await
2310 .unwrap();
2311
2312 let start = std::time::Instant::now();
2313 for _ in 0..1_000 {
2314 let h = mgr
2315 .acquire_for_domain_with_vendor("example.com", crate::types::VendorId::Akamai)
2316 .await
2317 .unwrap();
2318 h.mark_success();
2319 }
2320 let elapsed = start.elapsed();
2321 assert!(
2322 elapsed < std::time::Duration::from_secs(1),
2323 "1000 per-vendor acquisitions took {elapsed:?}; hot-path budget violated"
2324 );
2325 }
2326
2327 #[tokio::test]
2332 async fn add_proxy_with_metadata_stores_geo_fields() -> crate::error::ProxyResult<()> {
2333 let store = storage();
2334 let mgr = ProxyManager::with_round_robin(store.clone(), ProxyConfig::default())?;
2335 mgr.add_proxy_with_metadata(
2336 "http://cf-sf.test:8080",
2337 Some(13_335),
2338 Some("San Francisco"),
2339 Some("94110"),
2340 )
2341 .await?;
2342 let records = store.list().await?;
2343 assert_eq!(records.len(), 1);
2344 let record = records.first().expect("one record");
2345 assert_eq!(record.proxy.capabilities.asn, Some(13_335));
2346 assert_eq!(
2347 record.proxy.capabilities.city.as_deref(),
2348 Some("San Francisco")
2349 );
2350 assert_eq!(
2351 record.proxy.capabilities.postal_code.as_deref(),
2352 Some("94110")
2353 );
2354 Ok(())
2355 }
2356
2357 #[tokio::test]
2360 async fn add_proxy_with_metadata_rejects_invalid_geo() {
2361 let store = storage();
2362 let mgr = ProxyManager::with_round_robin(store.clone(), ProxyConfig::default()).unwrap();
2363 let err = mgr
2364 .add_proxy_with_metadata(
2365 "http://cf-sf.test:8080",
2366 Some(0), Some("San Francisco"),
2368 Some("94110"),
2369 )
2370 .await
2371 .expect_err("asn=0 must be rejected");
2372 assert!(matches!(
2373 err,
2374 crate::error::ProxyError::InvalidGeoMetadata { ref field, .. } if field == "asn"
2375 ));
2376 let err = mgr
2377 .add_proxy_with_metadata(
2378 "http://cf-sf.test:8080",
2379 Some(13_335),
2380 Some(""), Some("94110"),
2382 )
2383 .await
2384 .expect_err("empty city must be rejected");
2385 assert!(matches!(
2386 err,
2387 crate::error::ProxyError::InvalidGeoMetadata { ref field, .. } if field == "city"
2388 ));
2389 }
2390
2391 #[tokio::test]
2394 async fn add_proxy_with_metadata_round_trips_capability_filter() -> crate::error::ProxyResult<()>
2395 {
2396 let store = storage();
2397 let mgr = ProxyManager::with_round_robin(store.clone(), ProxyConfig::default())?;
2398 mgr.add_proxy_with_metadata(
2399 "http://cf-sf.test:8080",
2400 Some(13_335),
2401 Some("San Francisco"),
2402 Some("94110"),
2403 )
2404 .await?;
2405 let req = crate::types::CapabilityRequirement {
2406 require_asn: Some(13_335),
2407 ..Default::default()
2408 };
2409 let handle = mgr.acquire_with_capabilities(&req).await?;
2410 assert!(
2411 handle.proxy_url.contains("cf-sf.test"),
2412 "got url: {}",
2413 handle.proxy_url
2414 );
2415 handle.mark_success();
2416 Ok(())
2417 }
2418}