use std::{cell::RefCell, future::Future, time::Duration};
#[derive(Debug)]
pub struct RetryPolicy {
max_retries: u32,
base_delay: Duration,
attempts: RefCell<u32>,
}
impl RetryPolicy {
#[must_use]
pub fn new(max_retries: u32, base_delay: Duration) -> Self {
Self {
max_retries,
base_delay,
attempts: RefCell::new(0),
}
}
#[must_use]
pub fn attempts(&self) -> u32 {
*self.attempts.borrow()
}
#[must_use]
pub fn backoff_delay(&self, attempt: u32) -> Duration {
let factor = 5_u64.saturating_pow(attempt);
let secs = self.base_delay.as_secs().saturating_mul(factor);
Duration::from_secs(secs)
}
pub fn reset(&self) {
*self.attempts.borrow_mut() = 0;
}
}
impl Default for RetryPolicy {
fn default() -> Self {
Self::new(3, Duration::from_mins(5))
}
}
pub async fn with_retry<F, Fut, T, E>(policy: &RetryPolicy, f: F) -> Result<T, E>
where
F: Fn() -> Fut,
Fut: Future<Output = Result<T, E>>,
{
policy.reset();
let mut attempt: u32 = 0;
loop {
attempt += 1;
*policy.attempts.borrow_mut() = attempt;
match f().await {
Ok(value) => {
policy.reset();
return Ok(value);
}
Err(err) => {
if attempt > policy.max_retries {
return Err(err);
}
let delay = policy.backoff_delay(attempt - 1);
tokio::time::sleep(delay).await;
}
}
}
}
#[cfg(test)]
#[allow(
clippy::unwrap_used,
clippy::unwrap_in_result,
clippy::expect_used,
clippy::panic,
clippy::pedantic,
clippy::disallowed_methods,
clippy::indexing_slicing,
clippy::wildcard_enum_match_arm,
reason = "test module relaxes production lint strictness"
)]
mod tests {
use std::sync::atomic::{AtomicU32, Ordering};
use super::*;
#[test]
fn test_retry_policy_default_values() {
let policy = RetryPolicy::default();
assert_eq!(policy.max_retries, 3);
assert_eq!(policy.base_delay, Duration::from_mins(5));
}
#[test]
fn test_retry_policy_custom_values() {
let policy = RetryPolicy::new(5, Duration::from_secs(10));
assert_eq!(policy.max_retries, 5);
assert_eq!(policy.base_delay, Duration::from_secs(10));
}
#[test]
fn test_retry_policy_starts_at_zero_attempts() {
let policy = RetryPolicy::default();
assert_eq!(policy.attempts(), 0);
}
#[test]
fn test_retry_policy_backoff_delay_exponential() {
let policy = RetryPolicy::new(3, Duration::from_mins(5));
assert_eq!(policy.backoff_delay(0), Duration::from_mins(5));
assert_eq!(policy.backoff_delay(1), Duration::from_mins(25));
assert_eq!(policy.backoff_delay(2), Duration::from_mins(125));
}
#[test]
fn test_retry_policy_backoff_delay_protects_overflow() {
let policy = RetryPolicy::new(100, Duration::from_secs(u64::MAX));
let delay = policy.backoff_delay(10);
assert_eq!(delay, Duration::from_secs(u64::MAX));
}
#[test]
fn test_retry_policy_reset_clears_attempts() {
let policy = RetryPolicy::default();
*policy.attempts.borrow_mut() = 5;
assert_eq!(policy.attempts(), 5);
policy.reset();
assert_eq!(policy.attempts(), 0);
}
#[tokio::test]
async fn test_should_succeed_on_first_attempt() {
let policy = RetryPolicy::default();
let counter = AtomicU32::new(0);
let result = with_retry(&policy, || async {
counter.fetch_add(1, Ordering::SeqCst);
Ok::<_, String>("ok")
})
.await;
assert_eq!(result, Ok("ok"));
assert_eq!(counter.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn test_should_retry_up_to_max() {
let policy = RetryPolicy::new(2, Duration::from_millis(1));
let counter = AtomicU32::new(0);
let result = with_retry(&policy, || async {
let attempt = counter.fetch_add(1, Ordering::SeqCst);
if attempt < 2 {
Err::<(), _>("fail".to_string())
} else {
Ok::<_, String>(())
}
})
.await;
assert!(result.is_ok());
assert_eq!(counter.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn test_should_mark_failed_after_exhaustion() {
let policy = RetryPolicy::new(2, Duration::from_millis(1));
let result = with_retry(&policy, || async {
Err::<String, _>("permanent_fail".to_string())
})
.await;
assert_eq!(result, Err("permanent_fail".to_string()));
assert_eq!(policy.attempts(), 3);
}
#[tokio::test]
async fn test_should_reset_after_success() {
let policy = RetryPolicy::new(3, Duration::from_millis(1));
let call_count = AtomicU32::new(0);
let result = with_retry(&policy, || async {
let n = call_count.fetch_add(1, Ordering::SeqCst);
if n == 0 {
Err::<(), _>("first_fail".to_string())
} else {
Ok::<_, String>(())
}
})
.await;
assert!(result.is_ok());
assert_eq!(policy.attempts(), 0);
let result = with_retry(&policy, || async { Ok::<_, String>(()) }).await;
assert!(result.is_ok());
assert_eq!(policy.attempts(), 0);
}
#[tokio::test]
async fn test_should_not_retry_when_max_retries_zero() {
let policy = RetryPolicy::new(0, Duration::from_millis(1));
let counter = AtomicU32::new(0);
let result = with_retry(&policy, || async {
counter.fetch_add(1, Ordering::SeqCst);
Err::<String, _>("fail".to_string())
})
.await;
assert_eq!(result, Err("fail".to_string()));
assert_eq!(counter.load(Ordering::SeqCst), 1);
}
}