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().map(|server| {
if server > MAX_SERVER_DELAY {
tracing::warn!(
recommended = ?server,
cap = ?MAX_SERVER_DELAY,
"the participant recommended a retry delay above the cap; capping it"
);
}
server.min(MAX_SERVER_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;
Duration::try_from_secs_f64(backoff.as_secs_f64() * factor).unwrap_or(backoff)
}
pub const MAX_SERVER_DELAY: Duration = Duration::from_secs(60);
#[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(start_paused = true)]
async fn a_server_recommended_delay_is_capped() {
let status = {
use tonic_types::{ErrorDetails, StatusExt as _};
let mut details = ErrorDetails::new();
details.set_retry_info(Some(Duration::from_secs(365 * 86_400)));
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() >= MAX_SERVER_DELAY,
"{:?}",
started.elapsed()
);
assert!(
started.elapsed() <= MAX_SERVER_DELAY.mul_f64(1.5) + Duration::from_secs(1),
"slept {:?}",
started.elapsed()
);
}
#[tokio::test(start_paused = true)]
async fn an_absurd_json_retry_info_neither_panics_nor_hangs() {
for body in [
r#"{"errorCategory":1,"retryInfo":"1.8e19 seconds"}"#,
r#"{"errorCategory":1,"retryInfo":"1.3e19 seconds"}"#,
r#"{"cause":"busy","retryInfo":"18000000000000000000 seconds"}"#,
] {
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::Http {
status: 503,
body: body.to_string(),
url: None,
};
async move { if n == 1 { Err(err) } else { Ok(n) } }
})
.await;
assert_eq!(result.unwrap(), 2, "{body}");
assert!(
started.elapsed() <= MAX_SERVER_DELAY.mul_f64(1.5) + Duration::from_secs(1),
"{body}: slept {:?}",
started.elapsed()
);
}
}
#[test]
fn jitter_never_panics() {
for backoff in [Duration::ZERO, Duration::from_millis(1), Duration::MAX] {
let jittered = with_jitter(backoff);
assert!(jittered <= Duration::MAX);
}
}
#[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);
}
}