use std::fmt;
use async_trait::async_trait;
use http::Extensions;
use loop_retry::{IsRetryable, RetryConfig, RetryError, with_backoff};
use reqwest::{Request, Response, StatusCode};
use super::{Error, Middleware, Next, Result};
#[derive(Debug, Clone)]
pub struct RetryMiddleware {
config: RetryConfig,
}
impl RetryMiddleware {
pub fn new(config: RetryConfig) -> Self {
Self { config }
}
pub fn with_defaults() -> Self {
Self::new(RetryConfig::default())
}
}
pub fn is_retryable_status(status: StatusCode) -> bool {
status == StatusCode::REQUEST_TIMEOUT
|| status == StatusCode::TOO_MANY_REQUESTS
|| (status.is_server_error() && status != StatusCode::NOT_IMPLEMENTED)
}
pub fn is_retryable_error(err: &Error) -> bool {
match err {
Error::Reqwest(e) => e.is_timeout() || e.is_connect(),
Error::Middleware(_) => false,
}
}
enum AttemptFailure {
Transient(Error),
RetryableStatus(Response),
Permanent(Error),
}
impl fmt::Display for AttemptFailure {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Transient(e) | Self::Permanent(e) => write!(f, "{e}"),
Self::RetryableStatus(resp) => write!(f, "retryable status {}", resp.status()),
}
}
}
impl IsRetryable for AttemptFailure {
fn is_retryable(&self) -> bool {
match self {
Self::Transient(_) | Self::RetryableStatus(_) => true,
Self::Permanent(_) => false,
}
}
}
#[async_trait]
impl Middleware for RetryMiddleware {
async fn handle(
&self,
req: Request,
extensions: &mut Extensions,
next: Next<'_>,
) -> Result<Response> {
let mut slot = Some(req);
let shared = std::mem::take(extensions);
let outcome = with_backoff(&self.config, || {
let attempt_req = slot.take();
slot = attempt_req.as_ref().and_then(Request::try_clone);
let snapshot = shared.clone();
let next = next.clone();
async move {
let Some(req) = attempt_req else {
return Err(AttemptFailure::Permanent(Error::Middleware(
"request body is not cloneable; cannot retry".into(),
)));
};
let mut snapshot = snapshot;
match next.run(req, &mut snapshot).await {
Ok(resp) if is_retryable_status(resp.status()) => {
Err(AttemptFailure::RetryableStatus(resp))
}
Ok(resp) => Ok(resp),
Err(e) if is_retryable_error(&e) => Err(AttemptFailure::Transient(e)),
Err(e) => Err(AttemptFailure::Permanent(e)),
}
}
})
.await;
*extensions = shared;
match outcome {
Ok(resp) => Ok(resp),
Err(RetryError::FinalError(failure, _attempts)) => match failure {
AttemptFailure::RetryableStatus(resp) => Ok(resp),
AttemptFailure::Transient(e) | AttemptFailure::Permanent(e) => Err(e),
},
}
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used, clippy::expect_used)] use super::*;
#[test]
fn retryable_statuses_match_reqwest_retry_defaults() {
assert!(is_retryable_status(StatusCode::REQUEST_TIMEOUT));
assert!(is_retryable_status(StatusCode::TOO_MANY_REQUESTS));
assert!(is_retryable_status(StatusCode::INTERNAL_SERVER_ERROR));
assert!(is_retryable_status(StatusCode::BAD_GATEWAY));
assert!(is_retryable_status(StatusCode::SERVICE_UNAVAILABLE));
assert!(!is_retryable_status(StatusCode::NOT_IMPLEMENTED));
assert!(!is_retryable_status(StatusCode::NOT_FOUND));
assert!(!is_retryable_status(StatusCode::UNAUTHORIZED));
assert!(!is_retryable_status(StatusCode::OK));
}
#[test]
fn middleware_errors_are_never_retryable() {
let middleware_err = Error::Middleware("boom".into());
assert!(!is_retryable_error(&middleware_err));
}
#[test]
fn attempt_failure_classification_and_display() {
let permanent = AttemptFailure::Permanent(Error::Middleware("boom".into()));
assert!(permanent.to_string().contains("boom"));
assert!(!permanent.is_retryable());
assert!(AttemptFailure::Transient(Error::Middleware("x".into())).is_retryable());
}
}