use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
#[derive(Clone, Debug)]
pub struct Retry {
pub max_attempts: u32,
pub backoff: Backoff,
}
#[derive(Clone, Copy, Debug, Default)]
pub enum Backoff {
#[default]
None,
Fixed {
delay: Duration,
},
Exponential {
initial: Duration,
max: Duration,
jitter: f64,
},
}
impl Default for Retry {
fn default() -> Self {
Self {
max_attempts: 3,
backoff: Backoff::Fixed {
delay: Duration::from_secs(1),
},
}
}
}
impl Retry {
pub fn none() -> Self {
Self {
max_attempts: 0,
backoff: Backoff::None,
}
}
pub fn fixed(attempts: u32, delay: Duration) -> Self {
Self {
max_attempts: attempts,
backoff: Backoff::Fixed { delay },
}
}
pub fn exponential(attempts: u32) -> RetryBuilder {
RetryBuilder {
max_attempts: attempts,
..Default::default()
}
}
pub fn compute_delay(&self, attempt: u32) -> Duration {
match &self.backoff {
Backoff::None => Duration::ZERO,
Backoff::Fixed { delay } => *delay,
Backoff::Exponential {
initial,
max,
jitter,
} => {
let shift = attempt.min(31);
let multiplier = 1u32.checked_shl(shift).unwrap_or(u32::MAX);
let base = initial.saturating_mul(multiplier);
let capped = base.min(*max);
let factor = jitter_factor(*jitter);
Duration::from_secs_f64(capped.as_secs_f64() * factor)
}
}
}
}
#[derive(Debug, Clone)]
pub struct RetryBuilder {
max_attempts: u32,
initial: Duration,
max: Duration,
jitter: f64,
}
impl Default for RetryBuilder {
fn default() -> Self {
Self {
max_attempts: 3,
initial: Duration::from_secs(1),
max: Duration::from_secs(5),
jitter: 0.25,
}
}
}
impl RetryBuilder {
pub fn initial_delay(mut self, delay: Duration) -> Self {
self.initial = delay;
self
}
pub fn max_delay(mut self, delay: Duration) -> Self {
self.max = delay;
self
}
pub fn jitter(mut self, jitter: f64) -> Self {
self.jitter = jitter.clamp(0.0, 1.0);
self
}
pub fn build(self) -> Retry {
Retry {
max_attempts: self.max_attempts,
backoff: Backoff::Exponential {
initial: self.initial,
max: self.max,
jitter: self.jitter,
},
}
}
}
impl From<RetryBuilder> for Retry {
fn from(builder: RetryBuilder) -> Self {
builder.build()
}
}
static JITTER_COUNTER: AtomicU64 = AtomicU64::new(0);
fn jitter_factor(jitter: f64) -> f64 {
if jitter <= 0.0 {
return 1.0;
}
let counter = JITTER_COUNTER.fetch_add(1, Ordering::Relaxed);
let hash = counter.wrapping_mul(0x5851f42d4c957f2d);
let random = (hash >> 11) as f64 / ((1u64 << 53) as f64);
1.0 + (random - 0.5) * 2.0 * jitter
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_retry_none() {
let retry = Retry::none();
assert_eq!(retry.max_attempts, 0);
assert!(matches!(retry.backoff, Backoff::None));
}
#[test]
fn test_retry_default() {
let retry = Retry::default();
assert_eq!(retry.max_attempts, 3);
assert!(
matches!(retry.backoff, Backoff::Fixed { delay } if delay == Duration::from_secs(1))
);
}
#[test]
fn test_retry_fixed() {
let retry = Retry::fixed(5, Duration::from_millis(200));
assert_eq!(retry.max_attempts, 5);
assert!(
matches!(retry.backoff, Backoff::Fixed { delay } if delay == Duration::from_millis(200))
);
}
#[test]
fn test_retry_exponential_builder() {
let retry = Retry::exponential(4)
.initial_delay(Duration::from_millis(50))
.max_delay(Duration::from_secs(1))
.jitter(0.1)
.build();
assert_eq!(retry.max_attempts, 4);
match retry.backoff {
Backoff::Exponential {
initial,
max,
jitter,
} => {
assert_eq!(initial, Duration::from_millis(50));
assert_eq!(max, Duration::from_secs(1));
assert!((jitter - 0.1).abs() < f64::EPSILON);
}
_ => panic!("expected Exponential"),
}
}
#[test]
fn test_jitter_clamped() {
let retry = Retry::exponential(1).jitter(-0.5).build();
match retry.backoff {
Backoff::Exponential { jitter, .. } => assert_eq!(jitter, 0.0),
_ => panic!("expected Exponential"),
}
let retry = Retry::exponential(1).jitter(2.0).build();
match retry.backoff {
Backoff::Exponential { jitter, .. } => assert_eq!(jitter, 1.0),
_ => panic!("expected Exponential"),
}
}
#[test]
fn test_compute_delay_none() {
let retry = Retry::none();
assert_eq!(retry.compute_delay(0), Duration::ZERO);
assert_eq!(retry.compute_delay(5), Duration::ZERO);
}
#[test]
fn test_compute_delay_default() {
let retry = Retry::default();
assert_eq!(retry.compute_delay(0), Duration::from_secs(1));
assert_eq!(retry.compute_delay(5), Duration::from_secs(1));
}
#[test]
fn test_compute_delay_fixed() {
let retry = Retry::fixed(3, Duration::from_millis(100));
assert_eq!(retry.compute_delay(0), Duration::from_millis(100));
assert_eq!(retry.compute_delay(1), Duration::from_millis(100));
assert_eq!(retry.compute_delay(10), Duration::from_millis(100));
}
#[test]
fn test_compute_delay_exponential_no_jitter() {
let retry = Retry::exponential(5)
.initial_delay(Duration::from_millis(100))
.max_delay(Duration::from_secs(10))
.jitter(0.0)
.build();
assert_eq!(retry.compute_delay(0), Duration::from_millis(100));
assert_eq!(retry.compute_delay(1), Duration::from_millis(200));
assert_eq!(retry.compute_delay(2), Duration::from_millis(400));
assert_eq!(retry.compute_delay(3), Duration::from_millis(800));
}
#[test]
fn test_compute_delay_exponential_capped() {
let retry = Retry::exponential(10)
.initial_delay(Duration::from_millis(100))
.max_delay(Duration::from_millis(500))
.jitter(0.0)
.build();
assert_eq!(retry.compute_delay(0), Duration::from_millis(100));
assert_eq!(retry.compute_delay(1), Duration::from_millis(200));
assert_eq!(retry.compute_delay(2), Duration::from_millis(400));
assert_eq!(retry.compute_delay(3), Duration::from_millis(500));
assert_eq!(retry.compute_delay(10), Duration::from_millis(500));
}
#[test]
fn test_compute_delay_exponential_with_jitter() {
let retry = Retry::exponential(3)
.initial_delay(Duration::from_millis(100))
.max_delay(Duration::from_secs(1))
.jitter(0.25)
.build();
for _ in 0..10 {
let delay = retry.compute_delay(0);
let millis = delay.as_millis();
assert!((75..=125).contains(&millis), "delay was {}ms", millis);
}
}
#[test]
fn test_jitter_factor_range() {
for _ in 0..100 {
let factor = jitter_factor(0.5);
assert!((0.5..=1.5).contains(&factor), "factor was {}", factor);
}
}
#[test]
fn test_jitter_factor_zero() {
assert_eq!(jitter_factor(0.0), 1.0);
assert_eq!(jitter_factor(-0.1), 1.0);
}
#[test]
fn test_from_builder() {
let builder = Retry::exponential(2).initial_delay(Duration::from_millis(50));
let retry: Retry = builder.into();
assert_eq!(retry.max_attempts, 2);
}
}