use std::collections::HashSet;
use std::time::Duration;
use rand::Rng;
use reqwest::{Method, Request, Response};
use reqwest_middleware::Next;
#[derive(Clone, Debug)]
pub struct RetryPolicy {
pub max_retries: u32,
pub base_delay: Duration,
pub max_delay: Duration,
pub retry_status: Vec<u16>,
}
impl Default for RetryPolicy {
fn default() -> Self {
Self {
max_retries: 2,
base_delay: Duration::from_millis(250),
max_delay: Duration::from_secs(10),
retry_status: vec![429, 502, 503, 504],
}
}
}
impl RetryPolicy {
pub fn disabled() -> Self {
Self {
max_retries: 0,
..Self::default()
}
}
}
fn idempotent(method: &Method) -> bool {
matches!(
*method,
Method::GET | Method::HEAD | Method::PUT | Method::DELETE | Method::OPTIONS
)
}
fn retry_after(resp: &Response) -> Option<Duration> {
let header = resp
.headers()
.get(reqwest::header::RETRY_AFTER)?
.to_str()
.ok()?;
if let Ok(secs) = header.trim().parse::<u64>() {
return Some(Duration::from_secs(secs));
}
let when = httpdate::parse_http_date(header.trim()).ok()?;
Some(
when.duration_since(std::time::SystemTime::now())
.unwrap_or(Duration::ZERO),
)
}
pub(crate) struct RetryMiddleware {
policy: RetryPolicy,
codes: HashSet<u16>,
}
impl RetryMiddleware {
pub(crate) fn new(policy: RetryPolicy) -> Self {
let codes = policy.retry_status.iter().copied().collect();
RetryMiddleware { policy, codes }
}
fn backoff(&self, attempt: u32) -> Duration {
let factor = 1u32.checked_shl(attempt.min(16)).unwrap_or(u32::MAX);
let ceiling = self
.policy
.base_delay
.saturating_mul(factor)
.min(self.policy.max_delay);
let ceiling_ms = u64::try_from(ceiling.as_millis()).unwrap_or(u64::MAX);
Duration::from_millis(rand::rng().random_range(0..=ceiling_ms))
}
fn retry_wait(
&self,
attempt: u32,
idempotent: bool,
result: &reqwest_middleware::Result<Response>,
) -> Option<Duration> {
if attempt >= self.policy.max_retries {
return None;
}
match result {
Err(_) => idempotent.then(|| self.backoff(attempt)),
Ok(resp) => {
let status = resp.status().as_u16();
if !self.codes.contains(&status) {
return None;
}
if !idempotent && status != 429 {
return None;
}
Some(
retry_after(resp)
.map(|d| d.min(self.policy.max_delay))
.unwrap_or_else(|| self.backoff(attempt)),
)
}
}
}
}
#[async_trait::async_trait]
impl reqwest_middleware::Middleware for RetryMiddleware {
async fn handle(
&self,
req: Request,
extensions: &mut http::Extensions,
next: Next<'_>,
) -> reqwest_middleware::Result<Response> {
if self.policy.max_retries == 0 {
return next.run(req, extensions).await;
}
let idempotent = idempotent(req.method());
let mut attempt: u32 = 0;
loop {
let Some(attempt_req) = req.try_clone() else {
return next.run(req, extensions).await;
};
let result = next.clone().run(attempt_req, extensions).await;
let Some(wait) = self.retry_wait(attempt, idempotent, &result) else {
return result;
};
drop(result);
tokio::time::sleep(wait).await;
attempt += 1;
}
}
}
pub(crate) async fn with_retry<F, Fut, R, E>(
_policy: &RetryPolicy,
_idempotent: bool,
mut f: F,
) -> Result<R, E>
where
F: FnMut() -> Fut,
Fut: std::future::Future<Output = Result<R, E>>,
{
f().await
}