use std::future::Future;
use std::time::Duration;
use tokio::time::Instant;
use tokio_util::sync::CancellationToken;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct RetryPolicy {
pub initial_delay: Duration,
pub max_delay: Duration,
pub budget: Duration,
}
pub(crate) const INTERACTIVE_RETRY: RetryPolicy = RetryPolicy {
initial_delay: Duration::from_millis(500),
max_delay: Duration::from_secs(5),
budget: Duration::from_secs(30),
};
pub(crate) const BACKGROUND_RETRY: RetryPolicy = RetryPolicy {
initial_delay: Duration::from_secs(2),
max_delay: Duration::from_secs(60),
budget: Duration::from_secs(5 * 60),
};
const MIN_ATTEMPT: Duration = Duration::from_secs(1);
pub(crate) async fn with_retry<T, E, F, Fut>(
mut attempt: F,
is_retryable: impl Fn(&E) -> bool,
policy: RetryPolicy,
cancel: &CancellationToken,
) -> Result<T, E>
where
F: FnMut(Duration) -> Fut,
Fut: Future<Output = Result<T, E>>,
{
let deadline = Instant::now() + policy.budget;
let mut delay = policy.initial_delay;
loop {
let error = match attempt(deadline.saturating_duration_since(Instant::now())).await {
Ok(value) => return Ok(value),
Err(error) => error,
};
let wait = delay / 2 + delay.mul_f64(rand::random::<f64>()) / 2;
if !is_retryable(&error)
|| cancel.is_cancelled()
|| Instant::now() + wait + MIN_ATTEMPT > deadline
{
return Err(error);
}
tokio::select! {
() = tokio::time::sleep(wait) => {}
() = cancel.cancelled() => return Err(error),
}
delay = (delay * 2).min(policy.max_delay);
}
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicU32, Ordering};
use super::*;
#[tokio::test(start_paused = true)]
async fn returns_the_first_success() {
let calls = AtomicU32::new(0);
let result: Result<u32, ()> = with_retry(
|_| async { Ok(calls.fetch_add(1, Ordering::SeqCst)) },
|_| true,
INTERACTIVE_RETRY,
&CancellationToken::new(),
)
.await;
assert_eq!(result, Ok(0));
}
#[tokio::test(start_paused = true)]
async fn passes_the_remaining_budget_and_stays_within_it() {
let started = Instant::now();
let mut remaining = Vec::new();
let result: Result<(), ()> = with_retry(
|left| {
remaining.push(left);
async { Err(()) }
},
|_| true,
INTERACTIVE_RETRY,
&CancellationToken::new(),
)
.await;
assert_eq!(result, Err(()));
assert!(remaining.len() > 3);
assert_eq!(remaining[0], INTERACTIVE_RETRY.budget);
assert!(remaining.windows(2).all(|w| w[1] < w[0]));
assert!(remaining.iter().all(|left| *left >= MIN_ATTEMPT));
assert!(started.elapsed() <= INTERACTIVE_RETRY.budget);
}
#[tokio::test(start_paused = true)]
async fn does_not_retry_a_final_error() {
let calls = AtomicU32::new(0);
let result: Result<(), u32> = with_retry(
|_| async {
calls.fetch_add(1, Ordering::SeqCst);
Err(1)
},
|_| false,
INTERACTIVE_RETRY,
&CancellationToken::new(),
)
.await;
assert_eq!(result, Err(1));
assert_eq!(calls.load(Ordering::SeqCst), 1);
}
#[tokio::test(start_paused = true)]
async fn stops_when_cancelled() {
let cancel = CancellationToken::new();
let calls = AtomicU32::new(0);
let retry = with_retry(
|_| async {
calls.fetch_add(1, Ordering::SeqCst);
Err::<(), _>(())
},
|_| true,
BACKGROUND_RETRY,
&cancel,
);
let cancel_soon = async {
tokio::time::sleep(Duration::from_millis(100)).await;
cancel.cancel();
};
let (result, ()) = tokio::join!(retry, cancel_soon);
assert_eq!(result, Err(()));
assert_eq!(calls.load(Ordering::SeqCst), 1);
let result: Result<(), ()> =
with_retry(|_| async { Err(()) }, |_| true, BACKGROUND_RETRY, &cancel).await;
assert_eq!(result, Err(()));
}
}