use std::future::Future;
use std::time::Duration;
use crate::Result;
#[derive(Clone, Debug)]
#[non_exhaustive]
pub struct RetryConfig {
pub max_attempts: u32,
pub initial_backoff: Duration,
pub max_backoff: Duration,
pub attempt_timeout: Option<Duration>,
}
impl Default for RetryConfig {
fn default() -> Self {
Self {
max_attempts: 3,
initial_backoff: Duration::from_millis(100),
max_backoff: Duration::from_secs(5),
attempt_timeout: None,
}
}
}
impl RetryConfig {
#[must_use]
pub fn with_max_attempts(mut self, max_attempts: u32) -> Self {
self.max_attempts = max_attempts;
self
}
#[must_use]
pub fn with_initial_backoff(mut self, initial_backoff: Duration) -> Self {
self.initial_backoff = initial_backoff;
self
}
#[must_use]
pub fn with_max_backoff(mut self, max_backoff: Duration) -> Self {
self.max_backoff = max_backoff;
self
}
#[must_use]
pub fn with_attempt_timeout(mut self, attempt_timeout: Duration) -> Self {
self.attempt_timeout = Some(attempt_timeout);
self
}
}
pub async fn run_with_retry<T, F, Fut>(config: Option<&RetryConfig>, mut op: F) -> Result<T>
where
F: FnMut() -> Fut,
Fut: Future<Output = Result<T>>,
{
let Some(config) = config else {
return op().await;
};
let mut attempt = 1u32;
let mut backoff = config.initial_backoff;
loop {
let outcome = match config.attempt_timeout {
Some(timeout) => match tokio::time::timeout(timeout, op()).await {
Ok(result) => result,
Err(_) => Err(crate::Error::Timeout),
},
None => op().await,
};
match outcome {
Ok(value) => return Ok(value),
Err(err) if err.is_retriable() && attempt < config.max_attempts => {
let recommended = err.retry_delay();
let delay = recommended.map_or(backoff, |server| server.max(backoff));
let delay = match recommended {
Some(minimum) => with_jitter(delay).max(minimum),
None => with_jitter(delay),
};
tokio::time::sleep(delay).await;
backoff = (backoff * 2).min(config.max_backoff);
attempt += 1;
}
Err(err) => return Err(err),
}
}
}
fn with_jitter(backoff: Duration) -> Duration {
let nanos = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map_or(0, |d| d.subsec_nanos());
let factor = 0.5 + f64::from(nanos % 1024) / 1024.0;
backoff.mul_f64(factor)
}
#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
use super::*;
use crate::Error;
use std::cell::Cell;
fn fast() -> RetryConfig {
RetryConfig::default()
.with_initial_backoff(Duration::from_millis(1))
.with_max_backoff(Duration::from_millis(1))
}
#[tokio::test]
async fn retries_retriable_errors_then_succeeds() {
let calls = Cell::new(0);
let result: Result<u32> = run_with_retry(Some(&fast()), || {
calls.set(calls.get() + 1);
let n = calls.get();
async move { if n < 3 { Err(Error::Timeout) } else { Ok(n) } }
})
.await;
assert_eq!(result.unwrap(), 3);
assert_eq!(calls.get(), 3);
}
#[tokio::test(start_paused = true)]
async fn the_backoff_doubles_each_attempt_up_to_the_cap() {
let cfg = RetryConfig::default()
.with_max_attempts(6)
.with_initial_backoff(Duration::from_millis(10))
.with_max_backoff(Duration::from_millis(80));
let starts = std::cell::RefCell::new(Vec::new());
let calls = Cell::new(0u32);
let _: Result<u32> = run_with_retry(Some(&cfg), || {
starts.borrow_mut().push(tokio::time::Instant::now());
calls.set(calls.get() + 1);
async move { Err::<u32, _>(Error::Timeout) }
})
.await;
let t = starts.borrow();
let gaps: Vec<u64> = t
.windows(2)
.map(|w| u64::try_from((w[1] - w[0]).as_millis()).unwrap_or(u64::MAX))
.collect();
let expected = [10u64, 20, 40, 80, 80];
assert_eq!(gaps.len(), expected.len(), "one gap per retry: {gaps:?}");
for (g, e) in gaps.iter().zip(expected) {
assert!(
*g * 2 >= e && *g * 2 <= e * 3,
"gap {g}ms outside jitter band of expected {e}ms; gaps={gaps:?}"
);
}
}
#[tokio::test(start_paused = true)]
async fn it_stops_after_max_attempts_rather_than_looping() {
let cfg = RetryConfig::default()
.with_max_attempts(4)
.with_initial_backoff(Duration::from_millis(1))
.with_max_backoff(Duration::from_millis(1));
let calls = Cell::new(0u32);
let result: Result<u32> = run_with_retry(Some(&cfg), || {
calls.set(calls.get() + 1);
async move { Err::<u32, _>(Error::Timeout) }
})
.await;
assert!(result.is_err());
assert_eq!(
calls.get(),
4,
"exactly max_attempts calls, no more, no fewer"
);
}
#[tokio::test(start_paused = true)]
async fn a_server_recommended_delay_stretches_the_backoff() {
let delay = Duration::from_secs(3);
let status = {
use tonic_types::{ErrorDetails, StatusExt as _};
let mut details = ErrorDetails::new();
details.set_retry_info(Some(delay));
tonic::Status::with_error_details(tonic::Code::Unavailable, "wait", details)
};
let started = tokio::time::Instant::now();
let calls = Cell::new(0);
let result: Result<u32> = run_with_retry(Some(&fast()), || {
calls.set(calls.get() + 1);
let n = calls.get();
let err = Error::from(status.clone());
async move { if n == 1 { Err(err) } else { Ok(n) } }
})
.await;
assert_eq!(result.unwrap(), 2);
assert!(
started.elapsed() >= delay,
"the server's delay is a minimum, slept only {:?}",
started.elapsed()
);
}
#[tokio::test]
async fn does_not_retry_non_retriable_errors() {
let calls = Cell::new(0);
let result: Result<u32> = run_with_retry(Some(&fast()), || {
calls.set(calls.get() + 1);
async move { Err(Error::InvalidRequest("nope".to_string())) }
})
.await;
assert!(result.is_err());
assert_eq!(calls.get(), 1, "non-retriable errors are not retried");
}
#[tokio::test]
async fn gives_up_after_max_attempts() {
let calls = Cell::new(0);
let result: Result<u32> = run_with_retry(Some(&fast()), || {
calls.set(calls.get() + 1);
async move { Err(Error::Timeout) }
})
.await;
assert!(result.is_err());
assert_eq!(calls.get(), 3, "stops at max_attempts");
}
#[tokio::test]
async fn attempt_timeout_bounds_a_hung_attempt() {
let calls = Cell::new(0);
let config = fast()
.with_max_attempts(2)
.with_attempt_timeout(Duration::from_millis(5));
let result: Result<u32> = run_with_retry(Some(&config), || {
calls.set(calls.get() + 1);
async move {
tokio::time::sleep(Duration::from_secs(30)).await;
Ok(1)
}
})
.await;
assert!(result.is_err(), "every attempt times out");
assert_eq!(calls.get(), 2, "the hung attempt is bounded and retried");
}
#[tokio::test]
async fn without_config_runs_exactly_once() {
let calls = Cell::new(0);
let result: Result<u32> = run_with_retry(None, || {
calls.set(calls.get() + 1);
async move { Err(Error::Timeout) }
})
.await;
assert!(result.is_err());
assert_eq!(calls.get(), 1);
}
}