use std::time::Duration;
use serde::{Deserialize, Serialize};
use crate::core::{Digest, EffectKey, RunId};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct RetryPolicy {
pub max_attempts: u32,
pub initial_backoff: Duration,
pub max_backoff: Duration,
pub multiplier: u32,
pub jitter: bool,
pub max_advice: Duration,
}
impl Default for RetryPolicy {
fn default() -> Self {
Self {
max_attempts: 3,
initial_backoff: Duration::from_millis(100),
max_backoff: Duration::from_secs(10),
multiplier: 2,
jitter: true,
max_advice: Duration::from_secs(60),
}
}
}
impl RetryPolicy {
#[must_use]
pub const fn never() -> Self {
Self {
max_attempts: 1,
initial_backoff: Duration::ZERO,
max_backoff: Duration::ZERO,
multiplier: 1,
jitter: false,
max_advice: Duration::ZERO,
}
}
#[must_use]
pub fn attempts(n: u32) -> Self {
Self {
max_attempts: n.max(1),
..Self::default()
}
}
#[must_use]
pub fn with_backoff(mut self, initial: Duration, max: Duration) -> Self {
self.initial_backoff = initial;
self.max_backoff = max;
self
}
#[must_use]
pub fn without_jitter(mut self) -> Self {
self.jitter = false;
self
}
#[must_use]
pub const fn wait_at_most(mut self, advice: Duration) -> Self {
self.max_advice = advice;
self
}
#[must_use]
pub fn wait_before(
&self,
run: RunId,
key: EffectKey,
attempt: u32,
advice: Option<Duration>,
) -> Duration {
if attempt <= 1 {
return Duration::ZERO;
}
match advice {
Some(named) if !named.is_zero() => named.min(self.max_advice),
_ => self.backoff(run, key, attempt),
}
}
#[must_use]
pub fn permits(&self, attempt: u32) -> bool {
attempt < self.max_attempts
}
#[must_use]
pub fn backoff(&self, run: RunId, key: EffectKey, attempt: u32) -> Duration {
if attempt <= 1 {
return Duration::ZERO;
}
let steps = attempt - 2;
let mut delay = self.initial_backoff;
for _ in 0..steps {
delay = delay
.saturating_mul(self.multiplier.max(1))
.min(self.max_backoff);
if delay >= self.max_backoff {
break;
}
}
let delay = delay.min(self.max_backoff);
if !self.jitter || delay.is_zero() {
return delay;
}
let mut seed = Vec::with_capacity(64);
seed.extend_from_slice(run.to_string().as_bytes());
seed.extend_from_slice(&key.to_hex().into_bytes());
seed.extend_from_slice(&attempt.to_be_bytes());
let digest = Digest::of(&seed);
let spread = u64::from_be_bytes(digest.as_bytes()[..8].try_into().expect("8 bytes"));
let half = delay.as_nanos() / 2;
let extra = (half.saturating_mul(u128::from(spread))) / u128::from(u64::MAX);
Duration::from_nanos(u64::try_from(half + extra).unwrap_or(u64::MAX))
}
}
#[must_use]
pub fn retry_after_seconds(value: &str) -> Option<u64> {
let seconds = value.trim().parse::<u64>().ok()?;
(seconds > 0).then_some(seconds)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::StepId;
fn key() -> EffectKey {
EffectKey::derive(
StepId(1),
crate::core::Phase::Forward,
0,
1,
"test.effect",
b"{}",
)
}
#[test]
fn the_first_attempt_never_waits() {
assert_eq!(
RetryPolicy::default().backoff(RunId::generate(), key(), 1),
Duration::ZERO
);
}
#[test]
fn without_jitter_makes_the_schedule_reproducible() {
let jittered = RetryPolicy::default();
let fixed = RetryPolicy::default().without_jitter();
assert!(jittered.jitter, "the default is jittered");
assert!(!fixed.jitter);
assert_eq!(
fixed.max_attempts, jittered.max_attempts,
"turning jitter off must not move the attempt ceiling"
);
assert_eq!(fixed.initial_backoff, jittered.initial_backoff);
assert_eq!(fixed.max_backoff, jittered.max_backoff);
let key = key();
let (a, b) = (RunId::generate(), RunId::generate());
assert_eq!(fixed.backoff(a, key, 3), fixed.backoff(b, key, 3));
}
#[test]
fn backoff_grows_and_then_stops_at_the_ceiling() {
let p = RetryPolicy {
max_attempts: 10,
initial_backoff: Duration::from_millis(100),
max_backoff: Duration::from_millis(800),
multiplier: 2,
jitter: false,
..RetryPolicy::default()
};
let run = RunId::generate();
let at = |n| p.backoff(run, key(), n);
assert_eq!(at(2), Duration::from_millis(100));
assert_eq!(at(3), Duration::from_millis(200));
assert_eq!(at(4), Duration::from_millis(400));
assert_eq!(at(5), Duration::from_millis(800));
assert_eq!(at(6), Duration::from_millis(800), "ceiling holds");
assert_eq!(at(50), Duration::from_millis(800), "and keeps holding");
}
#[test]
fn an_absurd_schedule_saturates_instead_of_panicking() {
let p = RetryPolicy {
max_attempts: u32::MAX,
initial_backoff: Duration::from_secs(1),
max_backoff: Duration::from_mins(1),
multiplier: u32::MAX,
jitter: true,
..RetryPolicy::default()
};
assert!(p.backoff(RunId::generate(), key(), u32::MAX) <= Duration::from_mins(1));
}
#[test]
fn jitter_stays_within_half_the_delay() {
let p = RetryPolicy {
jitter: true,
..RetryPolicy::default()
};
let run = RunId::generate();
for attempt in 2..8 {
let d = p.backoff(run, key(), attempt);
let plain = RetryPolicy { jitter: false, ..p }.backoff(run, key(), attempt);
assert!(
d >= plain / 2 && d <= plain,
"attempt {attempt}: {d:?} outside [{:?}, {plain:?}]",
plain / 2
);
}
}
#[test]
fn jitter_decorrelates_runs_but_repeats_for_one_run() {
let (a, b) = (RunId::generate(), RunId::generate());
let p = RetryPolicy::default();
assert_ne!(
p.backoff(a, key(), 3),
p.backoff(b, key(), 3),
"two runs retrying the same effect must not reconverge"
);
assert_eq!(
p.backoff(a, key(), 3),
p.backoff(a, key(), 3),
"and the schedule must be a pure function, not a draw"
);
}
#[test]
fn never_permits_no_second_attempt() {
assert!(!RetryPolicy::never().permits(1));
assert!(RetryPolicy::attempts(3).permits(1));
assert!(RetryPolicy::attempts(3).permits(2));
assert!(!RetryPolicy::attempts(3).permits(3));
}
}