1use std::pin::Pin;
2use std::sync::atomic::{AtomicUsize, Ordering};
3use std::sync::Arc;
4use std::time::{Duration, Instant};
5
6use futures::stream::Stream;
7use futures::StreamExt;
8use reqwest::header::USER_AGENT;
9use reqwest::{Client, StatusCode};
10use serde::{Deserialize, Serialize};
11use tokio::time::sleep;
12use url::Url;
13
14use crate::{CrawlConfig, CrawlError, FetchResult, RedirectHop};
15
16#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct RetryPolicy {
35 pub max_retries: usize,
37 #[serde(with = "crate::duration_ms")]
39 pub initial_backoff: Duration,
40 #[serde(with = "crate::duration_ms")]
42 pub max_backoff: Duration,
43 pub backoff_multiplier: f64,
45 pub retryable_statuses: Vec<u16>,
47}
48
49impl Default for RetryPolicy {
50 fn default() -> Self {
51 Self {
52 max_retries: 3,
53 initial_backoff: Duration::from_secs(1),
54 max_backoff: Duration::from_secs(30),
55 backoff_multiplier: 2.0,
56 retryable_statuses: vec![429, 500, 502, 503, 504],
57 }
58 }
59}
60
61impl RetryPolicy {
62 pub fn backoff_duration(&self, attempt: usize) -> Duration {
79 let base = self.initial_backoff.as_secs_f64();
80 let backoff = base * self.backoff_multiplier.powi(attempt as i32);
81 let capped = backoff.min(self.max_backoff.as_secs_f64());
82 Duration::from_secs_f64(capped)
83 }
84
85 pub fn is_retryable(&self, status: u16) -> bool {
89 self.retryable_statuses.contains(&status)
90 }
91}
92
93#[derive(Debug)]
112pub struct UserAgentRotator {
113 agents: Vec<String>,
114 index: AtomicUsize,
115}
116
117impl UserAgentRotator {
118 pub fn new(agents: Vec<String>) -> Self {
124 assert!(
125 !agents.is_empty(),
126 "UserAgentRotator requires at least one user-agent"
127 );
128 Self {
129 agents,
130 index: AtomicUsize::new(0),
131 }
132 }
133
134 pub fn next(&self) -> &str {
139 let idx = self.index.fetch_add(1, Ordering::AcqRel);
140 &self.agents[idx % self.agents.len()]
141 }
142
143 pub fn len(&self) -> usize {
145 self.agents.len()
146 }
147
148 pub fn is_empty(&self) -> bool {
150 self.agents.is_empty()
151 }
152}
153
154impl Default for UserAgentRotator {
155 fn default() -> Self {
156 Self::new(vec![format!("crawlkit/{}", env!("CARGO_PKG_VERSION"))])
157 }
158}
159
160#[derive(Debug, Clone)]
174pub struct HttpClientConfig {
175 pub timeout: Duration,
177 pub max_redirects: usize,
179 pub retry_policy: RetryPolicy,
181 pub user_agent: Arc<UserAgentRotator>,
183 pub max_body_size: usize,
185 pub pool_max_idle_per_host: usize,
187 pub pool_max_idle: usize,
189 pub tcp_keepalive: Option<Duration>,
191}
192
193impl From<&CrawlConfig> for HttpClientConfig {
194 fn from(config: &CrawlConfig) -> Self {
195 Self {
196 timeout: config.request_timeout,
197 max_redirects: config.max_redirects,
198 retry_policy: RetryPolicy::default(),
199 user_agent: Arc::new(UserAgentRotator::new(vec![config.user_agent.clone()])),
200 max_body_size: 10 * 1024 * 1024, pool_max_idle_per_host: 16,
202 pool_max_idle: 32,
203 tcp_keepalive: Some(Duration::from_secs(60)),
204 }
205 }
206}
207
208pub struct HttpClient {
232 client: Client,
233 config: HttpClientConfig,
234}
235
236impl HttpClient {
237 pub fn new(config: HttpClientConfig) -> Result<Self, CrawlError> {
247 let mut builder = Client::builder()
248 .timeout(config.timeout)
249 .redirect(reqwest::redirect::Policy::limited(config.max_redirects))
250 .user_agent(config.user_agent.next())
251 .https_only(true)
252 .http1_only()
253 .pool_max_idle_per_host(config.pool_max_idle_per_host)
254 .pool_idle_timeout(Duration::from_secs(90))
255 .connect_timeout(Duration::from_secs(10));
256
257 if let Some(keepalive) = config.tcp_keepalive {
258 builder = builder.tcp_keepalive(keepalive);
259 }
260
261 let client = builder.build()?;
262
263 Ok(Self { client, config })
264 }
265
266 pub fn from_crawl_config(config: &CrawlConfig) -> Result<Self, CrawlError> {
275 Self::new(HttpClientConfig::from(config))
276 }
277
278 pub fn high_throughput(config: HttpClientConfig) -> Result<Self, CrawlError> {
282 let cfg = HttpClientConfig {
283 pool_max_idle_per_host: 64,
284 pool_max_idle: 128,
285 tcp_keepalive: Some(Duration::from_secs(60)),
286 ..config
287 };
288 Self::new(cfg)
289 }
290
291 pub async fn fetch(&self, url: &Url) -> Result<FetchResult, CrawlError> {
302 self.fetch_with_redirects(url, self.config.max_redirects)
303 .await
304 }
305
306 pub async fn fetch_with_redirects(
315 &self,
316 url: &Url,
317 max_hops: usize,
318 ) -> Result<FetchResult, CrawlError> {
319 let mut current_url = url.clone();
320 let mut hops: Vec<RedirectHop> = Vec::new();
321
322 for _ in 0..=max_hops {
323 match self.fetch_once(¤t_url).await {
324 Ok((final_url, status, headers, body, elapsed)) => {
325 if status.is_redirection() {
326 let next_url = headers
327 .iter()
328 .find(|(k, _)| k.eq_ignore_ascii_case("location"))
329 .map(|(_, v)| v.clone());
330
331 match next_url {
332 Some(loc) => {
333 let resolved = current_url.join(&loc)?;
334 hops.push(RedirectHop {
335 from: current_url.clone(),
336 to: resolved.clone(),
337 status_code: status.as_u16(),
338 });
339 current_url = resolved;
340 continue;
341 }
342 None => {
343 let body_size = body.len();
345 return Ok(FetchResult {
346 final_url,
347 status_code: status.as_u16(),
348 headers,
349 body,
350 response_time: elapsed,
351 body_size,
352 fetched_at: chrono::Utc::now(),
353 });
354 }
355 }
356 }
357
358 let body_size = body.len();
359 return Ok(FetchResult {
360 final_url,
361 status_code: status.as_u16(),
362 headers,
363 body,
364 response_time: elapsed,
365 body_size,
366 fetched_at: chrono::Utc::now(),
367 });
368 }
369 Err(CrawlError::RequestFailed(e)) => {
370 return Err(CrawlError::RequestFailed(e));
371 }
372 Err(e) => return Err(e),
373 }
374 }
375
376 Err(CrawlError::TooManyRedirects(max_hops))
377 }
378
379 async fn fetch_once(
383 &self,
384 url: &Url,
385 ) -> Result<(Url, StatusCode, Vec<(String, String)>, String, Duration), CrawlError> {
386 let mut last_error: Option<CrawlError> = None;
387 let max_retries = self.config.retry_policy.max_retries;
388
389 for attempt in 0..=max_retries {
390 let start = Instant::now();
391 let user_agent = self.config.user_agent.next();
392
393 let result = self
394 .client
395 .get(url.as_str())
396 .header(USER_AGENT, user_agent)
397 .send()
398 .await;
399
400 match result {
401 Ok(response) => {
402 let status = response.status();
403 let elapsed = start.elapsed();
404 let headers: Vec<(String, String)> = response
405 .headers()
406 .iter()
407 .map(|(k, v)| {
408 (
409 k.as_str().to_string(),
410 String::from_utf8_lossy(v.as_bytes()).to_string(),
411 )
412 })
413 .collect();
414
415 if self.config.retry_policy.is_retryable(status.as_u16())
416 && attempt < max_retries
417 {
418 let backoff = self.config.retry_policy.backoff_duration(attempt);
419
420 if status == StatusCode::TOO_MANY_REQUESTS {
422 if let Some(retry_after) = headers
423 .iter()
424 .find(|(k, _)| k.eq_ignore_ascii_case("retry-after"))
425 .and_then(|(_, v)| v.parse::<u64>().ok())
426 {
427 let wait = Duration::from_secs(retry_after).max(backoff);
428 tracing::warn!(
429 url = %url,
430 status = status.as_u16(),
431 retry_after = retry_after,
432 "429 Too Many Requests, waiting before retry"
433 );
434 sleep(wait).await;
435 continue;
436 }
437 }
438
439 tracing::warn!(
440 url = %url,
441 status = status.as_u16(),
442 attempt = attempt + 1,
443 backoff_ms = backoff.as_millis(),
444 "Retrying after retryable status"
445 );
446 sleep(backoff).await;
447 continue;
448 }
449
450 let final_url = response.url().clone();
452
453 let body = if self.config.max_body_size > 0 {
454 let bytes = response.bytes().await.map_err(CrawlError::RequestFailed)?;
455 let limited = &bytes[..bytes.len().min(self.config.max_body_size)];
456 String::from_utf8_lossy(limited).to_string()
457 } else {
458 response.text().await.map_err(CrawlError::RequestFailed)?
459 };
460
461 return Ok((final_url, status, headers, body, elapsed));
462 }
463 Err(e) => {
464 if (e.is_timeout() || e.is_connect()) && attempt < max_retries {
465 let backoff = self.config.retry_policy.backoff_duration(attempt);
466 tracing::warn!(
467 url = %url,
468 error = %e,
469 attempt = attempt + 1,
470 backoff_ms = backoff.as_millis(),
471 "Retrying after network error"
472 );
473 sleep(backoff).await;
474 last_error = Some(CrawlError::RequestFailed(e));
475 continue;
476 }
477 return Err(CrawlError::RequestFailed(e));
478 }
479 }
480 }
481
482 Err(last_error.unwrap_or(CrawlError::MaxRetriesExceeded(max_retries)))
483 }
484
485 pub fn inner(&self) -> &Client {
487 &self.client
488 }
489
490 pub fn config(&self) -> &HttpClientConfig {
492 &self.config
493 }
494
495 pub async fn fetch_stream<F>(
505 &self,
506 url: &Url,
507 mut on_chunk: F,
508 ) -> Result<FetchResult, CrawlError>
509 where
510 F: FnMut(&str) + Send,
511 {
512 let mut current_url = url.clone();
513 let mut hops: Vec<RedirectHop> = Vec::new();
514
515 for _ in 0..=self.config.max_redirects {
516 let start = Instant::now();
517 let user_agent = self.config.user_agent.next();
518
519 let response = self
520 .client
521 .get(current_url.as_str())
522 .header(USER_AGENT, user_agent)
523 .send()
524 .await
525 .map_err(CrawlError::RequestFailed)?;
526
527 let status = response.status();
528 let elapsed = start.elapsed();
529 let headers: Vec<(String, String)> = response
530 .headers()
531 .iter()
532 .map(|(k, v)| {
533 (
534 k.as_str().to_string(),
535 String::from_utf8_lossy(v.as_bytes()).to_string(),
536 )
537 })
538 .collect();
539
540 if status.is_redirection() {
541 let next_url = headers
542 .iter()
543 .find(|(k, _)| k.eq_ignore_ascii_case("location"))
544 .map(|(_, v)| v.clone());
545
546 match next_url {
547 Some(loc) => {
548 let resolved = current_url.join(&loc)?;
549 hops.push(RedirectHop {
550 from: current_url.clone(),
551 to: resolved.clone(),
552 status_code: status.as_u16(),
553 });
554 current_url = resolved;
555 continue;
556 }
557 None => {
558 let final_url = response.url().clone();
559 return Ok(FetchResult {
560 final_url,
561 status_code: status.as_u16(),
562 headers,
563 body: String::new(),
564 response_time: elapsed,
565 body_size: 0,
566 fetched_at: chrono::Utc::now(),
567 });
568 }
569 }
570 }
571
572 let final_url = response.url().clone();
573 let mut body = String::new();
574 let mut stream = response.bytes_stream();
575 let mut total_size: usize = 0;
576
577 while let Some(chunk_result) = stream.next().await {
578 let chunk = chunk_result.map_err(CrawlError::RequestFailed)?;
579 total_size += chunk.len();
580
581 if self.config.max_body_size > 0 && total_size > self.config.max_body_size {
582 break;
583 }
584
585 let chunk_str = String::from_utf8_lossy(&chunk);
586 on_chunk(&chunk_str);
587 body.push_str(&chunk_str);
588 }
589
590 return Ok(FetchResult {
591 final_url,
592 status_code: status.as_u16(),
593 headers,
594 body,
595 response_time: elapsed,
596 body_size: total_size,
597 fetched_at: chrono::Utc::now(),
598 });
599 }
600
601 Err(CrawlError::TooManyRedirects(self.config.max_redirects))
602 }
603
604 pub async fn fetch_reader(&self, url: &Url) -> Result<FetchStreamReader, CrawlError> {
613 let start = Instant::now();
614 let user_agent = self.config.user_agent.next();
615
616 let response = self
617 .client
618 .get(url.as_str())
619 .header(USER_AGENT, user_agent)
620 .send()
621 .await
622 .map_err(CrawlError::RequestFailed)?;
623
624 let status = response.status();
625 let elapsed = start.elapsed();
626 let headers: Vec<(String, String)> = response
627 .headers()
628 .iter()
629 .map(|(k, v)| {
630 (
631 k.as_str().to_string(),
632 String::from_utf8_lossy(v.as_bytes()).to_string(),
633 )
634 })
635 .collect();
636 let final_url = response.url().clone();
637 let max_body_size = self.config.max_body_size;
638
639 let stream = response.bytes_stream().take_while(move |result| {
640 let should_continue = match result {
641 Ok(_bytes) => {
642 true
645 }
646 Err(_) => false,
647 };
648 async move { should_continue }
649 });
650
651 Ok(FetchStreamReader {
652 final_url,
653 status_code: status.as_u16(),
654 headers,
655 response_time: elapsed,
656 stream: Box::pin(stream),
657 body_size: 0,
658 max_body_size,
659 })
660 }
661}
662
663pub struct FetchStreamReader {
686 pub final_url: Url,
687 pub status_code: u16,
688 pub headers: Vec<(String, String)>,
689 pub response_time: Duration,
690 stream: Pin<Box<dyn Stream<Item = Result<bytes::Bytes, reqwest::Error>> + Send>>,
691 pub body_size: usize,
692 max_body_size: usize,
693}
694
695impl FetchStreamReader {
696 pub async fn next_chunk(&mut self) -> Result<Option<Vec<u8>>, CrawlError> {
705 if self.max_body_size > 0 && self.body_size >= self.max_body_size {
706 return Ok(None);
707 }
708
709 match self.stream.next().await {
710 Some(Ok(chunk)) => {
711 let chunk: bytes::Bytes = chunk;
712 let len = chunk.len();
713 self.body_size += len;
714 Ok(Some(chunk.to_vec()))
715 }
716 Some(Err(e)) => Err(CrawlError::RequestFailed(e)),
717 None => Ok(None),
718 }
719 }
720
721 pub async fn read_body(&mut self) -> Result<String, CrawlError> {
730 let mut body = String::new();
731 while let Some(chunk) = self.next_chunk().await? {
732 body.push_str(&String::from_utf8_lossy(&chunk));
733 }
734 Ok(body)
735 }
736
737 pub async fn into_fetch_result(mut self) -> Result<FetchResult, CrawlError> {
746 let body = self.read_body().await?;
747 let body_size = self.body_size;
748 Ok(FetchResult {
749 final_url: self.final_url,
750 status_code: self.status_code,
751 headers: self.headers,
752 body,
753 response_time: self.response_time,
754 body_size,
755 fetched_at: chrono::Utc::now(),
756 })
757 }
758}
759
760#[cfg(test)]
761mod tests {
762 use super::*;
763
764 #[test]
765 fn test_retry_policy_default() {
766 let policy = RetryPolicy::default();
767 assert_eq!(policy.max_retries, 3);
768 assert_eq!(policy.initial_backoff, Duration::from_secs(1));
769 assert_eq!(policy.max_backoff, Duration::from_secs(30));
770 assert!((policy.backoff_multiplier - 2.0).abs() < f64::EPSILON);
771 }
772
773 #[test]
774 fn test_retry_policy_backoff_duration() {
775 let policy = RetryPolicy::default();
776 assert_eq!(policy.backoff_duration(0), Duration::from_secs(1));
777 assert_eq!(policy.backoff_duration(1), Duration::from_secs(2));
778 assert_eq!(policy.backoff_duration(2), Duration::from_secs(4));
779 assert_eq!(policy.backoff_duration(3), Duration::from_secs(8));
780 assert_eq!(policy.backoff_duration(10), Duration::from_secs(30));
782 }
783
784 #[test]
785 fn test_retry_policy_is_retryable() {
786 let policy = RetryPolicy::default();
787 assert!(policy.is_retryable(429));
788 assert!(policy.is_retryable(500));
789 assert!(policy.is_retryable(502));
790 assert!(policy.is_retryable(503));
791 assert!(policy.is_retryable(504));
792 assert!(!policy.is_retryable(200));
793 assert!(!policy.is_retryable(404));
794 }
795
796 #[test]
797 fn test_user_agent_rotator() {
798 let rotator = UserAgentRotator::new(vec![
799 "agent-1".to_string(),
800 "agent-2".to_string(),
801 "agent-3".to_string(),
802 ]);
803 assert_eq!(rotator.len(), 3);
804 assert!(!rotator.is_empty());
805 assert_eq!(rotator.next(), "agent-1");
806 assert_eq!(rotator.next(), "agent-2");
807 assert_eq!(rotator.next(), "agent-3");
808 assert_eq!(rotator.next(), "agent-1"); }
810
811 #[test]
812 fn test_user_agent_rotator_default() {
813 let rotator = UserAgentRotator::default();
814 assert_eq!(rotator.len(), 1);
815 let agent = rotator.next().to_string();
816 assert!(agent.starts_with("crawlkit/"));
817 }
818
819 #[test]
820 fn test_http_client_config_from_crawl_config() {
821 let crawl_config = CrawlConfig::default();
822 let http_config = HttpClientConfig::from(&crawl_config);
823 assert_eq!(http_config.timeout, Duration::from_secs(30));
824 assert_eq!(http_config.max_redirects, 20);
825 assert_eq!(http_config.max_body_size, 10 * 1024 * 1024);
826 }
827
828 #[tokio::test]
829 async fn test_http_client_creation() {
830 let config = HttpClientConfig::from(&CrawlConfig::default());
831 let client = HttpClient::new(config);
832 assert!(client.is_ok());
833 }
834}