1use crate::error::ProxyError;
2use crate::error::Result;
3use async_trait::async_trait;
4use serde::{Deserialize, Serialize};
5use std::cmp::{Ordering, PartialEq};
6use std::collections::HashMap;
7use std::fmt::{Display, Formatter};
8use std::sync::Arc;
9use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
10use time::OffsetDateTime;
11use time::format_description::well_known::Rfc3339;
12use tokio::sync::{Mutex, RwLock};
13use url::Url;
14#[derive(Clone)]
15pub struct RateLimitTracker {
16 requests_in_window: u32,
17 window_start: Instant, window_duration: Duration, }
20
21impl RateLimitTracker {
22 pub fn new() -> Self {
24 Self {
25 requests_in_window: 0,
26 window_start: Instant::now(),
27 window_duration: Duration::from_millis(1000), }
29 }
30
31 pub fn record_request(&mut self) {
33 if self.window_start.elapsed() >= self.window_duration {
35 self.requests_in_window = 0;
36 self.window_start = Instant::now();
37 }
38 self.requests_in_window += 1;
39 }
40
41 pub fn is_rate_limited(&mut self, rate_limit: f32) -> bool {
43 if rate_limit <= 0.0 {
45 return false;
46 }
47 if self.window_start.elapsed() >= self.window_duration {
49 self.requests_in_window = 0;
50 self.window_start = Instant::now();
51 return false;
52 }
53 let cap = rate_limit.floor() as u32;
54 self.requests_in_window >= cap
55 }
56
57 pub fn get_current_rate(&self) -> f32 {
59 let elapsed = self.window_start.elapsed();
60 if elapsed >= self.window_duration || elapsed.as_millis() == 0 {
61 return 0.0;
62 }
63 self.requests_in_window as f32 / elapsed.as_secs_f32()
64 }
65
66 pub fn remaining_in_window(&self) -> Duration {
68 let elapsed = self.window_start.elapsed();
69 if elapsed >= self.window_duration {
70 Duration::from_millis(0)
71 } else {
72 self.window_duration - elapsed
73 }
74 }
75
76 pub fn current_window_count(&self) -> u32 {
78 if self.window_start.elapsed() >= self.window_duration {
79 0
80 } else {
81 self.requests_in_window
82 }
83 }
84}
85
86impl Default for RateLimitTracker {
87 fn default() -> Self {
88 Self::new()
89 }
90}
91
92#[derive(Serialize, Deserialize, Debug, Clone)]
93pub struct IpProvider {
94 pub name: String,
95 pub url: String,
96 pub retry_codes: Vec<u16>,
97 pub timeout: u64,
98 pub rate_limit: f32,
99 pub provider_expire_time: Option<String>, pub proxy_expire_time: u64, pub weight: Option<u32>, }
103
104#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
105pub struct IpProxy {
106 pub ip: String,
107 pub port: u16,
108 pub username: Option<String>,
109 pub password: Option<String>,
110 pub proxy_type: Option<String>, pub rate_limit: f32, }
113
114impl Display for IpProxy {
115 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
116 let proxy_type = self.proxy_type.as_deref().unwrap_or("http");
118 match (&self.username, &self.password) {
119 (Some(username), Some(password)) => {
120 write!(
121 f,
122 "{}://{}:{}@{}:{}",
123 proxy_type, username, password, self.ip, self.port
124 )
125 }
126 (Some(username), None) => {
127 write!(f, "{}://{}@{}:{}", proxy_type, username, self.ip, self.port)
128 }
129 _ => {
130 write!(f, "{}://{}:{}", proxy_type, self.ip, self.port)
131 }
132 }
133 }
134}
135#[derive(Serialize, Deserialize, Debug, Clone)]
136pub struct Tunnel {
137 pub name: String,
138 pub endpoint: String,
139 pub username: Option<String>,
140 pub password: Option<String>,
141 pub tunnel_type: String,
142 pub expire_time: String,
143 pub rate_limit: f32, }
145
146impl Display for Tunnel {
147 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
148 match (&self.username, &self.password) {
150 (Some(username), Some(password)) => {
151 write!(
152 f,
153 "{}://{}:{}@{}",
154 self.tunnel_type, username, password, self.endpoint
155 )
156 }
157 (Some(username), None) => {
158 write!(f, "{}://{}@{}", self.tunnel_type, username, self.endpoint)
159 }
160 _ => {
161 write!(f, "{}://{}", self.tunnel_type, self.endpoint)
162 }
163 }
164 }
165}
166#[derive(Serialize, Deserialize, Debug, Clone)]
167pub enum ProxyEnum {
168 Tunnel(Tunnel),
169 IpProxy(IpProxy),
170}
171
172impl Display for ProxyEnum {
173 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
174 let str = match self {
175 ProxyEnum::Tunnel(tunnel) => tunnel.to_string(),
176 ProxyEnum::IpProxy(ip_proxy) => ip_proxy.to_string(),
177 };
178 write!(f, "{str}")
179 }
180}
181impl PartialEq<Tunnel> for ProxyEnum {
182 fn eq(&self, other: &Tunnel) -> bool {
183 if let ProxyEnum::Tunnel(tunnel) = self {
184 tunnel.endpoint == other.endpoint
185 && tunnel.username == other.username
186 && tunnel.password == other.password
187 && tunnel.tunnel_type == other.tunnel_type
188 } else {
189 false
190 }
191 }
192}
193impl PartialEq<IpProxy> for ProxyEnum {
194 fn eq(&self, other: &IpProxy) -> bool {
195 if let ProxyEnum::IpProxy(ip_proxy) = self {
196 ip_proxy.ip == other.ip
197 && ip_proxy.port == other.port
198 && ip_proxy.username == other.username
199 && ip_proxy.password == other.password
200 && ip_proxy.proxy_type == other.proxy_type
201 && (ip_proxy.rate_limit - other.rate_limit).abs() < f32::EPSILON } else {
203 false
204 }
205 }
206}
207impl PartialEq for ProxyEnum {
208 fn eq(&self, other: &ProxyEnum) -> bool {
209 match self {
210 ProxyEnum::Tunnel(tunnel) => {
211 if let ProxyEnum::Tunnel(other_tunnel) = other {
212 tunnel.endpoint == other_tunnel.endpoint
213 && tunnel.username == other_tunnel.username
214 && tunnel.password == other_tunnel.password
215 && tunnel.tunnel_type == other_tunnel.tunnel_type
216 } else {
217 false
218 }
219 }
220 ProxyEnum::IpProxy(ip_proxy) => {
221 if let ProxyEnum::IpProxy(other_ip_proxy) = other {
222 ip_proxy.ip == other_ip_proxy.ip
223 && ip_proxy.port == other_ip_proxy.port
224 && ip_proxy.password == other_ip_proxy.password
225 && ip_proxy.username == other_ip_proxy.username
226 && ip_proxy.proxy_type == other_ip_proxy.proxy_type
227 } else {
228 false
229 }
230 }
231 }
232 }
233}
234#[derive(Serialize, Deserialize, Debug, Clone)]
235pub struct ProxyConfig {
236 pub tunnel: Option<Vec<Tunnel>>,
237 pub direct: Option<Vec<DirectProxy>>,
238 pub ip_provider: Option<Vec<IpProvider>>,
239 pub pool_config: Option<PoolConfig>,
240}
241
242#[derive(Serialize, Deserialize, Debug, Clone)]
243pub struct DirectProxy {
244 pub name: Option<String>,
245 pub url: String,
246 pub rate_limit: Option<f32>,
247 pub expire_time: Option<String>,
248}
249
250impl DirectProxy {
251 fn to_static_ip_proxy(&self, index: usize) -> Result<StaticIpProxyEntry> {
252 let parsed = Url::parse(&self.url).map_err(|e| {
253 ProxyError::InvalidConfig(
254 format!("invalid direct proxy url '{}': {e}", self.url).into(),
255 )
256 })?;
257
258 let scheme = parsed.scheme().to_ascii_lowercase();
259 let proxy_type = match scheme.as_str() {
260 "http" => "http",
261 "https" => "https",
262 "ws" => "http",
264 "wss" => "https",
265 _ => {
266 return Err(ProxyError::InvalidConfig(
267 format!(
268 "unsupported direct proxy scheme '{}', expected http/https/ws/wss",
269 scheme
270 )
271 .into(),
272 ));
273 }
274 }
275 .to_string();
276
277 let host = parsed.host_str().ok_or_else(|| {
278 ProxyError::InvalidConfig(format!("direct proxy missing host: {}", self.url).into())
279 })?;
280 let port = parsed.port_or_known_default().ok_or_else(|| {
281 ProxyError::InvalidConfig(format!("direct proxy missing port: {}", self.url).into())
282 })?;
283
284 let username = if parsed.username().is_empty() {
285 None
286 } else {
287 Some(parsed.username().to_string())
288 };
289
290 Ok(StaticIpProxyEntry {
291 provider_name: self
292 .name
293 .clone()
294 .unwrap_or_else(|| format!("direct_{}", index)),
295 proxy: IpProxy {
296 ip: host.to_string(),
297 port,
298 username,
299 password: parsed.password().map(|x| x.to_string()),
300 proxy_type: Some(proxy_type),
301 rate_limit: self.rate_limit.unwrap_or(10.0),
302 },
303 rate_limit: self.rate_limit.unwrap_or(10.0),
304 })
305 }
306}
307
308#[derive(Debug, Clone)]
309struct StaticIpProxyEntry {
310 provider_name: String,
311 proxy: IpProxy,
312 rate_limit: f32,
313}
314
315impl StaticIpProxyEntry {
316 fn into_proxy_item(self) -> ProxyItem {
317 ProxyItem::new_for_static_ip_proxy(self.proxy, self.provider_name, self.rate_limit)
318 }
319}
320
321impl ProxyConfig {
322 pub fn load_from_toml(toml_str: &str) -> Result<Self> {
323 toml::from_str(toml_str).map_err(|e| ProxyError::InvalidConfig(e.to_string().into()))
324 }
325
326 pub async fn build_proxy_pool(&self) -> ProxyPool {
327 let config = self.pool_config.clone().unwrap_or_default();
328 let mut builder = ProxyPoolBuilder::new(config);
329
330 if let Some(tunnels) = &self.tunnel {
331 builder = builder.with_tunnels(tunnels.clone());
332 }
333
334 if let Some(direct) = &self.direct {
335 builder = builder.with_direct_proxies(direct.clone());
336 }
337
338 if let Some(providers) = &self.ip_provider {
339 builder = builder.with_ip_providers(providers.clone());
340 }
341
342 builder.build().await
343 }
344}
345
346#[derive(Serialize, Deserialize, Debug, Clone)]
347pub struct PoolConfig {
348 pub min_size: usize,
349 pub max_size: usize,
350 pub max_errors: u32,
351 pub health_check_interval_secs: u64,
352 pub refill_threshold: f32, }
354
355impl Default for PoolConfig {
356 fn default() -> Self {
357 Self {
358 min_size: 5,
359 max_size: 50,
360 max_errors: 3,
361 health_check_interval_secs: 300,
362 refill_threshold: 0.3,
363 }
364 }
365}
366
367#[async_trait]
368pub trait IpProxyLoader: Send + Sync {
369 async fn get_ip_proxies(&self) -> Result<Vec<IpProxy>>;
370 fn is_retry_code(&self, code: &u16) -> bool;
371 fn get_name(&self) -> String;
372 fn get_weight(&self) -> u32;
373 fn get_config(&self) -> &IpProvider;
374 async fn health_check(&self, proxy: &IpProxy) -> bool;
375}
376
377pub struct ProxyPoolBuilder {
379 config: PoolConfig,
380 tunnels: Vec<Tunnel>,
381 direct_proxies: Vec<DirectProxy>,
382 ip_providers: Vec<IpProvider>,
383}
384
385impl ProxyPoolBuilder {
386 pub fn new(config: PoolConfig) -> Self {
387 Self {
388 config,
389 tunnels: Vec::new(),
390 direct_proxies: Vec::new(),
391 ip_providers: Vec::new(),
392 }
393 }
394
395 pub fn with_tunnels(mut self, tunnels: Vec<Tunnel>) -> Self {
396 self.tunnels = tunnels;
397 self
398 }
399
400 pub fn with_tunnel(mut self, tunnel: Tunnel) -> Self {
401 self.tunnels.push(tunnel);
402 self
403 }
404
405 pub fn with_ip_providers(mut self, providers: Vec<IpProvider>) -> Self {
406 self.ip_providers = providers;
407 self
408 }
409
410 pub fn with_direct_proxies(mut self, proxies: Vec<DirectProxy>) -> Self {
411 self.direct_proxies = proxies;
412 self
413 }
414
415 pub async fn build(self) -> ProxyPool {
416 let pool = ProxyPool::new(self.config);
417 for tunnel in &self.tunnels {
418 pool.add_tunnel(tunnel.clone()).await;
419 }
420 for (idx, direct) in self.direct_proxies.into_iter().enumerate() {
421 match direct.to_static_ip_proxy(idx) {
422 Ok(entry) => {
423 pool.add_static_ip_proxy(entry.into_proxy_item()).await;
424 }
425 Err(e) => {
426 log::warn!(
427 "[ProxyConfig] skip invalid direct proxy '{}': {}",
428 direct.url,
429 e
430 );
431 }
432 }
433 }
434 for provider in &self.ip_providers {
435 let loader = crate::proxy_impl::build_ip_proxy_loader(provider.clone());
436 pool.add_ip_provider(loader).await;
437 }
438 pool
439 }
440}
441#[derive(Clone)]
442pub struct ProxyItem {
443 pub proxy: ProxyEnum,
444 pub error_count: u32,
445 pub success_count: u32,
446 pub last_used: Option<Duration>,
447 pub expire_time: Duration,
448 pub provider_name: String,
449 pub response_time: Option<Duration>, pub success_rate: f32, pub rate_limit_tracker: RateLimitTracker, pub provider_rate_limit: f32, }
454
455impl ProxyItem {
456 pub fn new_for_tunnel(tunnel: Tunnel) -> Self {
457 let expire_time = if let Ok(datetime) = OffsetDateTime::parse(&tunnel.expire_time, &Rfc3339)
458 {
459 (datetime
460 - OffsetDateTime::from_unix_timestamp(0).unwrap_or(OffsetDateTime::UNIX_EPOCH))
461 .unsigned_abs()
462 } else {
463 Duration::from_secs(360 * 24 * 60 * 60)
464 + SystemTime::now()
465 .duration_since(UNIX_EPOCH)
466 .unwrap_or_default()
467 };
468
469 let rate = tunnel.rate_limit; let name = tunnel.name.clone();
471 Self {
472 proxy: ProxyEnum::Tunnel(tunnel),
473 error_count: 0,
474 success_count: 0,
475 last_used: None,
476 expire_time,
477 provider_name: name,
478 response_time: None,
479 success_rate: 1.0,
480 rate_limit_tracker: RateLimitTracker::new(),
481 provider_rate_limit: rate, }
483 }
484 pub fn new_for_ip_proxy(ip_proxy: IpProxy, ip_provider: &IpProvider) -> Self {
485 let expire_time = Duration::from_secs(ip_provider.proxy_expire_time)
486 + SystemTime::now()
487 .duration_since(UNIX_EPOCH)
488 .unwrap_or_default(); Self {
490 proxy: ProxyEnum::IpProxy(ip_proxy),
491 error_count: 0,
492 success_count: 0,
493 last_used: None,
494 expire_time,
495 provider_name: ip_provider.name.clone(),
496 response_time: None,
497 success_rate: 1.0,
498 rate_limit_tracker: RateLimitTracker::new(),
499 provider_rate_limit: ip_provider.rate_limit, }
501 }
502
503 pub fn new_for_static_ip_proxy(
504 ip_proxy: IpProxy,
505 provider_name: String,
506 rate_limit: f32,
507 ) -> Self {
508 let expire_time = Duration::from_secs(360 * 24 * 60 * 60)
509 + SystemTime::now()
510 .duration_since(UNIX_EPOCH)
511 .unwrap_or_default();
512 Self {
513 proxy: ProxyEnum::IpProxy(ip_proxy),
514 error_count: 0,
515 success_count: 0,
516 last_used: None,
517 expire_time,
518 provider_name,
519 response_time: None,
520 success_rate: 1.0,
521 rate_limit_tracker: RateLimitTracker::new(),
522 provider_rate_limit: rate_limit,
523 }
524 }
525
526 pub fn is_expired(&self) -> bool {
527 let now = SystemTime::now()
528 .duration_since(UNIX_EPOCH)
529 .unwrap_or_default();
530 now > self.expire_time
531 }
532
533 pub fn is_valid(&self, max_errors: u32) -> bool {
534 self.error_count < max_errors && !self.is_expired()
535 }
536
537 pub fn record_success(&mut self, response_time: Duration) {
538 self.success_count += 1;
539 self.response_time = Some(response_time);
540 self.last_used = Some(
541 SystemTime::now()
542 .duration_since(UNIX_EPOCH)
543 .unwrap_or_default(),
544 );
545 self.update_success_rate();
546 }
547
548 pub fn record_error(&mut self) {
549 self.error_count += 1;
550 self.last_used = Some(
551 SystemTime::now()
552 .duration_since(UNIX_EPOCH)
553 .unwrap_or_default(),
554 );
555 self.update_success_rate();
556 self.rate_limit_tracker.record_request();
558 }
559 fn update_success_rate(&mut self) {
560 let total = self.success_count + self.error_count;
561 if total > 0 {
562 self.success_rate = self.success_count as f32 / total as f32;
563 }
564 }
565
566 pub fn quality_score(&self) -> f32 {
567 let mut score = self.success_rate * 100.0;
568
569 if let Some(response_time) = self.response_time {
571 let response_ms = response_time.as_millis() as f32;
572 score -= response_ms / 100.0; }
574
575 let now = SystemTime::now()
577 .duration_since(UNIX_EPOCH)
578 .unwrap_or_default();
579 let time_since_last_use = if let Some(last_used) = self.last_used {
580 (now - last_used).as_secs()
581 } else {
582 0
583 };
584 score -= time_since_last_use as f32 / 3600.0; score.max(0.0)
587 }
588
589 pub fn is_rate_limited(&mut self) -> bool {
590 let actual_rate_limit = match &self.proxy {
592 ProxyEnum::IpProxy(ip_proxy) => {
593 if ip_proxy.rate_limit > 0.0 {
594 ip_proxy.rate_limit
595 } else {
596 self.provider_rate_limit
597 }
598 }
599 ProxyEnum::Tunnel(tunnel) => tunnel.rate_limit,
600 };
601 let limited = self.rate_limit_tracker.is_rate_limited(actual_rate_limit);
602 if limited {
603 let remaining = self.rate_limit_tracker.remaining_in_window();
604 let count = self.rate_limit_tracker.current_window_count();
605 log::warn!(
606 "Proxy rate limited: provider={}, proxy={}, limit={:.2}/s, count_in_window={}, remaining={:?}",
607 self.provider_name,
608 self.proxy,
609 actual_rate_limit,
610 count,
611 remaining
612 );
613 }
614 limited
615 }
616}
617
618impl PartialEq for ProxyItem {
619 fn eq(&self, other: &Self) -> bool {
620 match &self.proxy {
621 ProxyEnum::IpProxy(ip_proxy) => {
622 if let ProxyEnum::IpProxy(other_ip_proxy) = &other.proxy {
623 ip_proxy.ip == other_ip_proxy.ip && ip_proxy.port == other_ip_proxy.port
624 } else {
625 false
626 }
627 }
628 ProxyEnum::Tunnel(tunnel) => {
629 if let ProxyEnum::Tunnel(other_tunnel) = &other.proxy {
630 tunnel.endpoint == other_tunnel.endpoint
631 } else {
632 false
633 }
634 }
635 }
636 }
637}
638
639impl Eq for ProxyItem {}
640
641impl PartialOrd for ProxyItem {
642 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
643 Some(self.cmp(other))
644 }
645}
646
647impl Ord for ProxyItem {
648 fn cmp(&self, other: &Self) -> Ordering {
649 other
650 .quality_score()
651 .partial_cmp(&self.quality_score())
652 .unwrap_or(Ordering::Equal)
653 }
654}
655
656#[derive(Debug, Clone)]
657pub struct PoolStats {
658 pub total_proxies: usize,
659 pub valid_proxies: usize,
660 pub error_proxies: usize,
661 pub expired_proxies: usize,
662 pub avg_success_rate: f32,
663 pub providers: HashMap<String, ProviderStats>,
664}
665
666#[derive(Debug, Clone)]
667pub struct ProviderStats {
668 pub name: String,
669 pub total_proxies: usize,
670 pub valid_proxies: usize,
671 pub avg_success_rate: f32,
672 pub avg_response_time: Option<Duration>,
673}
674
675type IpProvidersMap = HashMap<String, Arc<Box<dyn IpProxyLoader>>>;
676
677pub struct ProxyPool {
678 pub config: PoolConfig,
679 pub pools: Arc<RwLock<HashMap<String, Vec<ProxyItem>>>>,
680 pub ip_providers: Arc<Mutex<IpProvidersMap>>,
681 pub stats: Arc<RwLock<PoolStats>>,
682}
683
684impl ProxyPool {
685 pub fn new(config: PoolConfig) -> Self {
686 Self {
687 config,
688 pools: Arc::new(RwLock::new(HashMap::new())),
689 ip_providers: Arc::new(Mutex::new(HashMap::new())),
690 stats: Arc::new(RwLock::new(PoolStats {
691 total_proxies: 0,
692 valid_proxies: 0,
693 error_proxies: 0,
694 expired_proxies: 0,
695 avg_success_rate: 0.0,
696 providers: HashMap::new(),
697 })),
698 }
699 }
700 pub async fn add_tunnel(&self, tunnel: Tunnel) {
701 let proxy_name = tunnel.name.clone();
702 let proxy_item = ProxyItem::new_for_tunnel(tunnel);
703 let mut pools = self.pools.write().await;
704 pools
705 .entry(proxy_name)
706 .or_insert_with(Vec::new)
707 .push(proxy_item);
708 }
709 pub async fn add_ip_provider(&self, provider: Box<dyn IpProxyLoader>) {
710 let name = provider.get_name();
711 let mut ip_providers = self.ip_providers.lock().await;
712 ip_providers.insert(name.clone(), Arc::new(provider));
713 let mut pools = self.pools.write().await;
714 pools.insert(name.clone(), Vec::new());
715 }
716
717 pub async fn add_static_ip_proxy(&self, item: ProxyItem) {
718 let mut pools = self.pools.write().await;
719 pools
720 .entry(item.provider_name.clone())
721 .or_insert_with(Vec::new)
722 .push(item);
723 }
724
725 pub async fn get_proxy(&self, provider_name: Option<&str>) -> Result<ProxyEnum> {
727 if let Some(name) = provider_name {
728 return self.get_ip_proxy_from_provider(name).await;
729 }
730 if let Some(tunnel) = self.get_best_tunnel().await {
732 return Ok(tunnel);
733 }
734
735 self.get_best_ip_proxy().await
737 }
738
739 pub async fn get_best_tunnel(&self) -> Option<ProxyEnum> {
741 let mut pools = self.pools.write().await;
743
744 let mut tunnel_items = Vec::new();
746 for pool in pools.values_mut() {
747 for item in pool.iter_mut() {
748 if matches!(item.proxy, ProxyEnum::Tunnel(_)) {
749 tunnel_items.push(item);
750 }
751 }
752 }
753
754 if tunnel_items.is_empty() {
755 return None;
756 }
757 tunnel_items.sort();
758 for item in tunnel_items.iter_mut() {
768 if !item.is_rate_limited() {
769 item.rate_limit_tracker.record_request();
770 return Some(item.proxy.clone());
771 }
772 }
773
774 if let Some(min_remaining) = tunnel_items
776 .iter()
777 .map(|i| i.rate_limit_tracker.remaining_in_window())
778 .min()
779 {
780 let sleep_dur = if min_remaining > Duration::from_millis(0) {
781 min_remaining
782 } else {
783 Duration::from_millis(50)
784 };
785 tokio::time::sleep(sleep_dur).await;
786 for item in tunnel_items.iter_mut() {
788 if !item.is_rate_limited() {
789 item.rate_limit_tracker.record_request();
790 return Some(item.proxy.clone());
791 }
792 }
793 }
794
795 None
796 }
797
798 async fn get_ip_proxy_from_provider(&self, provider_name: &str) -> Result<ProxyEnum> {
800 self.ensure_pool_size(provider_name).await?;
801
802 {
804 let mut pools = self.pools.write().await;
805 let pool = pools.get_mut(provider_name).ok_or_else(|| {
806 ProxyError::InvalidConfig(format!("Provider {provider_name} not found").into())
807 })?;
808
809 pool.retain(|item| item.is_valid(self.config.max_errors));
811
812 pool.sort();
814 for item in pool.iter_mut() {
816 if !item.is_rate_limited() {
817 item.rate_limit_tracker.record_request();
818 return Ok(item.proxy.clone());
819 }
820 }
821 }
822
823 self.refill_pool(provider_name).await?;
825
826 {
827 let mut pools = self.pools.write().await;
828 let pool = pools.get_mut(provider_name).ok_or_else(|| {
829 ProxyError::InvalidConfig(format!("Provider {provider_name} not found").into())
830 })?;
831
832 pool.retain(|item| item.is_valid(self.config.max_errors));
834
835 pool.sort();
837 for item in pool.iter_mut() {
839 if !item.is_rate_limited() {
840 item.rate_limit_tracker.record_request();
841 return Ok(item.proxy.clone());
842 }
843 }
844 if let Some(min_remaining) = pool
846 .iter()
847 .map(|i| i.rate_limit_tracker.remaining_in_window())
848 .min()
849 {
850 let sleep_dur = if min_remaining > Duration::from_millis(0) {
851 min_remaining
852 } else {
853 Duration::from_millis(50)
854 };
855 tokio::time::sleep(sleep_dur).await;
856 for item in pool.iter_mut() {
857 if !item.is_rate_limited() {
858 item.rate_limit_tracker.record_request();
859 return Ok(item.proxy.clone());
860 }
861 }
862 }
863 }
864
865 Err(ProxyError::InvalidConfig("No valid proxy available".into()))
866 }
867
868 async fn get_best_ip_proxy(&self) -> Result<ProxyEnum> {
870 let mut pools = self.pools.write().await;
871 let mut proxy_items = pools
872 .iter_mut()
873 .flat_map(|(_, v)| v)
874 .filter(|x| matches!(x.proxy, ProxyEnum::IpProxy(_)))
875 .collect::<Vec<_>>();
876 proxy_items.sort();
877 for item in proxy_items.iter_mut() {
878 if !item.is_rate_limited() {
879 item.rate_limit_tracker.record_request();
880 return Ok(item.proxy.clone());
881 }
882 }
883
884 let mut providers: Vec<_> = {
887 let providers = self.ip_providers.lock().await;
888 providers
889 .iter()
890 .map(|(name, provider)| (name.clone(), provider.get_weight()))
891 .collect()
892 };
893 providers.sort_by_key(|x| std::cmp::Reverse(x.1)); if let Some((provider_name, _)) = providers.first() {
895 if let Some(min_remaining) = proxy_items
897 .iter()
898 .map(|i| i.rate_limit_tracker.remaining_in_window())
899 .min()
900 {
901 let sleep_dur = if min_remaining > Duration::from_millis(0) {
902 min_remaining
903 } else {
904 Duration::from_millis(50)
905 };
906 tokio::time::sleep(sleep_dur).await;
907 }
908 self.get_ip_proxy_from_provider(provider_name).await
909 } else if !proxy_items.is_empty() {
910 let mut_idx = 0usize;
912 if !proxy_items[mut_idx].is_rate_limited() {
913 proxy_items[mut_idx].rate_limit_tracker.record_request();
914 Ok(proxy_items[mut_idx].proxy.clone())
915 } else {
916 Err(ProxyError::InvalidConfig("No valid proxy available".into()))
917 }
918 } else {
919 Err(ProxyError::InvalidConfig("No valid proxy available".into()))
920 }
921 }
922
923 pub async fn report_proxy_result(
925 &self,
926 proxy: &ProxyEnum,
927 success: bool,
928 response_time: Option<Duration>,
929 ) -> Result<()> {
930 match proxy {
931 ProxyEnum::Tunnel(tunnel) => {
932 self.report_tunnel_result(tunnel, success, response_time)
933 .await
934 }
935 ProxyEnum::IpProxy(ip_proxy) => {
936 self.report_ip_proxy_result(ip_proxy, success, response_time)
937 .await
938 }
939 }
940 }
941
942 async fn report_tunnel_result(
944 &self,
945 tunnel: &Tunnel,
946 success: bool,
947 response_time: Option<Duration>,
948 ) -> Result<()> {
949 let mut found = false;
950 {
951 let mut pools = self.pools.write().await;
952 for pool in pools.values_mut() {
953 for item in pool.iter_mut() {
954 if let ProxyEnum::Tunnel(ref t) = item.proxy
955 && t.endpoint == tunnel.endpoint
956 {
957 if success {
958 item.record_success(
959 response_time.unwrap_or(Duration::from_millis(1000)),
960 );
961 } else {
962 item.record_error();
963 }
964 found = true;
965 }
966 }
967 }
968 } if !found {
970 return Err(ProxyError::InvalidConfig(
971 format!("Tunnel {} not found", tunnel.endpoint).into(),
972 ));
973 }
974 self.update_stats().await;
976 Ok(())
977 }
978
979 async fn report_ip_proxy_result(
981 &self,
982 proxy: &IpProxy,
983 success: bool,
984 response_time: Option<Duration>,
985 ) -> Result<()> {
986 let mut proxy_found = false;
987 {
988 let mut pools = self.pools.write().await;
990 for pool in pools.values_mut() {
991 for item in pool.iter_mut() {
992 if item.proxy.eq(proxy) {
993 if success {
994 item.record_success(
995 response_time.unwrap_or(Duration::from_millis(1000)),
996 );
997 } else {
998 item.record_error();
999 }
1000 proxy_found = true;
1001 break;
1002 }
1003 }
1004 if proxy_found {
1005 break;
1006 }
1007 }
1008
1009 for pool in pools.values_mut() {
1011 pool.retain(|item| item.is_valid(self.config.max_errors));
1012 }
1013 } if !proxy_found {
1016 return Err(ProxyError::InvalidConfig(
1017 format!(
1018 "Proxy {}:{} not found in any provider",
1019 proxy.ip, proxy.port
1020 )
1021 .into(),
1022 ));
1023 }
1024 self.update_stats().await;
1026 Ok(())
1027 }
1028
1029 pub async fn report_success(
1031 &self,
1032 proxy: &ProxyEnum,
1033 response_time: Option<Duration>,
1034 ) -> Result<()> {
1035 self.report_proxy_result(proxy, true, response_time).await
1036 }
1037
1038 pub async fn report_failure(&self, proxy: &ProxyEnum) -> Result<()> {
1040 self.report_proxy_result(proxy, false, None).await
1041 }
1042
1043 async fn ensure_pool_size(&self, provider_name: &str) -> Result<()> {
1045 let current_size = {
1046 let pools = self.pools.read().await;
1047 pools.get(provider_name).map(|p| p.len()).unwrap_or(0)
1048 };
1049
1050 let threshold = (self.config.max_size as f32 * self.config.refill_threshold) as usize;
1051
1052 if current_size < self.config.min_size || current_size < threshold {
1053 self.refill_pool(provider_name).await?;
1054 }
1055
1056 Ok(())
1057 }
1058
1059 async fn refill_pool(&self, provider_name: &str) -> Result<()> {
1061 let provider: Arc<Box<dyn IpProxyLoader>> = {
1063 let providers = self.ip_providers.lock().await;
1064 providers.get(provider_name).cloned().ok_or_else(|| {
1065 ProxyError::InvalidConfig(format!("Provider {provider_name} not found").into())
1066 })?
1067 };
1068 if let Some(expire_time) = &provider.get_config().provider_expire_time {
1070 let now = OffsetDateTime::now_utc().unix_timestamp();
1071 if let Ok(expire_timestamp) = OffsetDateTime::parse(expire_time, &Rfc3339) {
1072 if now >= expire_timestamp.unix_timestamp() {
1073 return Err(ProxyError::ProxyProviderExpired);
1074 }
1075 } else {
1076 return Err(ProxyError::InvalidConfig(
1077 format!(
1078 "Provider {} expire time is not available",
1079 provider.get_config().name
1080 )
1081 .into(),
1082 ));
1083 };
1084 }
1085 let new_proxies = provider.get_ip_proxies().await?;
1086 {
1087 let mut pools = self.pools.write().await;
1088 let pool = pools.get_mut(provider_name).unwrap();
1089 for proxy in new_proxies {
1090 pool.push(ProxyItem::new_for_ip_proxy(proxy, provider.get_config()));
1091 }
1092 }
1093 Ok(())
1094 }
1095
1096 async fn update_stats(&self) {
1098 let pools = self.pools.read().await;
1099 let mut stats = self.stats.write().await;
1100
1101 stats.total_proxies = 0;
1102 stats.valid_proxies = 0;
1103 stats.error_proxies = 0;
1104 stats.expired_proxies = 0;
1105 stats.providers.clear();
1106
1107 let mut total_success_rate = 0.0;
1108 let mut total_providers = 0;
1109
1110 for (provider_name, pool) in pools.iter() {
1111 let mut provider_stats = ProviderStats {
1112 name: provider_name.clone(),
1113 total_proxies: pool.len(),
1114 valid_proxies: 0,
1115 avg_success_rate: 0.0,
1116 avg_response_time: None,
1117 };
1118
1119 let mut provider_success_rate = 0.0;
1120 let mut response_times = Vec::new();
1121
1122 for item in pool.iter() {
1123 stats.total_proxies += 1;
1124
1125 if item.is_valid(self.config.max_errors) {
1126 stats.valid_proxies += 1;
1127 provider_stats.valid_proxies += 1;
1128 } else if item.is_expired() {
1129 stats.expired_proxies += 1;
1130 } else {
1131 stats.error_proxies += 1;
1132 }
1133
1134 provider_success_rate += item.success_rate;
1135 if let Some(response_time) = item.response_time {
1136 response_times.push(response_time);
1137 }
1138 }
1139
1140 if !pool.is_empty() {
1141 provider_stats.avg_success_rate = provider_success_rate / pool.len() as f32;
1142 total_success_rate += provider_stats.avg_success_rate;
1143 total_providers += 1;
1144 }
1145
1146 if !response_times.is_empty() {
1147 let avg_ms = response_times.iter().map(|d| d.as_millis()).sum::<u128>()
1148 / response_times.len() as u128;
1149 provider_stats.avg_response_time = Some(Duration::from_millis(avg_ms as u64));
1150 }
1151
1152 stats
1153 .providers
1154 .insert(provider_name.clone(), provider_stats);
1155 }
1156
1157 if total_providers > 0 {
1158 stats.avg_success_rate = total_success_rate / total_providers as f32;
1159 }
1160 }
1161
1162 pub async fn get_pool_status(&self) -> HashMap<String, usize> {
1164 let pools = self.pools.read().await;
1165 pools
1166 .iter()
1167 .map(|(name, pool)| (name.clone(), pool.len()))
1168 .collect()
1169 }
1170
1171 pub async fn get_stats(&self) -> PoolStats {
1173 self.stats.read().await.clone()
1174 }
1175
1176 pub async fn health_check(&self) -> Result<()> {
1178 let providers: Vec<String> = {
1179 let providers = self.ip_providers.lock().await;
1180 providers.keys().cloned().collect()
1181 };
1182 for provider_name in providers {
1183 let provider: Option<Arc<Box<dyn IpProxyLoader>>> = {
1184 let providers = self.ip_providers.lock().await;
1185 providers.get(&provider_name).cloned()
1186 };
1187 if let Some(provider) = provider {
1188 let pool = {
1189 let pools = self.pools.read().await;
1190 pools.get(&provider_name).unwrap_or(&Vec::new()).clone()
1191 };
1192 let mut healthy_proxies = Vec::new();
1193 for item in pool.into_iter() {
1194 if let ProxyEnum::IpProxy(ref proxy) = item.proxy {
1195 if provider.health_check(proxy).await {
1197 healthy_proxies.push(item.clone());
1198 }
1199 }
1200 }
1201 {
1202 let mut pools = self.pools.write().await;
1203 if let Some(pool) = pools.get_mut(&provider_name) {
1204 *pool = healthy_proxies;
1205 }
1206 }
1207 }
1208 }
1209 self.update_stats().await;
1210 Ok(())
1211 }
1212}