babelforce_manager_sdk/
retry.rs1use std::collections::HashSet;
11use std::time::Duration;
12
13use rand::Rng;
14use reqwest::{Method, Request, Response};
15use reqwest_middleware::Next;
16
17#[derive(Clone, Debug)]
20pub struct RetryPolicy {
21 pub max_retries: u32,
23 pub base_delay: Duration,
25 pub max_delay: Duration,
27 pub retry_status: Vec<u16>,
29}
30
31impl Default for RetryPolicy {
32 fn default() -> Self {
33 Self {
34 max_retries: 2,
35 base_delay: Duration::from_millis(250),
36 max_delay: Duration::from_secs(10),
37 retry_status: vec![429, 502, 503, 504],
38 }
39 }
40}
41
42impl RetryPolicy {
43 pub fn disabled() -> Self {
45 Self {
46 max_retries: 0,
47 ..Self::default()
48 }
49 }
50}
51
52fn idempotent(method: &Method) -> bool {
54 matches!(
55 *method,
56 Method::GET | Method::HEAD | Method::PUT | Method::DELETE | Method::OPTIONS
57 )
58}
59
60fn retry_after(resp: &Response) -> Option<Duration> {
62 let header = resp
63 .headers()
64 .get(reqwest::header::RETRY_AFTER)?
65 .to_str()
66 .ok()?;
67 if let Ok(secs) = header.trim().parse::<u64>() {
68 return Some(Duration::from_secs(secs));
69 }
70 let when = httpdate::parse_http_date(header.trim()).ok()?;
71 Some(
72 when.duration_since(std::time::SystemTime::now())
73 .unwrap_or(Duration::ZERO),
74 )
75}
76
77pub(crate) struct RetryMiddleware {
81 policy: RetryPolicy,
82 codes: HashSet<u16>,
83}
84
85impl RetryMiddleware {
86 pub(crate) fn new(policy: RetryPolicy) -> Self {
87 let codes = policy.retry_status.iter().copied().collect();
88 RetryMiddleware { policy, codes }
89 }
90
91 fn backoff(&self, attempt: u32) -> Duration {
93 let factor = 1u32.checked_shl(attempt.min(16)).unwrap_or(u32::MAX);
94 let ceiling = self
95 .policy
96 .base_delay
97 .saturating_mul(factor)
98 .min(self.policy.max_delay);
99 let ceiling_ms = u64::try_from(ceiling.as_millis()).unwrap_or(u64::MAX);
100 Duration::from_millis(rand::rng().random_range(0..=ceiling_ms))
101 }
102
103 fn retry_wait(
105 &self,
106 attempt: u32,
107 idempotent: bool,
108 result: &reqwest_middleware::Result<Response>,
109 ) -> Option<Duration> {
110 if attempt >= self.policy.max_retries {
111 return None;
112 }
113 match result {
114 Err(_) => idempotent.then(|| self.backoff(attempt)),
116 Ok(resp) => {
117 let status = resp.status().as_u16();
118 if !self.codes.contains(&status) {
119 return None;
120 }
121 if !idempotent && status != 429 {
123 return None;
124 }
125 Some(
126 retry_after(resp)
127 .map(|d| d.min(self.policy.max_delay))
128 .unwrap_or_else(|| self.backoff(attempt)),
129 )
130 }
131 }
132 }
133}
134
135#[async_trait::async_trait]
136impl reqwest_middleware::Middleware for RetryMiddleware {
137 async fn handle(
138 &self,
139 req: Request,
140 extensions: &mut http::Extensions,
141 next: Next<'_>,
142 ) -> reqwest_middleware::Result<Response> {
143 if self.policy.max_retries == 0 {
144 return next.run(req, extensions).await;
145 }
146 let idempotent = idempotent(req.method());
147 let mut attempt: u32 = 0;
148 loop {
149 let Some(attempt_req) = req.try_clone() else {
151 return next.run(req, extensions).await;
152 };
153 let result = next.clone().run(attempt_req, extensions).await;
154 let Some(wait) = self.retry_wait(attempt, idempotent, &result) else {
155 return result;
156 };
157 drop(result);
159 tokio::time::sleep(wait).await;
160 attempt += 1;
161 }
162 }
163}
164
165pub(crate) async fn with_retry<F, Fut, R, E>(
171 _policy: &RetryPolicy,
172 _idempotent: bool,
173 mut f: F,
174) -> Result<R, E>
175where
176 F: FnMut() -> Fut,
177 Fut: std::future::Future<Output = Result<R, E>>,
178{
179 f().await
180}