Skip to main content

mocra_proxy/
proxy_pool.rs

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,     // 窗口开始(单调时钟)
18    window_duration: Duration, // 窗口长度
19}
20
21impl RateLimitTracker {
22    /// 创建新的限速跟踪器(毫秒级)
23    pub fn new() -> Self {
24        Self {
25            requests_in_window: 0,
26            window_start: Instant::now(),
27            window_duration: Duration::from_millis(1000), // 固定 1 秒窗口
28        }
29    }
30
31    /// 记录一个请求
32    pub fn record_request(&mut self) {
33        // 如果时间窗口已过期,重置计数器
34        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    /// 检查是否达到限速(毫秒级)
42    pub fn is_rate_limited(&mut self, rate_limit: f32) -> bool {
43        // 非正限速视为不限制
44        if rate_limit <= 0.0 {
45            return false;
46        }
47        // 若窗口已过,视为未限
48        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    /// 获取当前请求频率(毫秒级)
58    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    /// 剩余窗口时长(用于等待)
67    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    /// 当前窗口内已计数的请求数(若窗口已过期则返回0)
77    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>, // 提供商的过期时间
100    pub proxy_expire_time: u64,               // 当前供应商的代理过期时间
101    pub weight: Option<u32>,                  // 添加权重支持
102}
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>, // http, socks5, etc
111    pub rate_limit: f32,            // 每秒最大请求数
112}
113
114impl Display for IpProxy {
115    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
116        // 格式化为reqwest可以直接使用的代理URL格式
117        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, // 每秒最大请求数
144}
145
146impl Display for Tunnel {
147    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
148        // 格式化为reqwest可以直接使用的代理URL格式
149        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 // 比较浮点数时使用容差
202        } 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            // websocket 代理统一归一到 http/https 代理协议
263            "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, // 当池大小低于这个比例时触发补充
353}
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
377/// 代理池构建器
378pub 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>,      // 响应时间
450    pub success_rate: f32,                    // 成功率
451    pub rate_limit_tracker: RateLimitTracker, // 限速跟踪器
452    pub provider_rate_limit: f32,             // 提供商设置的限速值
453}
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; // 默认每秒10个请求
470        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, // 使用隧道的限速值
482        }
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(); // 默认5分钟
489        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, // 使用提供商的限速值
500        }
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        // 失败同样计入本窗口请求数,保持与选择时的一致性
557        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        // 响应时间影响分数
570        if let Some(response_time) = self.response_time {
571            let response_ms = response_time.as_millis() as f32;
572            score -= response_ms / 100.0; // 响应时间越长,分数越低
573        }
574
575        // 最近使用时间影响分数
576        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; // 越久未使用,分数越低
585
586        score.max(0.0)
587    }
588
589    pub fn is_rate_limited(&mut self) -> bool {
590        // 确定使用的限速值
591        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    /// 获取代理,支持负载均衡和故障转移
726    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        // 首先尝试获取最佳隧道代理
731        if let Some(tunnel) = self.get_best_tunnel().await {
732            return Ok(tunnel);
733        }
734
735        // 如果没有可用的隧道代理,处理IP代理
736        self.get_best_ip_proxy().await
737    }
738
739    /// 获取质量最好的隧道代理
740    pub async fn get_best_tunnel(&self) -> Option<ProxyEnum> {
741        // 获取写锁,直接操作pools中的数据
742        let mut pools = self.pools.write().await;
743
744        // 收集所有隧道代理并按质量分数排序
745        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        // 按质量分数排序(降序)
759        // tunnel_items.sort_by(|a, b| {
760        //     b.quality_score()
761        //         .partial_cmp(&a.quality_score())
762        //         .unwrap_or(std::cmp::Ordering::Equal)
763        // });
764
765        // 尝试找到未达到限速的最佳隧道代理
766
767        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        // 全部被限,计算最短剩余时间等待
775        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            // 重试一次(不递归多次,交给调用方后续节奏)
787            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    /// 从指定提供商获取代理
799    async fn get_ip_proxy_from_provider(&self, provider_name: &str) -> Result<ProxyEnum> {
800        self.ensure_pool_size(provider_name).await?;
801
802        // 先获取代理池的副本,避免锁跨越await
803        {
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            // 移除过期和无效的代理
810            pool.retain(|item| item.is_valid(self.config.max_errors));
811
812            // 按质量分数排序
813            pool.sort();
814            // 尝试找到未达到限速的最佳代理
815            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        // 如果所有代理都达到限速,尝试获取新的代理
824        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            // 移除过期和无效的代理
833            pool.retain(|item| item.is_valid(self.config.max_errors));
834
835            // 按质量分数排序
836            pool.sort();
837            // 尝试找到未达到限速的最佳代理
838            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            // 如果仍然没有可用代理,计算最短剩余时间等待后重试一次
845            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    /// 获取最佳代理(跨所有提供商)
869    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        // 如果所有代理都达到限速,获取权重最高的IpProvider 执行self.get_ip_proxy_from_provider
885
886        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)); // 按权重降序排序
894        if let Some((provider_name, _)) = providers.first() {
895            // 计算全局最短剩余时间并等待一次
896            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            // 再尝试一次是否已解除限速(需要可变借用)
911            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    /// 报告代理使用结果
924    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    /// 报告隧道代理使用结果
943    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        } // 写锁作用域提前结束
969        if !found {
970            return Err(ProxyError::InvalidConfig(
971                format!("Tunnel {} not found", tunnel.endpoint).into(),
972            ));
973        }
974        // 更新统计信息
975        self.update_stats().await;
976        Ok(())
977    }
978
979    /// 报告IP代理使用结果
980    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            // 遍历所有提供商的池查找代理
989            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            // 清理所有池中的无效代理
1010            for pool in pools.values_mut() {
1011                pool.retain(|item| item.is_valid(self.config.max_errors));
1012            }
1013        } // 写锁作用域提前结束
1014        // 如果没有找到代理,返回错误
1015        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        // 更新统计信息
1025        self.update_stats().await;
1026        Ok(())
1027    }
1028
1029    /// 报告代理成功使用
1030    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    /// 报告代理失败使用
1039    pub async fn report_failure(&self, proxy: &ProxyEnum) -> Result<()> {
1040        self.report_proxy_result(proxy, false, None).await
1041    }
1042
1043    /// 确保池大小满足要求,具体补充多少由对应的struct决定
1044    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    /// 当所有代理IP超限了,获取新一批的代理并添加进去,之前的代理仍然存在
1060    async fn refill_pool(&self, provider_name: &str) -> Result<()> {
1061        // await 前 clone Arc 指针
1062        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        // check provider is available
1069        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    /// 更新统计信息
1097    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    /// 获取池状态
1163    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    /// 获取详细统计信息
1172    pub async fn get_stats(&self) -> PoolStats {
1173        self.stats.read().await.clone()
1174    }
1175
1176    /// 执行健康检查
1177    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                        // 执行健康检查
1196                        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}