#![deny(missing_docs)]
use crate::error::{Error, Result};
use std::future::Future;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
#[derive(Debug, Clone, Copy)]
pub struct RetryPolicy {
pub max_retries: u32,
pub base_delay: Duration,
pub max_delay: Duration,
}
impl Default for RetryPolicy {
fn default() -> Self {
Self {
max_retries: 3,
base_delay: Duration::from_millis(250),
max_delay: Duration::from_secs(10),
}
}
}
impl RetryPolicy {
pub fn none() -> Self {
Self {
max_retries: 0,
..Self::default()
}
}
fn backoff(&self, attempt: u32) -> Duration {
let factor = 1u32.checked_shl(attempt).unwrap_or(u32::MAX);
let capped = self.base_delay.saturating_mul(factor).min(self.max_delay);
let jitter_ceiling_ms = (capped.as_millis() / 4) as u64;
if jitter_ceiling_ms == 0 {
return capped;
}
capped + Duration::from_millis(jitter_source() % (jitter_ceiling_ms + 1))
}
}
fn jitter_source() -> u64 {
static COUNTER: AtomicU64 = AtomicU64::new(0);
COUNTER.fetch_add(0x9E37_79B9_7F4A_7C15, Ordering::Relaxed)
}
pub(crate) async fn retry_with<T, F>(
policy: &RetryPolicy,
is_retryable: impl Fn(&Error) -> bool,
mut op: impl FnMut() -> F,
) -> Result<T>
where
F: Future<Output = Result<T>>,
{
let mut attempt = 0u32;
loop {
match op().await {
Ok(value) => return Ok(value),
Err(err) => {
if attempt >= policy.max_retries || !is_retryable(&err) {
return Err(err);
}
let delay = policy.backoff(attempt);
tracing::warn!(
"transient network error ({}); retry {}/{}",
err,
attempt + 1,
policy.max_retries
);
if !delay.is_zero() {
tokio::time::sleep(delay).await;
}
attempt += 1;
}
}
}
}
pub(crate) async fn retry_http<T, F>(policy: &RetryPolicy, op: impl FnMut() -> F) -> Result<T>
where
F: Future<Output = Result<T>>,
{
retry_with(policy, is_retryable_http, op).await
}
pub(crate) fn is_retryable_http(err: &Error) -> bool {
match err {
Error::Http(e) => {
if let Some(status) = e.status() {
status.is_server_error() || status == reqwest::StatusCode::TOO_MANY_REQUESTS
} else {
e.is_timeout() || e.is_connect() || e.is_request() || e.is_body()
}
}
_ => false,
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::cell::Cell;
fn zero_delay(max_retries: u32) -> RetryPolicy {
RetryPolicy {
max_retries,
base_delay: Duration::ZERO,
max_delay: Duration::ZERO,
}
}
#[tokio::test]
async fn retries_until_success() {
let calls = Cell::new(0u32);
let result: Result<u32> = retry_with(
&zero_delay(3),
|_| true,
|| async {
calls.set(calls.get() + 1);
if calls.get() < 3 {
Err(Error::Cache("transient".to_string()))
} else {
Ok(42)
}
},
)
.await;
assert_eq!(result.unwrap(), 42);
assert_eq!(calls.get(), 3);
}
#[tokio::test]
async fn gives_up_after_max_retries() {
let calls = Cell::new(0u32);
let result: Result<u32> = retry_with(
&zero_delay(2),
|_| true,
|| async {
calls.set(calls.get() + 1);
Err(Error::Cache("always".to_string()))
},
)
.await;
assert!(result.is_err());
assert_eq!(calls.get(), 3); }
#[tokio::test]
async fn does_not_retry_non_retryable() {
let calls = Cell::new(0u32);
let result: Result<u32> = retry_with(
&zero_delay(5),
|_| false,
|| async {
calls.set(calls.get() + 1);
Err(Error::Cache("terminal".to_string()))
},
)
.await;
assert!(result.is_err());
assert_eq!(calls.get(), 1);
}
#[test]
fn crate_terminal_errors_are_not_retryable() {
assert!(!is_retryable_http(&Error::SizeLimitExceeded {
context: "download body".to_string(),
limit: 1,
actual: None,
}));
assert!(!is_retryable_http(&Error::Cache("bad cache".to_string())));
assert!(!is_retryable_http(&Error::ChecksumMismatch {
path: "x".to_string(),
expected: 1,
actual: 2,
}));
}
}