1use std::fmt;
2use std::sync::Arc;
3use std::time::Duration;
4
5use http::StatusCode;
6
7use crate::core::block::BlockRules;
8use crate::core::config::ProxyConfig;
9use crate::core::error::{Error, ErrorCategory, Kind, Result};
10use crate::core::response::Response;
11
12mod wait;
13
14pub use wait::WaitFormat;
15pub(crate) use wait::parse_retry_after;
16use wait::{RETRY_AFTER, WaitHeader};
17
18pub(crate) enum RetryPlan {
19 Stop,
20 Exhausted,
21 Declined,
22 Backoff(Duration),
23}
24
25const GATEWAY_STATUSES: [StatusCode; 3] = [
26 StatusCode::BAD_GATEWAY,
27 StatusCode::SERVICE_UNAVAILABLE,
28 StatusCode::GATEWAY_TIMEOUT,
29];
30
31pub(crate) fn is_gateway_status(code: u16) -> bool {
32 GATEWAY_STATUSES
33 .iter()
34 .any(|status| status.as_u16() == code)
35}
36
37#[derive(Debug, Clone, Copy)]
38#[non_exhaustive]
39pub enum RetryTrigger {
40 ConnectionError,
41 Status(u16),
42 ServerError,
43 Timeout,
44}
45
46#[derive(Debug, Clone)]
47#[non_exhaustive]
48pub struct RetryPolicy {
49 pub(crate) max_retries: u32,
50 pub(crate) initial_backoff: Duration,
51 pub(crate) max_backoff: Duration,
52 pub(crate) max_retry_after: Duration,
53 pub(crate) backoff_factor: f64,
54 pub(crate) jitter: bool,
55 pub(crate) retry_on: Vec<RetryTrigger>,
56 pub(crate) allow_non_idempotent: bool,
57 pub(crate) proxies: Vec<ProxyConfig>,
58 pub(crate) retry_if: Vec<RetryIf>,
59 pub(crate) wait_headers: Vec<WaitHeader>,
60 pub(crate) skip_blocks: Option<BlockRules>,
61 pub(crate) retry_unsent: bool,
62}
63
64type ResponsePredicate = dyn Fn(&Response) -> bool + Send + Sync + std::panic::RefUnwindSafe;
65
66#[derive(Clone)]
67pub(crate) struct RetryIf(Arc<ResponsePredicate>);
68
69impl fmt::Debug for RetryIf {
70 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
71 f.write_str("RetryIf(..)")
72 }
73}
74
75impl Default for RetryPolicy {
76 fn default() -> Self {
77 Self::none()
78 }
79}
80
81impl RetryPolicy {
82 pub fn none() -> Self {
83 Self {
84 max_retries: 0,
85 initial_backoff: Duration::from_millis(0),
86 max_backoff: Duration::from_millis(0),
87 max_retry_after: Duration::MAX,
88 backoff_factor: 1.0,
89 jitter: false,
90 retry_on: Vec::new(),
91 allow_non_idempotent: false,
92 proxies: Vec::new(),
93 retry_if: Vec::new(),
94 wait_headers: Vec::new(),
95 skip_blocks: None,
96 retry_unsent: false,
97 }
98 }
99
100 pub fn transient() -> Self {
101 Self {
102 max_retries: 3,
103 initial_backoff: Duration::from_millis(100),
104 max_backoff: Duration::from_secs(1),
105 max_retry_after: Duration::from_secs(60),
106 backoff_factor: 2.0,
107 jitter: true,
108 retry_on: [
109 RetryTrigger::ConnectionError,
110 RetryTrigger::Status(StatusCode::TOO_MANY_REQUESTS.as_u16()),
111 ]
112 .into_iter()
113 .chain(GATEWAY_STATUSES.map(|status| RetryTrigger::Status(status.as_u16())))
114 .chain([RetryTrigger::Timeout])
115 .collect(),
116 allow_non_idempotent: false,
117 proxies: Vec::new(),
118 retry_if: Vec::new(),
119 wait_headers: Vec::new(),
120 skip_blocks: None,
121 retry_unsent: false,
122 }
123 }
124
125 pub fn max_retries(mut self, n: u32) -> Self {
126 self.max_retries = n;
127 self
128 }
129
130 pub fn initial_backoff(mut self, d: Duration) -> Self {
131 self.initial_backoff = d;
132 self
133 }
134
135 pub fn max_backoff(mut self, d: Duration) -> Self {
136 self.max_backoff = d;
137 self
138 }
139
140 pub fn backoff_factor(mut self, factor: f64) -> Self {
141 self.backoff_factor = factor;
142 self
143 }
144
145 pub fn jitter(mut self, on: bool) -> Self {
146 self.jitter = on;
147 self
148 }
149
150 pub fn max_retry_after(mut self, d: Duration) -> Self {
151 self.max_retry_after = d;
152 self
153 }
154
155 pub fn on_status(mut self, code: u16) -> Self {
156 self.retry_on.push(RetryTrigger::Status(code));
157 self
158 }
159
160 pub fn retry_on(mut self, triggers: impl IntoIterator<Item = RetryTrigger>) -> Self {
161 self.retry_on = triggers.into_iter().collect();
162 self
163 }
164
165 pub fn allow_non_idempotent(mut self, allow: bool) -> Self {
166 self.allow_non_idempotent = allow;
167 self
168 }
169
170 pub fn rotate_proxies<P: Into<ProxyConfig>>(
171 mut self,
172 proxies: impl IntoIterator<Item = P>,
173 ) -> Self {
174 self.proxies = proxies.into_iter().map(Into::into).collect();
175 self
176 }
177
178 pub fn retry_if<F>(mut self, f: F) -> Self
179 where
180 F: Fn(&Response) -> bool + Send + Sync + std::panic::RefUnwindSafe + 'static,
181 {
182 self.retry_if.push(RetryIf(Arc::new(f)));
183 self
184 }
185
186 pub fn wait_header(mut self, name: impl Into<String>, format: WaitFormat) -> Self {
187 self.wait_headers.push(WaitHeader {
188 name: name.into(),
189 format,
190 });
191 self
192 }
193
194 pub fn skip_blocks(mut self, rules: BlockRules) -> Self {
195 match &mut self.skip_blocks {
196 Some(held) => held.extend(rules),
197 None => self.skip_blocks = Some(rules),
198 }
199 self
200 }
201
202 pub fn retry_unsent(mut self, enabled: bool) -> Self {
203 self.retry_unsent = enabled;
204 self
205 }
206
207 pub(crate) fn matches_response(&self, response: &Response) -> bool {
208 if self.is_block(response) {
209 return false;
210 }
211 self.matches_status(response.status().as_u16())
212 || self.retry_if.iter().any(|RetryIf(f)| f(response))
213 }
214
215 fn is_block(&self, response: &Response) -> bool {
216 self.skip_blocks
217 .as_ref()
218 .is_some_and(|rules| rules.check(response).is_some())
219 }
220
221 pub(crate) fn retries_unsent(&self, result: &Result<Response>) -> bool {
222 self.retry_unsent && result.as_ref().err().is_some_and(unsent)
223 }
224
225 pub(crate) fn server_wait(&self, response: &Response) -> Option<Duration> {
226 self.wait_in(|name| response.header(name))
227 }
228
229 fn wait_in<'a>(&self, header: impl Fn(&str) -> Option<&'a str>) -> Option<Duration> {
230 self.wait_headers
231 .iter()
232 .find_map(|wait| header(&wait.name).and_then(|value| wait.format.parse(value)))
233 .or_else(|| header(RETRY_AFTER).and_then(parse_retry_after))
234 }
235
236 pub(crate) fn proxy_for_retry(&self, retry: u32) -> Option<&ProxyConfig> {
237 let index = usize::try_from(retry.checked_sub(1)?).ok()?;
238 self.proxies.get(index.checked_rem(self.proxies.len())?)
239 }
240
241 pub(crate) fn plan(
242 &self,
243 result: &Result<Response>,
244 attempt: u32,
245 retryable_method: bool,
246 body_retryable: impl FnOnce() -> bool,
247 ) -> RetryPlan {
248 if self.is_none() {
249 return RetryPlan::Stop;
250 }
251 if attempt >= self.max_retries {
252 return match result {
253 Ok(response) if self.matches_response(response) => RetryPlan::Exhausted,
254 _ => RetryPlan::Stop,
255 };
256 }
257 if !self.wanted(result)
258 || !(retryable_method || self.retries_unsent(result))
259 || !body_retryable()
260 {
261 return RetryPlan::Stop;
262 }
263 match result.as_ref().ok().and_then(|r| self.server_wait(r)) {
264 Some(wait) if wait > self.max_retry_after => RetryPlan::Declined,
265 Some(wait) => RetryPlan::Backoff(wait),
266 None => RetryPlan::Backoff(self.backoff(attempt)),
267 }
268 }
269
270 fn wanted(&self, result: &Result<Response>) -> bool {
271 match result {
272 Ok(resp) => self.matches_response(resp),
273 Err(err) if err.is_timeout() => self.matches_timeout(),
274 Err(err) if err.is_retryable() => self.matches_connection_error(),
275 Err(_) => false,
276 }
277 }
278
279 #[cfg(feature = "html")]
280 pub(crate) fn or(self, fallback: RetryPolicy) -> RetryPolicy {
281 match self.is_none() {
282 true => fallback,
283 false => self,
284 }
285 }
286
287 pub(crate) fn is_none(&self) -> bool {
288 self.max_retries == 0
289 }
290
291 pub(crate) fn matches_status(&self, status: u16) -> bool {
292 for t in &self.retry_on {
293 match *t {
294 RetryTrigger::Status(s) if s == status => return true,
295 RetryTrigger::ServerError if (500..600).contains(&status) => return true,
296 _ => {}
297 }
298 }
299 false
300 }
301
302 pub(crate) fn matches_connection_error(&self) -> bool {
303 self.retry_on
304 .iter()
305 .any(|t| matches!(t, RetryTrigger::ConnectionError))
306 }
307
308 pub(crate) fn matches_timeout(&self) -> bool {
309 self.retry_on
310 .iter()
311 .any(|t| matches!(t, RetryTrigger::Timeout))
312 }
313
314 #[must_use]
315 pub fn backoff(&self, attempt: u32) -> Duration {
316 let base = self.initial_backoff.as_secs_f64();
317 let raw = base * self.backoff_factor.powi(attempt as i32);
318 let capped = raw.min(self.max_backoff.as_secs_f64());
319 let jittered = if self.jitter {
320 capped * cheap_jitter()
321 } else {
322 capped
323 };
324 Duration::try_from_secs_f64(jittered.max(0.0)).unwrap_or(self.max_backoff)
325 }
326}
327
328fn unsent(err: &Error) -> bool {
329 if err.follows_response() {
330 return false;
331 }
332 match err.category() {
333 ErrorCategory::Dns | ErrorCategory::Connect | ErrorCategory::Tls | ErrorCategory::Proxy => {
334 true
335 }
336 ErrorCategory::Timeout => err.kind() == Kind::Connect,
337 ErrorCategory::Protocol => err.is_refused_stream(),
338 _ => false,
339 }
340}
341
342fn cheap_jitter() -> f64 {
343 use rand::Rng;
344 rand::rng().random_range(0.0..=1.0)
345}
346
347#[cfg(test)]
348mod tests;