Skip to main content

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