use crate::{Error, Result};
use std::time::Duration;
#[derive(Debug, Clone)]
pub struct RetryPolicy {
pub max_attempts: u32,
pub initial_backoff_ms: u64,
pub max_backoff_ms: u64,
pub backoff_multiplier: f64,
}
impl RetryPolicy {
pub fn new(
max_attempts: u32,
initial_backoff_ms: u64,
max_backoff_ms: u64,
backoff_multiplier: f64,
) -> Self {
Self {
max_attempts,
initial_backoff_ms,
max_backoff_ms,
backoff_multiplier,
}
}
pub fn no_retry() -> Self {
Self {
max_attempts: 0,
initial_backoff_ms: 0,
max_backoff_ms: 0,
backoff_multiplier: 1.0,
}
}
pub fn fast() -> Self {
Self {
max_attempts: 3,
initial_backoff_ms: 10,
max_backoff_ms: 100,
backoff_multiplier: 2.0,
}
}
pub fn standard() -> Self {
Self {
max_attempts: 5,
initial_backoff_ms: 100,
max_backoff_ms: 5000,
backoff_multiplier: 2.0,
}
}
pub fn backoff_duration(&self, attempt: u32) -> Duration {
let backoff_ms = (self.initial_backoff_ms as f64
* self.backoff_multiplier.powi(attempt as i32))
.min(self.max_backoff_ms as f64) as u64;
Duration::from_millis(backoff_ms)
}
}
impl Default for RetryPolicy {
fn default() -> Self {
Self::standard()
}
}
pub fn retry_with_policy<F, T>(
policy: &RetryPolicy,
mut operation: F,
) -> Result<T>
where
F: FnMut() -> Result<T>,
{
let mut last_error = None;
match operation() {
Ok(result) => return Ok(result),
Err(e) => {
if !e.is_retryable() {
return Err(e);
}
last_error = Some(e);
}
}
for attempt in 0..policy.max_attempts {
let backoff = policy.backoff_duration(attempt);
std::thread::sleep(backoff);
match operation() {
Ok(result) => return Ok(result),
Err(e) => {
if !e.is_retryable() {
return Err(e);
}
last_error = Some(e);
}
}
}
Err(last_error.unwrap_or_else(|| {
Error::Internal("retry exhausted without error".to_string())
}))
}
pub fn retry<F, T>(operation: F) -> Result<T>
where
F: FnMut() -> Result<T>,
{
retry_with_policy(&RetryPolicy::default(), operation)
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::{Arc, Mutex};
#[test]
fn test_retry_policy_default() {
let policy = RetryPolicy::default();
assert_eq!(policy.max_attempts, 5);
assert_eq!(policy.initial_backoff_ms, 100);
assert_eq!(policy.max_backoff_ms, 5000);
assert_eq!(policy.backoff_multiplier, 2.0);
}
#[test]
fn test_retry_policy_fast() {
let policy = RetryPolicy::fast();
assert_eq!(policy.max_attempts, 3);
assert_eq!(policy.initial_backoff_ms, 10);
}
#[test]
fn test_retry_policy_no_retry() {
let policy = RetryPolicy::no_retry();
assert_eq!(policy.max_attempts, 0);
}
#[test]
fn test_backoff_duration_exponential() {
let policy = RetryPolicy::new(5, 100, 10000, 2.0);
assert_eq!(policy.backoff_duration(0).as_millis(), 100);
assert_eq!(policy.backoff_duration(1).as_millis(), 200);
assert_eq!(policy.backoff_duration(2).as_millis(), 400);
assert_eq!(policy.backoff_duration(3).as_millis(), 800);
}
#[test]
fn test_backoff_duration_respects_max() {
let policy = RetryPolicy::new(10, 100, 500, 2.0);
assert_eq!(policy.backoff_duration(5).as_millis(), 500);
assert_eq!(policy.backoff_duration(10).as_millis(), 500);
}
#[test]
fn test_retry_succeeds_immediately() {
let policy = RetryPolicy::fast();
let counter = Arc::new(Mutex::new(0));
let counter_clone = counter.clone();
let result = retry_with_policy(&policy, || {
let mut count = counter_clone.lock().unwrap();
*count += 1;
Ok::<i32, Error>(42)
});
assert!(result.is_ok());
assert_eq!(result.unwrap(), 42);
assert_eq!(*counter.lock().unwrap(), 1); }
#[test]
fn test_retry_succeeds_after_failures() {
let policy = RetryPolicy::fast();
let counter = Arc::new(Mutex::new(0));
let counter_clone = counter.clone();
let result = retry_with_policy(&policy, || {
let mut count = counter_clone.lock().unwrap();
*count += 1;
if *count < 3 {
Err(Error::Io(std::io::Error::new(
std::io::ErrorKind::TimedOut,
"timeout"
)))
} else {
Ok::<i32, Error>(42)
}
});
assert!(result.is_ok());
assert_eq!(result.unwrap(), 42);
assert_eq!(*counter.lock().unwrap(), 3); }
#[test]
fn test_retry_fails_after_max_attempts() {
let policy = RetryPolicy::new(2, 1, 10, 1.5); let counter = Arc::new(Mutex::new(0));
let counter_clone = counter.clone();
let result = retry_with_policy(&policy, || {
let mut count = counter_clone.lock().unwrap();
*count += 1;
Err::<i32, Error>(Error::Io(std::io::Error::new(
std::io::ErrorKind::TimedOut,
"timeout"
)))
});
assert!(result.is_err());
assert_eq!(*counter.lock().unwrap(), 3);
}
#[test]
fn test_retry_does_not_retry_non_retryable_error() {
let policy = RetryPolicy::fast();
let counter = Arc::new(Mutex::new(0));
let counter_clone = counter.clone();
let result = retry_with_policy(&policy, || {
let mut count = counter_clone.lock().unwrap();
*count += 1;
Err::<i32, Error>(Error::InvalidArgument("bad input".to_string()))
});
assert!(result.is_err());
assert_eq!(*counter.lock().unwrap(), 1);
match result {
Err(Error::InvalidArgument(_)) => (),
_ => panic!("Expected InvalidArgument error"),
}
}
#[test]
fn test_retry_no_retry_policy() {
let policy = RetryPolicy::no_retry();
let counter = Arc::new(Mutex::new(0));
let counter_clone = counter.clone();
let result = retry_with_policy(&policy, || {
let mut count = counter_clone.lock().unwrap();
*count += 1;
Err::<i32, Error>(Error::Io(std::io::Error::new(
std::io::ErrorKind::TimedOut,
"timeout"
)))
});
assert!(result.is_err());
assert_eq!(*counter.lock().unwrap(), 1); }
#[test]
fn test_retry_helper_function() {
let counter = Arc::new(Mutex::new(0));
let counter_clone = counter.clone();
let result = retry(|| {
let mut count = counter_clone.lock().unwrap();
*count += 1;
if *count < 2 {
Err(Error::Io(std::io::Error::new(
std::io::ErrorKind::TimedOut,
"timeout"
)))
} else {
Ok::<i32, Error>(100)
}
});
assert!(result.is_ok());
assert_eq!(result.unwrap(), 100);
}
}