use core::ops::Add;
use crate::compat::Duration;
use crate::state::RetryState;
use super::Jittered;
use super::Wait;
use super::strategies::{WaitExponential, WaitFixed, WaitLinear};
#[derive(Debug, Clone)]
pub struct WaitCapped<W> {
pub(super) inner: W,
pub(super) max: Duration,
}
impl<W: Wait> Wait for WaitCapped<W> {
fn next_wait(&self, state: &RetryState) -> Duration {
self.inner.next_wait(state).min(self.max)
}
}
#[derive(Debug, Clone)]
pub struct WaitCombine<A, B> {
left: A,
right: B,
}
impl<A, B> WaitCombine<A, B> {
#[must_use]
pub fn new(left: A, right: B) -> Self {
Self { left, right }
}
}
impl<A: Wait, B: Wait> Wait for WaitCombine<A, B> {
fn next_wait(&self, state: &RetryState) -> Duration {
let left = self.left.next_wait(state);
let right = self.right.next_wait(state);
left.saturating_add(right)
}
}
#[derive(Debug, Clone)]
pub struct WaitChain<A, B> {
first: A,
second: B,
after: u32,
}
impl<A, B> WaitChain<A, B> {
#[must_use]
pub fn new(first: A, second: B, after: u32) -> Self {
Self {
first,
second,
after,
}
}
}
impl<A: Wait, B: Wait> Wait for WaitChain<A, B> {
fn next_wait(&self, state: &RetryState) -> Duration {
if state.attempt <= self.after {
self.first.next_wait(state)
} else {
self.second.next_wait(state)
}
}
}
impl<W> WaitCapped<W> {
#[must_use]
pub fn jitter(self, max_jitter: Duration) -> WaitCapped<Jittered<W>> {
let WaitCapped { inner, max } = self;
WaitCapped {
inner: Jittered::additive(inner, max_jitter),
max,
}
}
}
macro_rules! impl_wait_add {
($ty:ty $(, $param:ident)*) => {
impl<$($param,)* Rhs> Add<Rhs> for $ty {
type Output = WaitCombine<Self, Rhs>;
fn add(self, rhs: Rhs) -> Self::Output {
WaitCombine::new(self, rhs)
}
}
};
}
impl_wait_add!(WaitFixed);
impl_wait_add!(WaitLinear);
impl_wait_add!(WaitExponential);
impl_wait_add!(WaitCombine<A, B>, A, B);
impl_wait_add!(WaitChain<A, B>, A, B);
impl_wait_add!(WaitCapped<W>, W);
impl_wait_add!(Jittered<W>, W);