crawlkit_extensions/cdp/
mod.rs

1//! CDP (Chrome DevTools Protocol) 客户端
2//!
3//! 通过 chromiumoxide 连接到已运行的 CDP 服务器（Lightpanda / Obscura / Chrome 等），
4//! 导航页面并提取 HTML。适用于需要 JavaScript 渲染的 SPA 站点。
5//!
6//! # 支持的 CDP 服务
7//!
8//! | 服务 | 启动方式 |
9//! |------|---------|
10//! | [Lightpanda](https://github.com/lightpanda-io/browser) | `lightpanda serve --port 9222` |
11//! | [Obscura](https://github.com/h4ckf0r0day/obscura) | `obscura serve --port 9222` |
12//! | Chrome/Chromium | `chrome --remote-debugging-port=9222` |
13//!
14//! # 示例
15//!
16//! ```rust,no_run
17//! use crawlkit_extensions::cdp::{CdpClient, CdpPool, CdpStrategy};
18//!
19//! # #[tokio::main]
20//! # async fn main() -> anyhow::Result<()> {
21//! // 单个客户端
22//! let client = CdpClient::builder()
23//!     .with_endpoint("http://127.0.0.1:9222")
24//!     .build()
25//!     .await?;
26//!
27//! // 多端点池（轮询策略）
28//! let pool = CdpPool::builder()
29//!     .with_endpoint("http://127.0.0.1:9222")
30//!     .with_endpoint("http://127.0.0.1:9223")
31//!     .with_strategy(CdpStrategy::RoundRobin)
32//!     .build()
33//!     .await?;
34//! # Ok(())
35//! # }
36//! ```
37
38use std::collections::HashMap;
39use std::sync::atomic::{AtomicUsize, Ordering};
40use std::sync::Arc;
41use std::time::Duration;
42
43use async_trait::async_trait;
44use chromiumoxide::browser::Browser;
45use crawlkit_core::{CrawlError, HttpClient, Response};
46use futures::StreamExt;
47use tokio::sync::RwLock;
48
49/// 端点选择策略
50#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
51pub enum CdpStrategy {
52    /// 随机选择一个可用端点
53    Random,
54    /// 轮询（Round-Robin），依次使用每个端点
55    #[default]
56    RoundRobin,
57    /// 故障转移（Failover），按优先级顺序，失败后切到下一个
58    Failover,
59}
60
61/// 单个 CDP 端点状态（池内部使用）
62struct CdpEndpoint {
63    client: CdpClient,
64    name: String,
65    healthy: bool,
66}
67
68// ============================================================
69// CdpClient（单端点）
70// ============================================================
71
72/// CDP 客户端构建器
73pub struct CdpClientBuilder {
74    endpoint: String,
75    name: Option<String>,
76    navigation_timeout: Duration,
77}
78
79impl Default for CdpClientBuilder {
80    fn default() -> Self {
81        Self {
82            endpoint: "http://127.0.0.1:9222".to_string(),
83            name: None,
84            navigation_timeout: Duration::from_secs(30),
85        }
86    }
87}
88
89impl CdpClientBuilder {
90    /// 设置 CDP 服务端点
91    ///
92    /// 支持 HTTP（`http://host:port`）或 WebSocket（`ws://host:port`）格式。
93    /// HTTP 端点会自动从 `/json/version` 获取 WebSocket URL。
94    pub fn with_endpoint(mut self, endpoint: impl Into<String>) -> Self {
95        self.endpoint = endpoint.into();
96        self
97    }
98
99    /// 设置端点名称（用于日志和调试）
100    pub fn with_name(mut self, name: impl Into<String>) -> Self {
101        self.name = Some(name.into());
102        self
103    }
104
105    /// 设置导航超时（默认 30 秒）
106    pub fn with_navigation_timeout(mut self, timeout: Duration) -> Self {
107        self.navigation_timeout = timeout;
108        self
109    }
110
111    /// 构建 CdpClient
112    ///
113    /// 连接到 CDP 服务器。如果服务器未运行，返回错误。
114    pub async fn build(self) -> anyhow::Result<CdpClient> {
115        let (browser, mut handler) = Browser::connect(&self.endpoint)
116            .await
117            .map_err(|e| anyhow::anyhow!("CDP 连接失败 ({}): {e}", self.endpoint))?;
118
119        tokio::spawn(async move {
120            while let Some(h) = handler.next().await {
121                if h.is_err() {
122                    break;
123                }
124            }
125        });
126
127        Ok(CdpClient {
128            browser,
129            navigation_timeout: self.navigation_timeout,
130            name: self.name.unwrap_or(self.endpoint),
131        })
132    }
133}
134
135/// CDP 客户端（单端点）
136///
137/// 通过 chromiumoxide 连接到 CDP 服务器，实现 `HttpClient` trait。
138pub struct CdpClient {
139    browser: Browser,
140    navigation_timeout: Duration,
141    name: String,
142}
143
144impl CdpClient {
145    /// 创建 Builder
146    pub fn builder() -> CdpClientBuilder {
147        CdpClientBuilder::default()
148    }
149
150    /// 获取内部 Browser 引用（用于高级操作）
151    pub fn browser(&self) -> &Browser {
152        &self.browser
153    }
154
155    /// 获取端点名称
156    pub fn endpoint_name(&self) -> &str {
157        &self.name
158    }
159}
160
161#[async_trait]
162impl HttpClient for CdpClient {
163    async fn get(
164        &self,
165        url: &str,
166        _headers: &HashMap<String, String>,
167    ) -> crawlkit_core::Result<Response> {
168        let page = self
169            .browser
170            .new_page(url)
171            .await
172            .map_err(|e| CrawlError::Http(format!("CDP 创建页面失败: {e}")))?;
173
174        match tokio::time::timeout(self.navigation_timeout, page.wait_for_navigation()).await {
175            Ok(Ok(_)) => {}
176            Ok(Err(e)) => {
177                eprintln!("[CDP:{}] 等待导航警告: {e}", self.name);
178            }
179            Err(_) => {
180                eprintln!(
181                    "[CDP:{}] 导航超时 ({:?})，继续提取",
182                    self.name, self.navigation_timeout
183                );
184            }
185        }
186
187        let html = page
188            .content()
189            .await
190            .map_err(|e| CrawlError::Http(format!("CDP 提取 HTML 失败: {e}")))?;
191
192        let final_url = page
193            .url()
194            .await
195            .unwrap_or(None)
196            .unwrap_or_else(|| url.to_string());
197
198        let _ = page.close().await;
199
200        Ok(Response {
201            url: final_url,
202            status: 200,
203            headers: Default::default(),
204            body: html,
205        })
206    }
207
208    async fn post(
209        &self,
210        url: &str,
211        headers: &HashMap<String, String>,
212        _body: Vec<u8>,
213    ) -> crawlkit_core::Result<Response> {
214        self.get(url, headers).await
215    }
216
217    fn name(&self) -> &str {
218        &self.name
219    }
220}
221
222// ============================================================
223// CdpPool（多端点池）
224// ============================================================
225
226/// CDP 连接池构建器
227pub struct CdpPoolBuilder {
228    endpoints: Vec<(String, Option<String>)>,
229    strategy: CdpStrategy,
230    navigation_timeout: Duration,
231}
232
233impl Default for CdpPoolBuilder {
234    fn default() -> Self {
235        Self {
236            endpoints: Vec::new(),
237            strategy: CdpStrategy::default(),
238            navigation_timeout: Duration::from_secs(30),
239        }
240    }
241}
242
243impl CdpPoolBuilder {
244    /// 添加一个 CDP 端点
245    pub fn with_endpoint(mut self, endpoint: impl Into<String>) -> Self {
246        self.endpoints.push((endpoint.into(), None));
247        self
248    }
249
250    /// 添加一个带名称的 CDP 端点
251    pub fn with_named_endpoint(
252        mut self,
253        name: impl Into<String>,
254        endpoint: impl Into<String>,
255    ) -> Self {
256        self.endpoints.push((endpoint.into(), Some(name.into())));
257        self
258    }
259
260    /// 设置端点选择策略（默认 RoundRobin）
261    pub fn with_strategy(mut self, strategy: CdpStrategy) -> Self {
262        self.strategy = strategy;
263        self
264    }
265
266    /// 设置导航超时（默认 30 秒，应用到所有端点）
267    pub fn with_navigation_timeout(mut self, timeout: Duration) -> Self {
268        self.navigation_timeout = timeout;
269        self
270    }
271
272    /// 构建 CdpPool
273    ///
274    /// 依次连接所有端点。至少需要一个端点连接成功，否则返回错误。
275    /// 连接失败的端点会被标记为不健康，后续请求会跳过。
276    pub async fn build(self) -> anyhow::Result<CdpPool> {
277        if self.endpoints.is_empty() {
278            anyhow::bail!("CdpPool 至少需要一个 CDP 端点");
279        }
280
281        let mut endpoints = Vec::new();
282        let mut last_error = None;
283
284        for (url, name) in &self.endpoints {
285            let ep_name = name.clone().unwrap_or_else(|| url.clone());
286            match CdpClient::builder()
287                .with_endpoint(url)
288                .with_name(&ep_name)
289                .with_navigation_timeout(self.navigation_timeout)
290                .build()
291                .await
292            {
293                Ok(client) => {
294                    endpoints.push(CdpEndpoint {
295                        client,
296                        name: ep_name,
297                        healthy: true,
298                    });
299                }
300                Err(e) => {
301                    eprintln!("[CdpPool] 端点 {ep_name} ({url}) 连接失败: {e}，标记为不健康");
302                    last_error = Some(e);
303                }
304            }
305        }
306
307        if endpoints.is_empty() {
308            anyhow::bail!(
309                "CdpPool 所有端点均连接失败，最后错误: {}",
310                last_error.map(|e| e.to_string()).unwrap_or_else(|| "未知".into())
311            );
312        }
313
314        Ok(CdpPool {
315            endpoints: Arc::new(RwLock::new(endpoints)),
316            strategy: self.strategy,
317            counter: AtomicUsize::new(0),
318        })
319    }
320}
321
322/// CDP 连接池
323///
324/// 管理多个 CDP 端点，支持随机、轮询、故障转移策略。
325/// 实现 `HttpClient` trait，可直接用于 `CompositeFetcher`。
326pub struct CdpPool {
327    endpoints: Arc<RwLock<Vec<CdpEndpoint>>>,
328    strategy: CdpStrategy,
329    counter: AtomicUsize,
330}
331
332impl CdpPool {
333    /// 创建 Builder
334    pub fn builder() -> CdpPoolBuilder {
335        CdpPoolBuilder::default()
336    }
337
338    /// 获取当前健康端点数量
339    pub async fn healthy_count(&self) -> usize {
340        self.endpoints
341            .read()
342            .await
343            .iter()
344            .filter(|e| e.healthy)
345            .count()
346    }
347
348    /// 获取端点总数
349    pub async fn total_count(&self) -> usize {
350        self.endpoints.read().await.len()
351    }
352
353    /// 选择下一个端点索引（返回时立即释放锁）
354    async fn select_index(&self) -> Option<usize> {
355        let snapshot: Vec<(usize, bool, String)> = {
356            let endpoints = self.endpoints.read().await;
357            endpoints
358                .iter()
359                .enumerate()
360                .map(|(i, e)| (i, e.healthy, e.name.clone()))
361                .collect()
362        };
363
364        let healthy: Vec<(usize, &str)> = snapshot
365            .iter()
366            .filter(|(_, healthy, _)| *healthy)
367            .map(|(i, _, name)| (*i, name.as_str()))
368            .collect();
369
370        if healthy.is_empty() {
371            // 所有端点不健康，重置并返回第一个
372            self.reset_all_health().await;
373            return Some(0);
374        }
375
376        match self.strategy {
377            CdpStrategy::Random => {
378                use rand::Rng;
379                let idx = rand::rng().random_range(0..healthy.len());
380                Some(healthy[idx].0)
381            }
382            CdpStrategy::RoundRobin => {
383                let counter = self.counter.fetch_add(1, Ordering::Relaxed);
384                let idx = counter % healthy.len();
385                Some(healthy[idx].0)
386            }
387            CdpStrategy::Failover => Some(healthy[0].0),
388        }
389    }
390
391    /// 标记端点为不健康
392    async fn mark_unhealthy(&self, index: usize) {
393        let mut endpoints = self.endpoints.write().await;
394        if let Some(ep) = endpoints.get_mut(index) {
395            eprintln!("[CdpPool] 端点 {} 标记为不健康", ep.name);
396            ep.healthy = false;
397        }
398    }
399
400    /// 重置所有端点为健康状态
401    async fn reset_all_health(&self) {
402        let mut endpoints = self.endpoints.write().await;
403        for ep in endpoints.iter_mut() {
404            ep.healthy = true;
405        }
406    }
407}
408
409#[async_trait]
410impl HttpClient for CdpPool {
411    async fn get(
412        &self,
413        url: &str,
414        headers: &HashMap<String, String>,
415    ) -> crawlkit_core::Result<Response> {
416        let total = self.total_count().await;
417        let mut last_error = None;
418
419        for _ in 0..total {
420            let idx = match self.select_index().await {
421                Some(i) => i,
422                None => break,
423            };
424
425            // 获取端点名称（快速操作，立即释放锁）
426            let name = {
427                let endpoints = self.endpoints.read().await;
428                endpoints[idx].name.clone()
429            };
430
431            // 发起请求（不持有锁）
432            let result = {
433                let endpoints = self.endpoints.read().await;
434                endpoints[idx].client.get(url, headers).await
435            };
436
437            match result {
438                Ok(resp) => return Ok(resp),
439                Err(e) => {
440                    eprintln!("[CdpPool] 端点 {name} 请求失败: {e}");
441                    self.mark_unhealthy(idx).await;
442                    last_error = Some(e);
443                }
444            }
445        }
446
447        Err(CrawlError::Http(format!(
448            "CdpPool: 所有 {total} 个端点均失败，最后错误: {}",
449            last_error.unwrap_or_else(|| CrawlError::Http("无可用端点".to_string()))
450        )))
451    }
452
453    async fn post(
454        &self,
455        url: &str,
456        headers: &HashMap<String, String>,
457        _body: Vec<u8>,
458    ) -> crawlkit_core::Result<Response> {
459        self.get(url, headers).await
460    }
461
462    fn name(&self) -> &str {
463        "cdp-pool"
464    }
465}
466
467#[cfg(test)]
468mod tests {
469    use super::*;
470
471    #[test]
472    fn test_cdp_strategy_default() {
473        assert_eq!(CdpStrategy::default(), CdpStrategy::RoundRobin);
474    }
475
476    #[test]
477    fn test_cdp_builder_defaults() {
478        let builder = CdpClientBuilder::default();
479        assert_eq!(builder.endpoint, "http://127.0.0.1:9222");
480        assert!(builder.name.is_none());
481        assert_eq!(builder.navigation_timeout, Duration::from_secs(30));
482    }
483
484    #[test]
485    fn test_cdp_pool_builder_defaults() {
486        let builder = CdpPoolBuilder::default();
487        assert!(builder.endpoints.is_empty());
488        assert_eq!(builder.strategy, CdpStrategy::RoundRobin);
489        assert_eq!(builder.navigation_timeout, Duration::from_secs(30));
490    }
491}
crawlkit_extensions/cdp/mod.rs

crawlkit_extensions/cdp/
mod.rs
