mocra_proxy/
proxy_manager.rs

1use crate::error::{ProxyError, Result};
2use crate::proxy_pool::*;
3use std::time::Duration;
4
5pub struct ProxyManager {
6    pool: ProxyPool,
7}
8
9impl ProxyManager {
10    pub async fn from_proxy_config(proxy_config: &ProxyConfig) -> Result<Self> {
11        let pool = proxy_config.build_proxy_pool().await;
12        Ok(Self { pool })
13    }
14
15    pub async fn from_config(config_str: &str) -> Result<Self> {
16        let proxy_setting = ProxyConfig::load_from_toml(config_str)?;
17        let pool = proxy_setting.build_proxy_pool().await;
18        Ok(Self { pool })
19    }
20
21    pub fn new() -> Self {
22        let pool = ProxyPool::new(PoolConfig::default());
23        Self { pool }
24    }
25
26    pub fn with_config(config: PoolConfig) -> Self {
27        let pool = ProxyPool::new(config);
28        Self { pool }
29    }
30
31    pub async fn get_proxy(&self, provider_name: Option<&str>) -> Result<ProxyEnum> {
32        self.pool.get_proxy(provider_name).await
33    }
34    pub async fn get_tunnel(&self) -> Result<ProxyEnum> {
35        self.pool
36            .get_best_tunnel()
37            .await
38            .ok_or_else(|| ProxyError::ProxyNotFound)
39    }
40
41    pub async fn report_proxy_result(
42        &self,
43        proxy: &ProxyEnum,
44        success: bool,
45        response_time: Option<Duration>,
46    ) -> Result<()> {
47        // Call the proxy pool's report_proxy_result directly; it dispatches on the proxy
48        // type automatically.
49        self.pool
50            .report_proxy_result(proxy, success, response_time)
51            .await
52    }
53
54    pub async fn report_success(
55        &self,
56        proxy: &ProxyEnum,
57        response_time: Option<Duration>,
58    ) -> Result<()> {
59        self.report_proxy_result(proxy, true, response_time).await
60    }
61
62    pub async fn report_failure(&self, proxy: &ProxyEnum) -> Result<()> {
63        self.report_proxy_result(proxy, false, None).await
64    }
65
66    pub async fn get_status(&self) -> std::collections::HashMap<String, usize> {
67        self.pool.get_pool_status().await
68    }
69
70    pub async fn get_detailed_stats(&self) -> PoolStats {
71        self.pool.get_stats().await
72    }
73
74    pub async fn health_check(&self) -> Result<()> {
75        self.pool.health_check().await
76    }
77
78    pub async fn add_ip_provider(&mut self, provider: Box<dyn IpProxyLoader>) {
79        self.pool.add_ip_provider(provider).await;
80    }
81    pub async fn add_tunnel(&mut self, tunnel: Tunnel) {
82        self.pool.add_tunnel(tunnel).await;
83    }
84}
85
86impl Default for ProxyManager {
87    fn default() -> Self {
88        Self::new()
89    }
90}
91
92#[cfg(test)]
93mod tests {
94
95    use crate::{ProxyConfig, ProxyManager};
96    use tokio::fs;
97
98    #[tokio::test]
99    #[ignore = "requires proxy.toml"]
100    async fn test() {
101        let config =
102            ProxyConfig::load_from_toml(&fs::read_to_string("proxy.toml").await.unwrap()).unwrap();
103        let mut manager = ProxyManager::new();
104        if let Some(tunnel) = config.tunnel {
105            for t in tunnel {
106                manager.add_tunnel(t).await;
107            }
108        }
109        let proxy = manager.get_tunnel().await.unwrap();
110        println!("{:?}", proxy.to_string());
111    }
112}
mocra_proxy/proxy_manager.rs

mocra_proxy/
proxy_manager.rs
