spider_lib/middlewares/
retry.rs1use 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#[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 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}