use std::time::Duration;
use crate::error::{Error, Result};
const DEFAULT_MAX_ATTEMPTS: u32 = 3;
const DEFAULT_INITIAL_BACKOFF: Duration = Duration::from_secs(1);
const DEFAULT_MAX_BACKOFF: Duration = Duration::from_secs(30);
const DEFAULT_MULTIPLIER: f64 = 2.0;
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct RetryPolicy {
pub max_attempts: u32,
pub initial_backoff: Duration,
pub max_backoff: Duration,
pub multiplier: f64,
}
impl Default for RetryPolicy {
fn default() -> Self {
Self {
max_attempts: DEFAULT_MAX_ATTEMPTS,
initial_backoff: DEFAULT_INITIAL_BACKOFF,
max_backoff: DEFAULT_MAX_BACKOFF,
multiplier: DEFAULT_MULTIPLIER,
}
}
}
impl RetryPolicy {
fn backoff(&self, attempt: u32) -> Duration {
let attempt = attempt.max(1);
let factor = self.multiplier.powi((attempt - 1) as i32);
let secs = self.initial_backoff.as_secs_f64() * factor;
let max_secs = self.max_backoff.as_secs_f64();
if !secs.is_finite() || secs >= max_secs {
return self.max_backoff;
}
if secs <= 0.0 {
return Duration::ZERO;
}
Duration::from_secs_f64(secs)
}
}
pub async fn retry<F, Fut, T>(policy: RetryPolicy, mut f: F) -> Result<T>
where
F: FnMut() -> Fut,
Fut: std::future::Future<Output = Result<T>>,
{
let max_attempts = policy.max_attempts.max(1);
let mut last_err: Option<Error> = None;
for attempt in 1..=max_attempts {
match f().await {
Ok(value) => return Ok(value),
Err(err) => {
if !err.is_retryable() {
return Err(err);
}
last_err = Some(err);
if attempt == max_attempts {
break;
}
let delay = policy.backoff(attempt);
tokio::time::sleep(delay).await;
}
}
}
Err(last_err.unwrap_or_else(|| Error::Other("retry: no attempts made".to_string())))
}
#[cfg(test)]
mod tests {
use super::*;
use std::cell::Cell;
fn fast_policy(max_attempts: u32) -> RetryPolicy {
RetryPolicy {
max_attempts,
initial_backoff: Duration::from_millis(0),
max_backoff: Duration::from_millis(0),
multiplier: 2.0,
}
}
#[tokio::test]
async fn test_succeeds_first_try() {
let calls = Cell::new(0);
let result: Result<i32> = retry(fast_policy(3), || {
calls.set(calls.get() + 1);
async { Ok(7) }
})
.await;
assert_eq!(result.unwrap(), 7);
assert_eq!(calls.get(), 1);
}
#[tokio::test]
async fn test_retries_then_succeeds() {
let calls = Cell::new(0);
let result: Result<i32> = retry(fast_policy(3), || {
let n = calls.get() + 1;
calls.set(n);
async move {
if n < 3 {
Err(Error::connection("transient"))
} else {
Ok(99)
}
}
})
.await;
assert_eq!(result.unwrap(), 99);
assert_eq!(calls.get(), 3);
}
#[tokio::test]
async fn test_gives_up_after_max_attempts() {
let calls = Cell::new(0);
let result: Result<i32> = retry(fast_policy(3), || {
calls.set(calls.get() + 1);
async { Err(Error::timeout()) }
})
.await;
assert!(result.is_err());
assert_eq!(calls.get(), 3, "should attempt exactly max_attempts times");
}
#[tokio::test]
async fn test_does_not_retry_non_retryable_auth() {
let calls = Cell::new(0);
let result: Result<i32> = retry(fast_policy(5), || {
calls.set(calls.get() + 1);
async { Err(Error::auth("bad credentials")) }
})
.await;
assert!(result.is_err());
assert_eq!(calls.get(), 1, "non-retryable error must not be retried");
}
#[tokio::test]
async fn test_does_not_retry_non_retryable_validation() {
let calls = Cell::new(0);
let result: Result<i32> = retry(fast_policy(5), || {
calls.set(calls.get() + 1);
async { Err(Error::validation("bad input")) }
})
.await;
assert!(result.is_err());
assert_eq!(calls.get(), 1);
}
#[tokio::test]
async fn test_max_attempts_zero_treated_as_one() {
let calls = Cell::new(0);
let result: Result<i32> = retry(fast_policy(0), || {
calls.set(calls.get() + 1);
async { Err(Error::connection("transient")) }
})
.await;
assert!(result.is_err());
assert_eq!(calls.get(), 1);
}
#[test]
fn test_default_policy_values() {
let p = RetryPolicy::default();
assert_eq!(p.max_attempts, 3);
assert_eq!(p.initial_backoff, Duration::from_secs(1));
assert_eq!(p.max_backoff, Duration::from_secs(30));
assert_eq!(p.multiplier, 2.0);
}
#[test]
fn test_backoff_exponential_and_clamped() {
let p = RetryPolicy {
max_attempts: 10,
initial_backoff: Duration::from_secs(1),
max_backoff: Duration::from_secs(10),
multiplier: 2.0,
};
assert_eq!(p.backoff(1), Duration::from_secs(1)); assert_eq!(p.backoff(2), Duration::from_secs(2)); assert_eq!(p.backoff(3), Duration::from_secs(4)); assert_eq!(p.backoff(4), Duration::from_secs(8)); assert_eq!(p.backoff(5), Duration::from_secs(10));
assert_eq!(p.backoff(100), Duration::from_secs(10));
}
#[test]
fn test_backoff_negative_multiplier_does_not_panic() {
let p = RetryPolicy {
max_attempts: 5,
initial_backoff: Duration::from_secs(1),
max_backoff: Duration::from_secs(10),
multiplier: -2.0,
};
assert_eq!(p.backoff(2), Duration::ZERO);
assert_eq!(p.backoff(1), Duration::from_secs(1));
for attempt in 1..=5 {
let _ = p.backoff(attempt);
}
}
}