use std::{fmt::Debug, sync::Arc, time::Duration};
use queuey_core::Backoff;
const DEFAULT_BASE: Duration = Duration::from_millis(500);
const DEFAULT_MAX: Duration = Duration::from_secs(30);
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Rebuilding {
Connection,
Consumer,
}
#[derive(Clone, Copy, Debug)]
#[non_exhaustive]
pub struct Attempt<'a> {
pub failures: u32,
pub error: Option<&'a (dyn std::error::Error + 'static)>,
pub rebuilding: Rebuilding,
}
impl<'a> Attempt<'a> {
#[must_use]
pub fn first(rebuilding: Rebuilding) -> Self {
Self {
failures: 0,
error: None,
rebuilding,
}
}
#[must_use]
pub fn after(
rebuilding: Rebuilding,
failures: u32,
error: &'a (dyn std::error::Error + 'static),
) -> Self {
Self {
failures,
error: Some(error),
rebuilding,
}
}
}
pub trait ReconnectPolicy: Send + Sync + Debug {
fn next_delay(&self, attempt: Attempt<'_>) -> Option<Duration>;
}
#[derive(Clone, Debug)]
pub struct BackoffPolicy {
pub max_attempts: Option<u32>,
pub backoff: Backoff,
}
impl Default for BackoffPolicy {
fn default() -> Self {
Self {
max_attempts: None,
backoff: Backoff::Exponential {
base: DEFAULT_BASE,
factor: 2.0,
max: DEFAULT_MAX,
jitter: true,
},
}
}
}
impl BackoffPolicy {
#[must_use]
pub fn max_attempts(mut self, attempts: Option<u32>) -> Self {
self.max_attempts = attempts;
self
}
#[must_use]
pub fn backoff(mut self, backoff: Backoff) -> Self {
self.backoff = backoff;
self
}
}
impl ReconnectPolicy for BackoffPolicy {
fn next_delay(&self, attempt: Attempt<'_>) -> Option<Duration> {
if self.max_attempts.is_some_and(|max| attempt.failures >= max) {
return None;
}
if attempt.failures == 0 {
return Some(Duration::ZERO);
}
Some(self.backoff.delay_for(attempt.failures))
}
}
pub(crate) fn default_policy() -> Arc<dyn ReconnectPolicy> {
Arc::new(BackoffPolicy::default())
}
#[cfg(test)]
mod tests {
use super::*;
fn delay(policy: &dyn ReconnectPolicy, failures: u32) -> Option<Duration> {
let error = std::io::Error::other("broker went away");
let attempt = if failures == 0 {
Attempt::first(Rebuilding::Connection)
} else {
Attempt::after(Rebuilding::Connection, failures, &error)
};
policy.next_delay(attempt)
}
#[test]
fn the_default_retries_forever_with_a_jittered_exponential_backoff() {
let policy = BackoffPolicy::default();
assert_eq!(policy.max_attempts, None);
let Backoff::Exponential {
base,
factor,
max,
jitter,
} = policy.backoff
else {
panic!("the default must be exponential");
};
assert_eq!(base, DEFAULT_BASE);
assert_eq!(factor, 2.0);
assert_eq!(max, DEFAULT_MAX);
assert!(jitter, "a fleet must not reconnect in lockstep");
}
#[test]
fn the_first_attempt_after_a_drop_is_immediate() {
assert_eq!(delay(&BackoffPolicy::default(), 0), Some(Duration::ZERO));
}
#[test]
fn an_unlimited_policy_always_offers_another_attempt() {
let policy = BackoffPolicy::default();
for failures in [0, 1, 100, u32::MAX] {
assert!(delay(&policy, failures).is_some(), "failures = {failures}");
}
}
#[test]
fn a_bounded_policy_stops_at_the_limit() {
let policy = BackoffPolicy::default().max_attempts(Some(3));
assert!(delay(&policy, 0).is_some());
assert!(delay(&policy, 1).is_some());
assert!(delay(&policy, 2).is_some());
assert_eq!(delay(&policy, 3), None, "the third failure is the last");
assert_eq!(delay(&policy, 4), None);
}
#[test]
fn a_zero_attempt_policy_never_reconnects() {
let policy = BackoffPolicy::default().max_attempts(Some(0));
assert_eq!(delay(&policy, 0), None);
}
#[test]
fn the_delay_grows_and_is_capped() {
let policy = BackoffPolicy::default().backoff(Backoff::Exponential {
base: Duration::from_millis(500),
factor: 2.0,
max: Duration::from_secs(30),
jitter: false,
});
assert_eq!(delay(&policy, 1), Some(Duration::from_millis(500)));
assert_eq!(delay(&policy, 2), Some(Duration::from_secs(1)));
assert_eq!(delay(&policy, 3), Some(Duration::from_secs(2)));
assert_eq!(delay(&policy, 20), Some(Duration::from_secs(30)), "capped");
}
#[test]
fn a_jittered_delay_never_exceeds_the_cap() {
let policy = BackoffPolicy::default();
for failures in 1..40 {
assert!(delay(&policy, failures).expect("unlimited") <= DEFAULT_MAX);
}
}
#[derive(Debug)]
struct Picky;
impl ReconnectPolicy for Picky {
fn next_delay(&self, attempt: Attempt<'_>) -> Option<Duration> {
if attempt.rebuilding == Rebuilding::Consumer && attempt.failures >= 1 {
return None;
}
let fatal = attempt
.error
.is_some_and(|error| error.to_string().contains("ACCESS_REFUSED"));
if fatal {
return None;
}
Some(Duration::from_secs(1))
}
}
#[test]
fn a_custom_policy_can_decide_on_the_error_and_the_target() {
let refused = std::io::Error::other("ACCESS_REFUSED - login was refused");
assert_eq!(
Picky.next_delay(Attempt::after(Rebuilding::Connection, 1, &refused)),
None,
"credentials will not fix themselves"
);
let flaky = std::io::Error::other("connection reset by peer");
assert_eq!(
Picky.next_delay(Attempt::after(Rebuilding::Connection, 9, &flaky)),
Some(Duration::from_secs(1)),
"a network blip is retried regardless of the count"
);
assert_eq!(
Picky.next_delay(Attempt::after(Rebuilding::Consumer, 1, &flaky)),
None,
"but a consumer is given up on"
);
}
#[test]
fn a_policy_is_usable_behind_the_arc_the_options_store_it_in() {
let policy: Arc<dyn ReconnectPolicy> = Arc::new(Picky);
assert_eq!(
policy.next_delay(Attempt::first(Rebuilding::Connection)),
Some(Duration::from_secs(1))
);
assert!(format!("{policy:?}").contains("Picky"));
}
}