use std::pin::Pin;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use futures::stream::Stream;
use futures::StreamExt;
use reqwest::header::USER_AGENT;
use reqwest::{Client, StatusCode};
use serde::{Deserialize, Serialize};
use tokio::time::sleep;
use url::Url;
use crate::{CrawlConfig, CrawlError, FetchResult, RedirectHop};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RetryPolicy {
pub max_retries: usize,
#[serde(with = "crate::duration_ms")]
pub initial_backoff: Duration,
#[serde(with = "crate::duration_ms")]
pub max_backoff: Duration,
pub backoff_multiplier: f64,
pub retryable_statuses: Vec<u16>,
}
impl Default for RetryPolicy {
fn default() -> Self {
Self {
max_retries: 3,
initial_backoff: Duration::from_secs(1),
max_backoff: Duration::from_secs(30),
backoff_multiplier: 2.0,
retryable_statuses: vec![429, 500, 502, 503, 504],
}
}
}
impl RetryPolicy {
pub fn backoff_duration(&self, attempt: usize) -> Duration {
let base = self.initial_backoff.as_secs_f64();
let backoff = base * self.backoff_multiplier.powi(attempt as i32);
let capped = backoff.min(self.max_backoff.as_secs_f64());
Duration::from_secs_f64(capped)
}
pub fn is_retryable(&self, status: u16) -> bool {
self.retryable_statuses.contains(&status)
}
}
#[derive(Debug)]
pub struct UserAgentRotator {
agents: Vec<String>,
index: AtomicUsize,
}
impl UserAgentRotator {
pub fn new(agents: Vec<String>) -> Self {
assert!(
!agents.is_empty(),
"UserAgentRotator requires at least one user-agent"
);
Self {
agents,
index: AtomicUsize::new(0),
}
}
pub fn next(&self) -> &str {
let idx = self.index.fetch_add(1, Ordering::AcqRel);
&self.agents[idx % self.agents.len()]
}
pub fn len(&self) -> usize {
self.agents.len()
}
pub fn is_empty(&self) -> bool {
self.agents.is_empty()
}
}
impl Default for UserAgentRotator {
fn default() -> Self {
Self::new(vec![format!("crawlkit/{}", env!("CARGO_PKG_VERSION"))])
}
}
#[derive(Debug, Clone)]
pub struct HttpClientConfig {
pub timeout: Duration,
pub max_redirects: usize,
pub retry_policy: RetryPolicy,
pub user_agent: Arc<UserAgentRotator>,
pub max_body_size: usize,
pub pool_max_idle_per_host: usize,
pub pool_max_idle: usize,
pub tcp_keepalive: Option<Duration>,
}
impl From<&CrawlConfig> for HttpClientConfig {
fn from(config: &CrawlConfig) -> Self {
Self {
timeout: config.request_timeout,
max_redirects: config.max_redirects,
retry_policy: RetryPolicy::default(),
user_agent: Arc::new(UserAgentRotator::new(vec![config.user_agent.clone()])),
max_body_size: 10 * 1024 * 1024, pool_max_idle_per_host: 16,
pool_max_idle: 32,
tcp_keepalive: Some(Duration::from_secs(60)),
}
}
}
pub struct HttpClient {
client: Client,
config: HttpClientConfig,
}
impl HttpClient {
pub fn new(config: HttpClientConfig) -> Result<Self, CrawlError> {
let mut builder = Client::builder()
.timeout(config.timeout)
.redirect(reqwest::redirect::Policy::limited(config.max_redirects))
.user_agent(config.user_agent.next())
.https_only(true)
.http1_only()
.pool_max_idle_per_host(config.pool_max_idle_per_host)
.pool_idle_timeout(Duration::from_secs(90))
.connect_timeout(Duration::from_secs(10));
if let Some(keepalive) = config.tcp_keepalive {
builder = builder.tcp_keepalive(keepalive);
}
let client = builder.build()?;
Ok(Self { client, config })
}
pub fn from_crawl_config(config: &CrawlConfig) -> Result<Self, CrawlError> {
Self::new(HttpClientConfig::from(config))
}
pub fn high_throughput(config: HttpClientConfig) -> Result<Self, CrawlError> {
let cfg = HttpClientConfig {
pool_max_idle_per_host: 64,
pool_max_idle: 128,
tcp_keepalive: Some(Duration::from_secs(60)),
..config
};
Self::new(cfg)
}
pub async fn fetch(&self, url: &Url) -> Result<FetchResult, CrawlError> {
self.fetch_with_redirects(url, self.config.max_redirects)
.await
}
pub async fn fetch_with_redirects(
&self,
url: &Url,
max_hops: usize,
) -> Result<FetchResult, CrawlError> {
let mut current_url = url.clone();
let mut hops: Vec<RedirectHop> = Vec::new();
for _ in 0..=max_hops {
match self.fetch_once(¤t_url).await {
Ok((final_url, status, headers, body, elapsed)) => {
if status.is_redirection() {
let next_url = headers
.iter()
.find(|(k, _)| k.eq_ignore_ascii_case("location"))
.map(|(_, v)| v.clone());
match next_url {
Some(loc) => {
let resolved = current_url.join(&loc)?;
hops.push(RedirectHop {
from: current_url.clone(),
to: resolved.clone(),
status_code: status.as_u16(),
});
current_url = resolved;
continue;
}
None => {
let body_size = body.len();
return Ok(FetchResult {
final_url,
status_code: status.as_u16(),
headers,
body,
response_time: elapsed,
body_size,
fetched_at: chrono::Utc::now(),
});
}
}
}
let body_size = body.len();
return Ok(FetchResult {
final_url,
status_code: status.as_u16(),
headers,
body,
response_time: elapsed,
body_size,
fetched_at: chrono::Utc::now(),
});
}
Err(CrawlError::RequestFailed(e)) => {
return Err(CrawlError::RequestFailed(e));
}
Err(e) => return Err(e),
}
}
Err(CrawlError::TooManyRedirects(max_hops))
}
async fn fetch_once(
&self,
url: &Url,
) -> Result<(Url, StatusCode, Vec<(String, String)>, String, Duration), CrawlError> {
let mut last_error: Option<CrawlError> = None;
let max_retries = self.config.retry_policy.max_retries;
for attempt in 0..=max_retries {
let start = Instant::now();
let user_agent = self.config.user_agent.next();
let result = self
.client
.get(url.as_str())
.header(USER_AGENT, user_agent)
.send()
.await;
match result {
Ok(response) => {
let status = response.status();
let elapsed = start.elapsed();
let headers: Vec<(String, String)> = response
.headers()
.iter()
.map(|(k, v)| {
(
k.as_str().to_string(),
String::from_utf8_lossy(v.as_bytes()).to_string(),
)
})
.collect();
if self.config.retry_policy.is_retryable(status.as_u16())
&& attempt < max_retries
{
let backoff = self.config.retry_policy.backoff_duration(attempt);
if status == StatusCode::TOO_MANY_REQUESTS {
if let Some(retry_after) = headers
.iter()
.find(|(k, _)| k.eq_ignore_ascii_case("retry-after"))
.and_then(|(_, v)| v.parse::<u64>().ok())
{
let wait = Duration::from_secs(retry_after).max(backoff);
tracing::warn!(
url = %url,
status = status.as_u16(),
retry_after = retry_after,
"429 Too Many Requests, waiting before retry"
);
sleep(wait).await;
continue;
}
}
tracing::warn!(
url = %url,
status = status.as_u16(),
attempt = attempt + 1,
backoff_ms = backoff.as_millis(),
"Retrying after retryable status"
);
sleep(backoff).await;
continue;
}
let final_url = response.url().clone();
let body = if self.config.max_body_size > 0 {
let bytes = response.bytes().await.map_err(CrawlError::RequestFailed)?;
let limited = &bytes[..bytes.len().min(self.config.max_body_size)];
String::from_utf8_lossy(limited).to_string()
} else {
response.text().await.map_err(CrawlError::RequestFailed)?
};
return Ok((final_url, status, headers, body, elapsed));
}
Err(e) => {
if (e.is_timeout() || e.is_connect()) && attempt < max_retries {
let backoff = self.config.retry_policy.backoff_duration(attempt);
tracing::warn!(
url = %url,
error = %e,
attempt = attempt + 1,
backoff_ms = backoff.as_millis(),
"Retrying after network error"
);
sleep(backoff).await;
last_error = Some(CrawlError::RequestFailed(e));
continue;
}
return Err(CrawlError::RequestFailed(e));
}
}
}
Err(last_error.unwrap_or(CrawlError::MaxRetriesExceeded(max_retries)))
}
pub fn inner(&self) -> &Client {
&self.client
}
pub fn config(&self) -> &HttpClientConfig {
&self.config
}
pub async fn fetch_stream<F>(
&self,
url: &Url,
mut on_chunk: F,
) -> Result<FetchResult, CrawlError>
where
F: FnMut(&str) + Send,
{
let mut current_url = url.clone();
let mut hops: Vec<RedirectHop> = Vec::new();
for _ in 0..=self.config.max_redirects {
let start = Instant::now();
let user_agent = self.config.user_agent.next();
let response = self
.client
.get(current_url.as_str())
.header(USER_AGENT, user_agent)
.send()
.await
.map_err(CrawlError::RequestFailed)?;
let status = response.status();
let elapsed = start.elapsed();
let headers: Vec<(String, String)> = response
.headers()
.iter()
.map(|(k, v)| {
(
k.as_str().to_string(),
String::from_utf8_lossy(v.as_bytes()).to_string(),
)
})
.collect();
if status.is_redirection() {
let next_url = headers
.iter()
.find(|(k, _)| k.eq_ignore_ascii_case("location"))
.map(|(_, v)| v.clone());
match next_url {
Some(loc) => {
let resolved = current_url.join(&loc)?;
hops.push(RedirectHop {
from: current_url.clone(),
to: resolved.clone(),
status_code: status.as_u16(),
});
current_url = resolved;
continue;
}
None => {
let final_url = response.url().clone();
return Ok(FetchResult {
final_url,
status_code: status.as_u16(),
headers,
body: String::new(),
response_time: elapsed,
body_size: 0,
fetched_at: chrono::Utc::now(),
});
}
}
}
let final_url = response.url().clone();
let mut body = String::new();
let mut stream = response.bytes_stream();
let mut total_size: usize = 0;
while let Some(chunk_result) = stream.next().await {
let chunk = chunk_result.map_err(CrawlError::RequestFailed)?;
total_size += chunk.len();
if self.config.max_body_size > 0 && total_size > self.config.max_body_size {
break;
}
let chunk_str = String::from_utf8_lossy(&chunk);
on_chunk(&chunk_str);
body.push_str(&chunk_str);
}
return Ok(FetchResult {
final_url,
status_code: status.as_u16(),
headers,
body,
response_time: elapsed,
body_size: total_size,
fetched_at: chrono::Utc::now(),
});
}
Err(CrawlError::TooManyRedirects(self.config.max_redirects))
}
pub async fn fetch_reader(&self, url: &Url) -> Result<FetchStreamReader, CrawlError> {
let start = Instant::now();
let user_agent = self.config.user_agent.next();
let response = self
.client
.get(url.as_str())
.header(USER_AGENT, user_agent)
.send()
.await
.map_err(CrawlError::RequestFailed)?;
let status = response.status();
let elapsed = start.elapsed();
let headers: Vec<(String, String)> = response
.headers()
.iter()
.map(|(k, v)| {
(
k.as_str().to_string(),
String::from_utf8_lossy(v.as_bytes()).to_string(),
)
})
.collect();
let final_url = response.url().clone();
let max_body_size = self.config.max_body_size;
let stream = response.bytes_stream().take_while(move |result| {
let should_continue = match result {
Ok(_bytes) => {
true
}
Err(_) => false,
};
async move { should_continue }
});
Ok(FetchStreamReader {
final_url,
status_code: status.as_u16(),
headers,
response_time: elapsed,
stream: Box::pin(stream),
body_size: 0,
max_body_size,
})
}
}
pub struct FetchStreamReader {
pub final_url: Url,
pub status_code: u16,
pub headers: Vec<(String, String)>,
pub response_time: Duration,
stream: Pin<Box<dyn Stream<Item = Result<bytes::Bytes, reqwest::Error>> + Send>>,
pub body_size: usize,
max_body_size: usize,
}
impl FetchStreamReader {
pub async fn next_chunk(&mut self) -> Result<Option<Vec<u8>>, CrawlError> {
if self.max_body_size > 0 && self.body_size >= self.max_body_size {
return Ok(None);
}
match self.stream.next().await {
Some(Ok(chunk)) => {
let chunk: bytes::Bytes = chunk;
let len = chunk.len();
self.body_size += len;
Ok(Some(chunk.to_vec()))
}
Some(Err(e)) => Err(CrawlError::RequestFailed(e)),
None => Ok(None),
}
}
pub async fn read_body(&mut self) -> Result<String, CrawlError> {
let mut body = String::new();
while let Some(chunk) = self.next_chunk().await? {
body.push_str(&String::from_utf8_lossy(&chunk));
}
Ok(body)
}
pub async fn into_fetch_result(mut self) -> Result<FetchResult, CrawlError> {
let body = self.read_body().await?;
let body_size = self.body_size;
Ok(FetchResult {
final_url: self.final_url,
status_code: self.status_code,
headers: self.headers,
body,
response_time: self.response_time,
body_size,
fetched_at: chrono::Utc::now(),
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_retry_policy_default() {
let policy = RetryPolicy::default();
assert_eq!(policy.max_retries, 3);
assert_eq!(policy.initial_backoff, Duration::from_secs(1));
assert_eq!(policy.max_backoff, Duration::from_secs(30));
assert!((policy.backoff_multiplier - 2.0).abs() < f64::EPSILON);
}
#[test]
fn test_retry_policy_backoff_duration() {
let policy = RetryPolicy::default();
assert_eq!(policy.backoff_duration(0), Duration::from_secs(1));
assert_eq!(policy.backoff_duration(1), Duration::from_secs(2));
assert_eq!(policy.backoff_duration(2), Duration::from_secs(4));
assert_eq!(policy.backoff_duration(3), Duration::from_secs(8));
assert_eq!(policy.backoff_duration(10), Duration::from_secs(30));
}
#[test]
fn test_retry_policy_is_retryable() {
let policy = RetryPolicy::default();
assert!(policy.is_retryable(429));
assert!(policy.is_retryable(500));
assert!(policy.is_retryable(502));
assert!(policy.is_retryable(503));
assert!(policy.is_retryable(504));
assert!(!policy.is_retryable(200));
assert!(!policy.is_retryable(404));
}
#[test]
fn test_user_agent_rotator() {
let rotator = UserAgentRotator::new(vec![
"agent-1".to_string(),
"agent-2".to_string(),
"agent-3".to_string(),
]);
assert_eq!(rotator.len(), 3);
assert!(!rotator.is_empty());
assert_eq!(rotator.next(), "agent-1");
assert_eq!(rotator.next(), "agent-2");
assert_eq!(rotator.next(), "agent-3");
assert_eq!(rotator.next(), "agent-1"); }
#[test]
fn test_user_agent_rotator_default() {
let rotator = UserAgentRotator::default();
assert_eq!(rotator.len(), 1);
let agent = rotator.next().to_string();
assert!(agent.starts_with("crawlkit/"));
}
#[test]
fn test_http_client_config_from_crawl_config() {
let crawl_config = CrawlConfig::default();
let http_config = HttpClientConfig::from(&crawl_config);
assert_eq!(http_config.timeout, Duration::from_secs(30));
assert_eq!(http_config.max_redirects, 20);
assert_eq!(http_config.max_body_size, 10 * 1024 * 1024);
}
#[tokio::test]
async fn test_http_client_creation() {
let config = HttpClientConfig::from(&CrawlConfig::default());
let client = HttpClient::new(config);
assert!(client.is_ok());
}
}