spider_middleware/
retry.rs1use 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#[derive(Debug, Clone)]
22pub struct RetryMiddleware {
23 pub max_retries: u32,
25 pub retry_http_codes: Vec<u16>,
27 pub backoff_factor: f64,
29 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 pub fn new() -> Self {
49 Self::default()
50 }
51
52 pub fn max_retries(mut self, max_retries: u32) -> Self {
54 self.max_retries = max_retries;
55 self
56 }
57
58 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 pub fn backoff_factor(mut self, backoff_factor: f64) -> Self {
66 self.backoff_factor = backoff_factor;
67 self
68 }
69
70 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}