use core::future::Future;
use core::time::Duration;
use crate::pipe::capabilities::{Clock, Retryable};
use crate::pipe::primitives::Pipe;
use crate::pipe::resilience::retry::{RetryAction, RetryController};
fn splitmix64(state: u64) -> u64 {
let state = state.wrapping_add(0x9e37_79b9_7f4a_7c15);
let state = (state ^ (state >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
let state = (state ^ (state >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
state ^ (state >> 31)
}
#[derive(Debug, Clone)]
pub struct Retry<Inner, Clk> {
inner: Inner,
controller: RetryController,
clock: Clk,
rand_seed: u64,
}
impl<Inner, Clk> Retry<Inner, Clk> {
#[must_use]
pub fn new(inner: Inner, controller: RetryController, clock: Clk, rand_seed: u64) -> Self {
Self {
inner,
controller,
clock,
rand_seed,
}
}
}
impl<Inner, Clk> Pipe for Retry<Inner, Clk>
where
Inner: Pipe,
Inner::In: Clone,
Inner::Out: Retryable,
Clk: Clock,
{
type In = Inner::In;
type Out = Inner::Out;
type Err = Inner::Err;
fn call(&self, input: Self::In) -> impl Future<Output = Result<Self::Out, Self::Err>> {
async move {
let mut prev_delay = Duration::ZERO;
let mut attempt = 0u32;
loop {
let outcome = self.inner.call(input.clone()).await;
let rand = splitmix64(self.rand_seed ^ u64::from(attempt));
let now_nanos = self.clock.now_nanos();
match self
.controller
.on_outcome(attempt, &outcome, now_nanos, rand, prev_delay)
{
RetryAction::Done | RetryAction::Exhausted => return outcome,
RetryAction::Retry { after } => {
self.clock.delay(after).await;
prev_delay = after;
attempt += 1;
}
}
}
}
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::*;
use crate::block_on;
use crate::pipe::clock::testing::RecordingClock;
use crate::pipe::resilience::backoff::{Backoff, Jitter};
use crate::pipe::resilience::deadline::Deadline;
use crate::pipe::retry_rules::RetryRules;
use std::sync::Arc;
use std::sync::atomic::{AtomicU32, Ordering};
#[derive(Debug, Clone, Copy, PartialEq)]
struct StatusOut(Option<u16>);
impl Retryable for StatusOut {
fn retry_status(&self) -> Option<u16> {
self.0
}
fn is_success(&self) -> bool {
self.0.is_none()
}
}
struct FlakyPipe {
calls: Arc<AtomicU32>,
succeed_at: u32,
}
impl Pipe for FlakyPipe {
type In = u32;
type Out = StatusOut;
type Err = &'static str;
fn call(&self, _input: u32) -> impl Future<Output = Result<StatusOut, &'static str>> {
let seen = self.calls.fetch_add(1, Ordering::Relaxed);
let succeed = seen >= self.succeed_at;
async move {
if succeed {
Ok(StatusOut(None))
} else {
Ok(StatusOut(Some(503)))
}
}
}
}
fn controller(max_attempts: u32, backoff: Backoff) -> RetryController {
RetryController {
rules: RetryRules::default(),
backoff,
jitter: Jitter::None,
max_attempts,
deadline: None,
}
}
#[test]
fn success_on_first_attempt_is_one_call_and_no_delays() {
let calls = Arc::new(AtomicU32::new(0));
let pipe = FlakyPipe {
calls: calls.clone(),
succeed_at: 0,
};
let clock = RecordingClock::new();
let retry = Retry::new(
pipe,
controller(3, Backoff::Constant(Duration::from_millis(50))),
clock,
0,
);
let result = block_on(Pipe::call(&retry, 0));
assert_eq!(result, Ok(StatusOut(None)));
assert_eq!(calls.load(Ordering::Relaxed), 1);
}
#[test]
fn two_failures_then_success_is_three_calls_and_two_delays() {
let calls = Arc::new(AtomicU32::new(0));
let pipe = FlakyPipe {
calls: calls.clone(),
succeed_at: 2,
};
let retry = Retry::new(
pipe,
controller(5, Backoff::Constant(Duration::from_millis(50))),
RecordingClock::new(),
0,
);
let result = block_on(Pipe::call(&retry, 0));
assert_eq!(result, Ok(StatusOut(None)));
assert_eq!(
calls.load(Ordering::Relaxed),
3,
"two failures then a success"
);
}
#[test]
fn exhausted_at_max_attempts_returns_last_outcome() {
let calls = Arc::new(AtomicU32::new(0));
let pipe = FlakyPipe {
calls: calls.clone(),
succeed_at: u32::MAX,
};
let retry = Retry::new(
pipe,
controller(3, Backoff::Constant(Duration::from_millis(50))),
RecordingClock::new(),
0,
);
let result = block_on(Pipe::call(&retry, 0));
assert_eq!(
result,
Ok(StatusOut(Some(503))),
"final retryable outcome propagates"
);
assert_eq!(
calls.load(Ordering::Relaxed),
3,
"exactly max_attempts calls"
);
}
#[test]
fn past_deadline_from_clock_exhausts_before_any_retry() {
let calls = Arc::new(AtomicU32::new(0));
let pipe = FlakyPipe {
calls: calls.clone(),
succeed_at: u32::MAX,
};
let mut controller = controller(10, Backoff::Constant(Duration::from_millis(50)));
controller.deadline = Some(Deadline::new(0, Duration::from_secs(1)));
let retry = Retry::new(pipe, controller, RecordingClock::at(2_000_000_000), 0);
let result = block_on(Pipe::call(&retry, 0));
assert_eq!(result, Ok(StatusOut(Some(503))));
assert_eq!(
calls.load(Ordering::Relaxed),
1,
"deadline cut it off before retrying"
);
}
#[test]
fn delays_follow_the_exponential_schedule() {
let calls = Arc::new(AtomicU32::new(0));
let pipe = FlakyPipe {
calls: calls.clone(),
succeed_at: u32::MAX,
};
let clock = RecordingClock::new();
let probe = clock.clone();
let backoff = Backoff::Exponential {
initial: Duration::from_millis(100),
factor: 2,
max: Duration::from_millis(2000),
};
let retry = Retry::new(pipe, controller(4, backoff), clock, 0);
let _ = block_on(Pipe::call(&retry, 0));
assert_eq!(
probe.delays(),
std::vec![
Duration::from_millis(100),
Duration::from_millis(200),
Duration::from_millis(400),
],
);
}
}