Skip to main content

spider_middleware/
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, trace, warn};
14
15use spider_util::error::SpiderError;
16use crate::middleware::{Middleware, MiddlewareAction};
17use spider_util::request::Request;
18use spider_util::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        trace!("Processing response for URL: {} with status: {}", response.url, response.status);
96
97        if self.retry_http_codes.contains(&response.status.as_u16()) {
98            let mut request = response.request_from_response();
99            let current_attempts = request.get_retry_attempts();
100
101            if current_attempts < self.max_retries {
102                request.increment_retry_attempts();
103                let delay = self.calculate_delay(current_attempts);
104                info!(
105                    "Retrying {} (status: {}, attempt {}/{}) after {:?}",
106                    request.url,
107                    response.status,
108                    current_attempts + 1,
109                    self.max_retries,
110                    delay
111                );
112                return Ok(MiddlewareAction::Retry(Box::new(request), delay));
113            } else {
114                warn!(
115                    "Max retries ({}) reached for {} (status: {}). Dropping response.",
116                    self.max_retries, request.url, response.status
117                );
118                return Ok(MiddlewareAction::Drop);
119            }
120        } else {
121            trace!("Response status {} is not in retry codes, continuing", response.status);
122        }
123
124        Ok(MiddlewareAction::Continue(response))
125    }
126
127    async fn handle_error(
128        &mut self,
129        request: &Request,
130        error: &SpiderError,
131    ) -> Result<MiddlewareAction<Request>, SpiderError> {
132        trace!("Handling error for request {}: {:?}", request.url, error);
133
134        if let SpiderError::ReqwestError(err_details) = error
135            && (err_details.is_connect || err_details.is_timeout)
136        {
137            let mut new_request = request.clone();
138            let current_attempts = new_request.get_retry_attempts();
139
140            if current_attempts < self.max_retries {
141                new_request.increment_retry_attempts();
142                let delay = self.calculate_delay(current_attempts);
143                info!(
144                    "Retrying {} (error: {}, attempt {}/{}) after {:?}",
145                    new_request.url,
146                    err_details.message,
147                    current_attempts + 1,
148                    self.max_retries,
149                    delay
150                );
151                return Ok(MiddlewareAction::Retry(Box::new(new_request), delay));
152            } else {
153                warn!(
154                    "Max retries ({}) reached for {} (error: {}). Dropping request.",
155                    self.max_retries, new_request.url, err_details.message
156                );
157                return Ok(MiddlewareAction::Drop);
158            }
159        } else {
160            trace!("Error is not a retryable error, returning original error");
161        }
162
163        Err(error.clone())
164    }
165}
166
167impl RetryMiddleware {
168    fn calculate_delay(&self, retries: u32) -> Duration {
169        let delay_secs = self.backoff_factor * (2.0f64.powi(retries as i32));
170        let delay = Duration::from_secs_f64(delay_secs);
171        delay.min(self.max_delay)
172    }
173}