1use 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
51pub enum CdpStrategy {
52 Random,
54 RoundRobin,
56 Failover,
58}
59
60impl Default for CdpStrategy {
61 fn default() -> Self {
62 Self::RoundRobin
63 }
64}
65
66struct CdpEndpoint {
68 client: CdpClient,
69 name: String,
70 healthy: bool,
71}
72
73pub 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 pub fn with_endpoint(mut self, endpoint: impl Into<String>) -> Self {
100 self.endpoint = endpoint.into();
101 self
102 }
103
104 pub fn with_name(mut self, name: impl Into<String>) -> Self {
106 self.name = Some(name.into());
107 self
108 }
109
110 pub fn with_navigation_timeout(mut self, timeout: Duration) -> Self {
112 self.navigation_timeout = timeout;
113 self
114 }
115
116 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
140pub struct CdpClient {
144 browser: Browser,
145 navigation_timeout: Duration,
146 name: String,
147}
148
149impl CdpClient {
150 pub fn builder() -> CdpClientBuilder {
152 CdpClientBuilder::default()
153 }
154
155 pub fn browser(&self) -> &Browser {
157 &self.browser
158 }
159
160 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
227pub 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 pub fn with_endpoint(mut self, endpoint: impl Into<String>) -> Self {
251 self.endpoints.push((endpoint.into(), None));
252 self
253 }
254
255 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 pub fn with_strategy(mut self, strategy: CdpStrategy) -> Self {
267 self.strategy = strategy;
268 self
269 }
270
271 pub fn with_navigation_timeout(mut self, timeout: Duration) -> Self {
273 self.navigation_timeout = timeout;
274 self
275 }
276
277 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
327pub struct CdpPool {
332 endpoints: Arc<RwLock<Vec<CdpEndpoint>>>,
333 strategy: CdpStrategy,
334 counter: AtomicUsize,
335}
336
337impl CdpPool {
338 pub fn builder() -> CdpPoolBuilder {
340 CdpPoolBuilder::default()
341 }
342
343 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 pub async fn total_count(&self) -> usize {
355 self.endpoints.read().await.len()
356 }
357
358 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 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 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 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 let name = {
432 let endpoints = self.endpoints.read().await;
433 endpoints[idx].name.clone()
434 };
435
436 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}