use num::checked_pow;
use std::iter::Iterator;
use std::time::{Duration, Instant};
pub trait BackoffSchedule: Iterator<Item = Duration> {}
impl<T> BackoffSchedule for T where T: Iterator<Item = Duration> {}
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub struct RetryWithBackoff {
attempt: usize,
initial_delay: Duration,
backoff_coefficient: u32,
max_attempt: Option<usize>,
max_delay: Option<Duration>,
expiration_time: Option<Instant>,
}
impl RetryWithBackoff {
pub fn default_setting() -> RetryWithBackoff {
RetryWithBackoff {
attempt: 0,
initial_delay: Duration::from_millis(1),
backoff_coefficient: 10,
max_delay: Some(Duration::from_millis(10000)),
max_attempt: None,
expiration_time: None,
}
}
pub fn initial_delay(mut self, initial_delay: Duration) -> RetryWithBackoff {
self.initial_delay = initial_delay;
self
}
pub fn backoff_coefficient(mut self, backoff_coefficient: u32) -> RetryWithBackoff {
self.backoff_coefficient = backoff_coefficient;
self
}
pub fn max_attempt(mut self, attempt: usize) -> RetryWithBackoff {
self.max_attempt = Some(attempt);
self
}
pub fn max_delay(mut self, duration: Duration) -> RetryWithBackoff {
self.max_delay = Some(duration);
self
}
pub fn expiration_time(mut self, time: Instant) -> RetryWithBackoff {
self.expiration_time = Some(time);
self
}
#[deprecated(since = "0.4.0", note = "please use `initial_delay` instead")]
pub fn from_millis(base: u64) -> RetryWithBackoff {
RetryWithBackoff {
attempt: 0,
initial_delay: Duration::from_millis(base),
backoff_coefficient: base as u32,
max_delay: None,
max_attempt: None,
expiration_time: None,
}
}
#[deprecated(since = "0.4.0", note = "please use `max_attempt` instead")]
pub fn max_tries(self, tries: i32) -> std::iter::Take<RetryWithBackoff> {
self.take(tries as usize)
}
}
impl Iterator for RetryWithBackoff {
type Item = Duration;
fn next(&mut self) -> Option<Duration> {
if let Some(expiration_time) = self.expiration_time {
if expiration_time <= Instant::now() {
return None;
}
}
if let Some(max_attempt) = self.max_attempt {
if self.attempt >= max_attempt {
return None;
}
}
self.attempt += 1;
let coefficient = checked_pow(self.backoff_coefficient, self.attempt - 1);
let delay = coefficient.and_then(|coefficient| self.initial_delay.checked_mul(coefficient));
if delay.is_some() && self.max_delay.is_some() {
if delay < self.max_delay {
delay
} else {
self.max_delay
}
} else {
delay.or(self.max_delay)
}
}
}
#[test]
fn test_uses_default_setting() {
let mut s = RetryWithBackoff::default_setting();
assert_eq!(s.next(), Some(Duration::from_millis(1)));
assert_eq!(s.next(), Some(Duration::from_millis(10)));
assert_eq!(s.next(), Some(Duration::from_millis(100)));
assert_eq!(s.next(), Some(Duration::from_millis(1000)));
}
#[test]
fn test_returns_some_exponential_base_10() {
let mut s = RetryWithBackoff::default_setting().initial_delay(Duration::from_millis(10));
assert_eq!(s.next(), Some(Duration::from_millis(10)));
assert_eq!(s.next(), Some(Duration::from_millis(100)));
assert_eq!(s.next(), Some(Duration::from_millis(1000)));
}
#[test]
#[allow(deprecated)]
fn test_returns_with_finite_retries() {
let mut s = RetryWithBackoff::from_millis(10).max_tries(5);
assert_eq!(s.next(), Some(Duration::from_millis(10)));
assert_eq!(s.next(), Some(Duration::from_millis(100)));
assert_eq!(s.next(), Some(Duration::from_millis(1000)));
assert_eq!(s.next(), Some(Duration::from_millis(10000)));
assert_eq!(s.next(), Some(Duration::from_millis(100000)));
assert_eq!(s.next(), None);
}
#[test]
fn test_returns_some_exponential_base_2() {
let mut s = RetryWithBackoff::default_setting()
.initial_delay(Duration::from_millis(2))
.backoff_coefficient(2);
assert_eq!(s.next(), Some(Duration::from_millis(2)));
assert_eq!(s.next(), Some(Duration::from_millis(4)));
assert_eq!(s.next(), Some(Duration::from_millis(8)));
}
#[test]
fn stops_increasing_at_max_delay() {
let mut s = RetryWithBackoff::default_setting()
.initial_delay(Duration::from_millis(2))
.max_delay(Duration::from_millis(4));
assert_eq!(s.next(), Some(Duration::from_millis(2)));
assert_eq!(s.next(), Some(Duration::from_millis(4)));
assert_eq!(s.next(), Some(Duration::from_millis(4)));
}
#[test]
fn returns_max_when_max_less_than_base() {
let mut s = RetryWithBackoff::default_setting()
.initial_delay(Duration::from_millis(20))
.max_delay(Duration::from_millis(10));
assert_eq!(s.next(), Some(Duration::from_millis(10)));
assert_eq!(s.next(), Some(Duration::from_millis(10)));
}
#[cfg(test)]
mod tests {
use std::{ops::Add, thread};
use super::*;
#[test]
fn test_without_max_attempt_without_max_delay_without_expiration_time() {
let mut s = RetryWithBackoff::default_setting()
.initial_delay(Duration::from_millis(1))
.backoff_coefficient(2);
assert_eq!(s.next(), Some(Duration::from_millis(1)));
assert_eq!(s.next(), Some(Duration::from_millis(2)));
assert_eq!(s.next(), Some(Duration::from_millis(4)));
assert_eq!(s.next(), Some(Duration::from_millis(8)));
}
#[test]
fn test_with_max_attempt_without_max_delay_without_expiration_time() {
let mut s = RetryWithBackoff::default_setting()
.initial_delay(Duration::from_millis(1))
.backoff_coefficient(2)
.max_attempt(4);
assert_eq!(s.next(), Some(Duration::from_millis(1)));
assert_eq!(s.next(), Some(Duration::from_millis(2)));
assert_eq!(s.next(), Some(Duration::from_millis(4)));
assert_eq!(s.next(), Some(Duration::from_millis(8)));
assert_eq!(s.next(), None);
}
#[test]
fn test_without_max_attempt_without_max_delay_with_expiration_time() {
let mut s = RetryWithBackoff::default_setting()
.initial_delay(Duration::from_millis(1))
.backoff_coefficient(2)
.max_delay(Duration::from_millis(12));
assert_eq!(s.next(), Some(Duration::from_millis(1)));
assert_eq!(s.next(), Some(Duration::from_millis(2)));
assert_eq!(s.next(), Some(Duration::from_millis(4)));
assert_eq!(s.next(), Some(Duration::from_millis(8)));
assert_eq!(s.next(), Some(Duration::from_millis(12)));
assert_eq!(s.next(), Some(Duration::from_millis(12)));
}
#[test]
fn test_without_max_attempt_with_max_delay_without_expiration_time() {
let sleep_duration = Duration::from_millis(10);
let mut s = RetryWithBackoff::default_setting()
.initial_delay(Duration::from_millis(1))
.backoff_coefficient(2)
.expiration_time(Instant::now().add(sleep_duration));
assert_eq!(s.next(), Some(Duration::from_millis(1)));
assert_eq!(s.next(), Some(Duration::from_millis(2)));
assert_eq!(s.next(), Some(Duration::from_millis(4)));
assert_eq!(s.next(), Some(Duration::from_millis(8)));
thread::sleep(sleep_duration);
assert_eq!(s.next(), None);
}
#[test]
fn test_with_max_attempt_with_max_delay_without_expiration_time() {
let mut s = RetryWithBackoff::default_setting()
.initial_delay(Duration::from_millis(1))
.backoff_coefficient(2)
.max_attempt(8)
.max_delay(Duration::from_millis(12));
assert_eq!(s.next(), Some(Duration::from_millis(1)));
assert_eq!(s.next(), Some(Duration::from_millis(2)));
assert_eq!(s.next(), Some(Duration::from_millis(4)));
assert_eq!(s.next(), Some(Duration::from_millis(8)));
assert_eq!(s.next(), Some(Duration::from_millis(12)));
assert_eq!(s.next(), Some(Duration::from_millis(12)));
assert_eq!(s.next(), Some(Duration::from_millis(12)));
assert_eq!(s.next(), Some(Duration::from_millis(12)));
assert_eq!(s.next(), None);
}
}