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