use std::time::Duration;
use tokio::time::sleep;
use tracing::{warn, info};
use std::future::Future;
#[derive(Debug, Clone)]
pub struct RetryConfig {
pub initial_delay: Duration,
pub max_delay: Duration,
pub factor: f64,
pub max_retries: Option<usize>,
}
impl Default for RetryConfig {
fn default() -> Self {
Self::daemon()
}
}
impl RetryConfig {
#[must_use]
pub fn startup() -> Self {
Self {
max_retries: Some(5),
initial_delay: Duration::from_millis(200),
max_delay: Duration::from_secs(2),
factor: 2.0,
}
}
#[must_use]
pub fn daemon() -> Self {
Self {
max_retries: None, initial_delay: Duration::from_secs(1),
max_delay: Duration::from_secs(300), factor: 2.0,
}
}
#[must_use]
pub fn query() -> Self {
Self {
max_retries: Some(3),
initial_delay: Duration::from_millis(100),
max_delay: Duration::from_secs(2),
factor: 2.0,
}
}
#[must_use]
pub fn batch_write() -> Self {
Self {
max_retries: Some(5),
initial_delay: Duration::from_millis(50 + (rand::random::<u64>() % 100)), max_delay: Duration::from_secs(5),
factor: 2.0,
}
}
#[cfg(test)]
pub fn test() -> Self {
Self {
max_retries: Some(3),
initial_delay: Duration::from_millis(1),
max_delay: Duration::from_millis(10),
factor: 2.0,
}
}
}
pub async fn retry<F, Fut, T, E>(
operation_name: &str,
config: &RetryConfig,
mut operation: F,
) -> Result<T, E>
where
F: FnMut() -> Fut,
Fut: Future<Output = Result<T, E>>,
E: std::fmt::Display,
{
let mut delay = config.initial_delay;
let mut attempts = 0;
loop {
match operation().await {
Ok(val) => {
if attempts > 0 {
info!("Operation '{}' succeeded after {} retries", operation_name, attempts);
}
return Ok(val);
}
Err(err) => {
attempts += 1;
if let Some(max) = config.max_retries {
if attempts >= max {
return Err(err);
}
}
if config.max_retries.is_none() {
warn!(
"Operation '{}' failed (attempt {}, will retry forever): {}. Next retry in {:?}...",
operation_name, attempts, err, delay
);
} else {
warn!(
"Operation '{}' failed (attempt {}/{}): {}. Retrying in {:?}...",
operation_name, attempts, config.max_retries.unwrap(), err, delay
);
}
sleep(delay).await;
delay = (delay.mul_f64(config.factor)).min(config.max_delay);
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
#[derive(Debug)]
struct TestError(String);
impl std::fmt::Display for TestError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
#[tokio::test]
async fn test_retry_succeeds_first_try() {
let result: Result<i32, TestError> = retry(
"test_op",
&RetryConfig::test(),
|| async { Ok(42) },
).await;
assert_eq!(result.unwrap(), 42);
}
#[tokio::test]
async fn test_retry_succeeds_after_failures() {
let attempts = Arc::new(AtomicUsize::new(0));
let attempts_clone = attempts.clone();
let result: Result<i32, TestError> = retry(
"test_op",
&RetryConfig::test(),
|| {
let a = attempts_clone.clone();
async move {
let count = a.fetch_add(1, Ordering::SeqCst) + 1;
if count < 3 {
Err(TestError(format!("fail {}", count)))
} else {
Ok(42)
}
}
},
).await;
assert_eq!(result.unwrap(), 42);
assert_eq!(attempts.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn test_retry_exhausts_retries() {
let attempts = Arc::new(AtomicUsize::new(0));
let attempts_clone = attempts.clone();
let config = RetryConfig {
max_retries: Some(3),
initial_delay: Duration::from_millis(1),
max_delay: Duration::from_millis(10),
factor: 2.0,
};
let result: Result<i32, TestError> = retry(
"test_op",
&config,
|| {
let a = attempts_clone.clone();
async move {
a.fetch_add(1, Ordering::SeqCst);
Err(TestError("always fail".to_string()))
}
},
).await;
assert!(result.is_err());
assert!(result.unwrap_err().0.contains("always fail"));
assert_eq!(attempts.load(Ordering::SeqCst), 3);
}
#[test]
fn test_retry_config_presets() {
let startup = RetryConfig::startup();
assert!(startup.max_retries.is_some());
assert_eq!(startup.max_retries.unwrap(), 5);
let daemon = RetryConfig::daemon();
assert!(daemon.max_retries.is_none());
let query = RetryConfig::query();
assert!(query.max_retries.is_some());
assert_eq!(query.max_retries.unwrap(), 3);
}
#[test]
fn test_delay_exponential_backoff() {
let config = RetryConfig {
initial_delay: Duration::from_millis(100),
max_delay: Duration::from_secs(10),
factor: 2.0,
max_retries: Some(5),
};
let mut delay = config.initial_delay;
assert_eq!(delay, Duration::from_millis(100));
delay = (delay.mul_f64(config.factor)).min(config.max_delay);
assert_eq!(delay, Duration::from_millis(200));
delay = (delay.mul_f64(config.factor)).min(config.max_delay);
assert_eq!(delay, Duration::from_millis(400));
}
#[test]
fn test_delay_caps_at_max() {
let config = RetryConfig {
initial_delay: Duration::from_secs(1),
max_delay: Duration::from_secs(5),
factor: 10.0, max_retries: Some(5),
};
let mut delay = config.initial_delay;
delay = (delay.mul_f64(config.factor)).min(config.max_delay);
assert_eq!(delay, Duration::from_secs(5));
}
}