Skip to main content

spider_lib/middlewares/
retry.rs

1//! Retry Middleware for handling failed requests.
2//!
3//! This module provides the `RetryMiddleware`, which automatically retries requests
4//! that fail due to specific HTTP status codes or network errors (e.g., connection issues, timeouts).
5//! It implements an exponential backoff strategy to space out retry attempts,
6//! respecting a configurable maximum number of retries and a maximum delay between attempts.
7//!
8//! The middleware intercepts responses and errors, deciding whether to re-enqueue
9//! the request for another attempt or to drop it if the retry limit is exceeded.
10
11use async_trait::async_trait;
12use std::time::Duration;
13use tracing::{info, warn};
14
15use crate::error::SpiderError;
16use crate::middleware::{Middleware, MiddlewareAction};
17use crate::request::Request;
18use crate::response::Response;
19
20/// Middleware that retries failed requests.
21#[derive(Debug, Clone)]
22pub struct RetryMiddleware {
23    /// Maximum number of times to retry a request.
24    pub max_retries: u32,
25    /// HTTP status codes that should trigger a retry.
26    pub retry_http_codes: Vec<u16>,
27    /// Factor for exponential backoff (delay = backoff_factor * (2^retries)).
28    pub backoff_factor: f64,
29    /// Maximum delay between retries.
30    pub max_delay: Duration,
31}
32
33impl Default for RetryMiddleware {
34    fn default() -> Self {
35        let middleware = RetryMiddleware {
36            max_retries: 3,
37            retry_http_codes: vec![500, 502, 503, 504, 408, 429],
38            backoff_factor: 1.0,
39            max_delay: Duration::from_secs(180),
40        };
41        info!("Initializing RetryMiddleware with config: {:?}", middleware);
42        middleware
43    }
44}
45
46impl RetryMiddleware {
47    /// Creates a new `RetryMiddleware` with default settings.
48    pub fn new() -> Self {
49        Self::default()
50    }
51
52    /// Sets the maximum number of times to retry a request.
53    pub fn max_retries(mut self, max_retries: u32) -> Self {
54        self.max_retries = max_retries;
55        self
56    }
57
58    /// Sets the HTTP status codes that should trigger a retry.
59    pub fn retry_http_codes(mut self, retry_http_codes: Vec<u16>) -> Self {
60        self.retry_http_codes = retry_http_codes;
61        self
62    }
63
64    /// Sets the factor for exponential backoff.
65    pub fn backoff_factor(mut self, backoff_factor: f64) -> Self {
66        self.backoff_factor = backoff_factor;
67        self
68    }
69
70    /// Sets the maximum delay between retries.
71    pub fn max_delay(mut self, max_delay: Duration) -> Self {
72        self.max_delay = max_delay;
73        self
74    }
75}
76
77#[async_trait]
78impl<C: Send + Sync> Middleware<C> for RetryMiddleware {
79    fn name(&self) -> &str {
80        "RetryMiddleware"
81    }
82
83    async fn process_request(
84        &mut self,
85        _client: &C,
86        request: Request,
87    ) -> Result<MiddlewareAction<Request>, SpiderError> {
88        Ok(MiddlewareAction::Continue(request))
89    }
90
91    async fn process_response(
92        &mut self,
93        response: Response,
94    ) -> Result<MiddlewareAction<Response>, SpiderError> {
95        if self.retry_http_codes.contains(&response.status.as_u16()) {
96            let mut request = response.request_from_response();
97            let current_attempts = request.get_retry_attempts();
98
99            if current_attempts < self.max_retries {
100                request.increment_retry_attempts();
101                let delay = self.calculate_delay(current_attempts);
102                warn!(
103                    "Retrying {} (status: {}, attempt {}/{}) after {:?}",
104                    request.url,
105                    response.status,
106                    current_attempts + 1,
107                    self.max_retries,
108                    delay
109                );
110                return Ok(MiddlewareAction::Retry(Box::new(request), delay));
111            } else {
112                warn!(
113                    "Max retries ({}) reached for {} (status: {}). Dropping response.",
114                    self.max_retries, request.url, response.status
115                );
116                return Ok(MiddlewareAction::Drop);
117            }
118        }
119
120        Ok(MiddlewareAction::Continue(response))
121    }
122
123    async fn handle_error(
124        &mut self,
125        request: &Request,
126        error: &SpiderError,
127    ) -> Result<MiddlewareAction<Request>, SpiderError> {
128        if let SpiderError::ReqwestError(err_details) = error
129            && (err_details.is_connect || err_details.is_timeout)
130        {
131            let mut new_request = request.clone();
132            let current_attempts = new_request.get_retry_attempts();
133
134            if current_attempts < self.max_retries {
135                new_request.increment_retry_attempts();
136                let delay = self.calculate_delay(current_attempts);
137                warn!(
138                    "Retrying {} (error: {}, attempt {}/{}) after {:?}",
139                    new_request.url,
140                    err_details.message,
141                    current_attempts + 1,
142                    self.max_retries,
143                    delay
144                );
145                return Ok(MiddlewareAction::Retry(Box::new(new_request), delay));
146            } else {
147                warn!(
148                    "Max retries ({}) reached for {} (error: {}).",
149                    self.max_retries, new_request.url, err_details.message
150                );
151            }
152        }
153
154        Err(error.clone())
155    }
156}
157
158impl RetryMiddleware {
159    fn calculate_delay(&self, retries: u32) -> Duration {
160        let delay_secs = self.backoff_factor * (2.0f64.powi(retries as i32));
161        let delay = Duration::from_secs_f64(delay_secs);
162        delay.min(self.max_delay)
163    }
164}