use std::{collections::VecDeque, pin::pin, sync::atomic::AtomicU8, time::Duration};
use nautilus_core::correctness::{check_in_range_inclusive_f64, check_predicate_true};
use rand::RngExt;
use crate::{dst, mode::ConnectionMode};
pub(crate) const RECONNECT_STABILITY_THRESHOLD: Duration = Duration::from_secs(10);
pub(crate) const RECONNECT_MIN_DELAY: Duration = Duration::from_secs(1);
pub(crate) const RECONNECT_MIN_DELAY_ATTEMPTS: usize = 3;
pub(crate) const RECONNECT_MIN_DELAY_WINDOW: Duration = Duration::from_mins(2);
#[derive(Debug, Default)]
pub(crate) struct ReconnectThrottle {
recent_attempts: VecDeque<dst::time::Instant>,
}
impl ReconnectThrottle {
pub(crate) fn gated_delay(&mut self, backoff_delay: Duration) -> Duration {
self.prune_expired();
if self.recent_attempts.len() >= RECONNECT_MIN_DELAY_ATTEMPTS {
backoff_delay.max(RECONNECT_MIN_DELAY)
} else {
backoff_delay
}
}
pub(crate) fn record_attempt(&mut self) {
self.prune_expired();
self.recent_attempts.push_back(dst::time::Instant::now());
}
fn prune_expired(&mut self) {
while self
.recent_attempts
.front()
.is_some_and(|oldest| oldest.elapsed() > RECONNECT_MIN_DELAY_WINDOW)
{
self.recent_attempts.pop_front();
}
}
}
#[derive(Clone, Debug)]
pub struct ExponentialBackoff {
delay_initial: Duration,
delay_max: Duration,
delay_current: Duration,
factor: f64,
jitter_ms: u64,
immediate_reconnect: bool,
immediate_reconnect_original: bool,
}
impl ExponentialBackoff {
pub fn new(
delay_initial: Duration,
delay_max: Duration,
factor: f64,
jitter_ms: u64,
immediate_first: bool,
) -> anyhow::Result<Self> {
check_predicate_true(!delay_initial.is_zero(), "delay_initial must be non-zero")?;
check_predicate_true(
delay_max >= delay_initial,
"delay_max must be >= delay_initial",
)?;
check_predicate_true(
delay_max.as_nanos() <= u128::from(u64::MAX),
"delay_max exceeds maximum representable duration (≈584 years)",
)?;
check_in_range_inclusive_f64(factor, 1.0, 100.0, "factor")?;
Ok(Self {
delay_initial,
delay_max,
delay_current: delay_initial,
factor,
jitter_ms,
immediate_reconnect: immediate_first,
immediate_reconnect_original: immediate_first,
})
}
pub fn next_duration(&mut self) -> Duration {
if self.immediate_reconnect && self.delay_current == self.delay_initial {
self.immediate_reconnect = false;
return Duration::ZERO;
}
let jitter = rand::rng().random_range(0..=self.jitter_ms);
let base = std::cmp::min(
self.delay_current,
self.delay_max
.saturating_sub(Duration::from_millis(self.jitter_ms)),
);
let delay_with_jitter = base + Duration::from_millis(jitter);
let floor = std::cmp::min(self.delay_initial, self.delay_max);
let clamped_delay = delay_with_jitter.clamp(floor, self.delay_max);
let current_nanos = self.delay_current.as_nanos() as u64;
let max_nanos = self.delay_max.as_nanos() as u64;
let next_nanos = (current_nanos as f64 * self.factor) as u64;
self.delay_current = Duration::from_nanos(next_nanos.min(max_nanos));
clamped_delay
}
pub const fn reset(&mut self) {
self.delay_current = self.delay_initial;
self.immediate_reconnect = self.immediate_reconnect_original;
}
#[must_use]
pub const fn current_delay(&self) -> Duration {
self.delay_current
}
}
pub(crate) async fn wait_reconnect_delay(
duration: Duration,
connection_mode: &AtomicU8,
state_notify: &tokio::sync::Notify,
) -> bool {
if duration.is_zero() {
return true;
}
tokio::select! {
biased;
() = dst::time::sleep(duration) => true,
() = async {
loop {
let mut notified = pin!(state_notify.notified());
notified.as_mut().enable();
if !ConnectionMode::from_atomic(connection_mode).is_reconnect() {
break;
}
notified.await;
}
} => false,
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use rstest::rstest;
use super::*;
#[rstest]
fn test_no_jitter_exponential_growth() {
let initial = Duration::from_millis(100);
let max = Duration::from_millis(1600);
let factor = 2.0;
let jitter = 0;
let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, false).unwrap();
let d1 = backoff.next_duration();
assert_eq!(d1, Duration::from_millis(100));
let d2 = backoff.next_duration();
assert_eq!(d2, Duration::from_millis(200));
let d3 = backoff.next_duration();
assert_eq!(d3, Duration::from_millis(400));
let d4 = backoff.next_duration();
assert_eq!(d4, Duration::from_millis(800));
let d5 = backoff.next_duration();
assert_eq!(d5, Duration::from_millis(1600));
let d6 = backoff.next_duration();
assert_eq!(d6, Duration::from_millis(1600));
}
#[rstest]
fn test_reset() {
let initial = Duration::from_millis(100);
let max = Duration::from_millis(1600);
let factor = 2.0;
let jitter = 0;
let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, false).unwrap();
let _ = backoff.next_duration(); backoff.reset();
let d = backoff.next_duration();
assert_eq!(d, Duration::from_millis(100));
}
#[rstest]
fn test_jitter_within_bounds() {
let initial = Duration::from_millis(100);
let max = Duration::from_secs(1);
let factor = 2.0;
let jitter = 50;
for _ in 0..10 {
let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, false).unwrap();
let base = backoff.delay_current;
let delay = backoff.next_duration();
let min_expected = base;
let max_expected = base + Duration::from_millis(jitter);
assert!(
delay >= min_expected,
"Delay {delay:?} is less than expected minimum {min_expected:?}"
);
assert!(
delay <= max_expected,
"Delay {delay:?} exceeds expected maximum {max_expected:?}"
);
}
}
#[rstest]
fn test_factor_less_than_two() {
let initial = Duration::from_millis(100);
let max = Duration::from_millis(200);
let factor = 1.5;
let jitter = 0;
let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, false).unwrap();
let d1 = backoff.next_duration();
assert_eq!(d1, Duration::from_millis(100));
let d2 = backoff.next_duration();
assert_eq!(d2, Duration::from_millis(150));
let d3 = backoff.next_duration();
assert_eq!(d3, Duration::from_millis(200));
let d4 = backoff.next_duration();
assert_eq!(d4, Duration::from_millis(200));
}
#[rstest]
fn test_max_delay_is_respected() {
let initial = Duration::from_millis(500);
let max = Duration::from_secs(1);
let factor = 3.0;
let jitter = 0;
let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, false).unwrap();
let d1 = backoff.next_duration();
assert_eq!(d1, Duration::from_millis(500));
let d2 = backoff.next_duration();
assert_eq!(d2, Duration::from_secs(1));
let d3 = backoff.next_duration();
assert_eq!(d3, Duration::from_secs(1));
}
#[rstest]
fn test_current_delay_getter() {
let initial = Duration::from_millis(100);
let max = Duration::from_millis(1600);
let factor = 2.0;
let jitter = 0;
let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, false).unwrap();
assert_eq!(backoff.current_delay(), initial);
let _ = backoff.next_duration();
assert_eq!(backoff.current_delay(), Duration::from_millis(200));
let _ = backoff.next_duration();
assert_eq!(backoff.current_delay(), Duration::from_millis(400));
backoff.reset();
assert_eq!(backoff.current_delay(), initial);
}
#[rstest]
fn test_validation_zero_initial_delay() {
let result = ExponentialBackoff::new(Duration::ZERO, Duration::from_secs(1), 2.0, 0, false);
assert!(result.is_err());
assert!(
result
.unwrap_err()
.to_string()
.contains("delay_initial must be non-zero")
);
}
#[rstest]
fn test_validation_max_less_than_initial() {
let result = ExponentialBackoff::new(
Duration::from_secs(1),
Duration::from_millis(500),
2.0,
0,
false,
);
assert!(result.is_err());
assert!(
result
.unwrap_err()
.to_string()
.contains("delay_max must be >= delay_initial")
);
}
#[rstest]
fn test_validation_factor_too_small() {
let result = ExponentialBackoff::new(
Duration::from_millis(100),
Duration::from_secs(1),
0.5,
0,
false,
);
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("factor"));
}
#[rstest]
fn test_validation_factor_too_large() {
let result = ExponentialBackoff::new(
Duration::from_millis(100),
Duration::from_secs(1),
150.0,
0,
false,
);
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("factor"));
}
#[rstest]
fn test_validation_delay_max_exceeds_u64_max_nanos() {
let max_valid = Duration::from_nanos(u64::MAX);
let too_large = max_valid + Duration::from_nanos(1);
let result = ExponentialBackoff::new(Duration::from_millis(100), too_large, 2.0, 0, false);
assert!(result.is_err());
assert!(
result
.unwrap_err()
.to_string()
.contains("delay_max exceeds maximum representable duration")
);
}
#[rstest]
fn test_immediate_first() {
let initial = Duration::from_millis(100);
let max = Duration::from_millis(1600);
let factor = 2.0;
let jitter = 0;
let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, true).unwrap();
let d1 = backoff.next_duration();
assert_eq!(
d1,
Duration::ZERO,
"Expected immediate reconnect (zero delay) on first call"
);
let d2 = backoff.next_duration();
assert_eq!(
d2, initial,
"Expected the delay to be the initial delay after immediate reconnect"
);
let d3 = backoff.next_duration();
let expected = initial * 2; assert_eq!(
d3, expected,
"Expected exponential growth from the initial delay"
);
}
#[rstest]
fn test_reset_restores_immediate_first() {
let initial = Duration::from_millis(100);
let max = Duration::from_millis(1600);
let factor = 2.0;
let jitter = 0;
let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, true).unwrap();
let d1 = backoff.next_duration();
assert_eq!(d1, Duration::ZERO);
let d2 = backoff.next_duration();
assert_eq!(d2, initial);
backoff.reset();
let d3 = backoff.next_duration();
assert_eq!(
d3,
Duration::ZERO,
"Reset should restore immediate_first behavior"
);
}
#[rstest]
fn test_jitter_never_exceeds_max_delay() {
let initial = Duration::from_millis(100);
let max = Duration::from_secs(1);
let factor = 2.0;
let jitter = 500;
let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, false).unwrap();
while backoff.current_delay() < max {
backoff.next_duration();
}
for _ in 0..100 {
let delay = backoff.next_duration();
assert!(
delay <= max,
"Delay with jitter {delay:?} exceeded max {max:?}"
);
}
}
#[rstest]
fn test_jitter_spreads_delays_at_cap() {
let initial = Duration::from_millis(100);
let max = Duration::from_secs(1);
let mut backoff = ExponentialBackoff::new(initial, max, 2.0, 500, false).unwrap();
while backoff.current_delay() < max {
backoff.next_duration();
}
let mut distinct = std::collections::HashSet::new();
for _ in 0..100 {
distinct.insert(backoff.next_duration());
}
assert!(
distinct.len() >= 2,
"Jitter must keep spreading delays once the backoff saturates at the cap"
);
}
#[rstest]
fn test_jitter_wider_than_max_never_returns_zero_delay() {
let max = Duration::from_millis(50);
let mut backoff =
ExponentialBackoff::new(Duration::from_millis(10), max, 2.0, 100, false).unwrap();
for _ in 0..200 {
let delay = backoff.next_duration();
assert!(
!delay.is_zero(),
"Non-immediate backoff delay must be positive"
);
assert!(delay <= max, "Delay {delay:?} exceeded max {max:?}");
}
}
#[cfg(not(all(feature = "simulation", madsim)))]
#[tokio::test(flavor = "current_thread", start_paused = true)]
async fn test_throttle_passes_delay_through_below_attempt_threshold() {
let mut throttle = ReconnectThrottle::default();
for _ in 0..RECONNECT_MIN_DELAY_ATTEMPTS {
assert_eq!(
throttle.gated_delay(Duration::ZERO),
Duration::ZERO,
"Cold window must not floor an immediate reconnect"
);
throttle.record_attempt();
}
}
#[cfg(not(all(feature = "simulation", madsim)))]
#[tokio::test(flavor = "current_thread", start_paused = true)]
async fn test_throttle_floors_delay_once_threshold_trips() {
let mut throttle = ReconnectThrottle::default();
for _ in 0..RECONNECT_MIN_DELAY_ATTEMPTS {
throttle.record_attempt();
}
assert_eq!(
throttle.gated_delay(Duration::ZERO),
RECONNECT_MIN_DELAY,
"Hot window must floor an immediate reconnect"
);
assert_eq!(
throttle.gated_delay(Duration::from_millis(25)),
RECONNECT_MIN_DELAY,
"Hot window must raise a sub-floor backoff delay"
);
assert_eq!(
throttle.gated_delay(Duration::from_secs(5)),
Duration::from_secs(5),
"Hot window must not lower a backoff delay above the floor"
);
}
#[cfg(not(all(feature = "simulation", madsim)))]
#[tokio::test(flavor = "current_thread", start_paused = true)]
async fn test_throttle_lifts_floor_after_window_expires() {
let mut throttle = ReconnectThrottle::default();
for _ in 0..RECONNECT_MIN_DELAY_ATTEMPTS {
throttle.record_attempt();
}
assert_eq!(throttle.gated_delay(Duration::ZERO), RECONNECT_MIN_DELAY);
dst::time::sleep(RECONNECT_MIN_DELAY_WINDOW + Duration::from_secs(1)).await;
assert_eq!(
throttle.gated_delay(Duration::ZERO),
Duration::ZERO,
"Floor must lift once the rolling window drains"
);
}
#[cfg(not(all(feature = "simulation", madsim)))]
#[tokio::test(flavor = "current_thread", start_paused = true)]
async fn test_throttle_window_is_purely_time_based() {
let mut throttle = ReconnectThrottle::default();
for _ in 0..RECONNECT_MIN_DELAY_ATTEMPTS {
throttle.record_attempt();
}
dst::time::sleep(Duration::from_mins(1)).await;
assert_eq!(
throttle.gated_delay(Duration::ZERO),
RECONNECT_MIN_DELAY,
"Stable uptime inside the window must not lift the floor"
);
dst::time::sleep(RECONNECT_MIN_DELAY_WINDOW).await;
assert_eq!(
throttle.gated_delay(Duration::ZERO),
Duration::ZERO,
"Floor must lift once attempts drain from the window"
);
}
#[cfg(not(all(feature = "simulation", madsim)))]
mod throttle_proptests {
use proptest::prelude::*;
use rstest::rstest;
use super::*;
fn build_paused_runtime() -> tokio::runtime::Runtime {
tokio::runtime::Builder::new_current_thread()
.enable_time()
.start_paused(true)
.build()
.unwrap()
}
#[derive(Clone, Debug)]
enum ThrottleOp {
Attempt,
AdvanceMs(u64),
Gate(u64),
}
fn throttle_op_strategy() -> impl Strategy<Value = ThrottleOp> {
prop_oneof![
3 => Just(ThrottleOp::Attempt),
2 => (0u64..=180_000).prop_map(ThrottleOp::AdvanceMs),
3 => (0u64..=10_000).prop_map(ThrottleOp::Gate),
]
}
proptest! {
#![proptest_config(ProptestConfig {
failure_persistence: Some(Box::new(
proptest::test_runner::FileFailurePersistence::Direct(
concat!(env!("CARGO_MANIFEST_DIR"), "/proptest-regressions/backoff.txt")
)
)),
..ProptestConfig::default()
})]
#[rstest]
fn test_throttle_threshold_boundary(
attempt_count in 0usize..=8,
input_ms in 0u64..=5_000,
) {
let runtime = build_paused_runtime();
runtime.block_on(async move {
let mut throttle = ReconnectThrottle::default();
for _ in 0..attempt_count {
throttle.record_attempt();
}
let input = Duration::from_millis(input_ms);
let expected = if attempt_count >= RECONNECT_MIN_DELAY_ATTEMPTS {
input.max(RECONNECT_MIN_DELAY)
} else {
input
};
assert_eq!(
throttle.gated_delay(input),
expected,
"Threshold mismatch at {attempt_count} attempts in window"
);
});
}
#[rstest]
fn test_throttle_matches_rolling_window_model(
ops in prop::collection::vec(throttle_op_strategy(), 1..=500),
) {
let runtime = build_paused_runtime();
runtime.block_on(async move {
let mut throttle = ReconnectThrottle::default();
let mut attempt_times_ms: Vec<u64> = Vec::new();
let mut now_ms = 0u64;
let window_ms = RECONNECT_MIN_DELAY_WINDOW.as_millis() as u64;
for op in ops {
match op {
ThrottleOp::Attempt => {
throttle.record_attempt();
attempt_times_ms.push(now_ms);
}
ThrottleOp::AdvanceMs(ms) => {
dst::time::sleep(Duration::from_millis(ms)).await;
now_ms += ms;
}
ThrottleOp::Gate(input_ms) => {
let input = Duration::from_millis(input_ms);
let kept = attempt_times_ms
.iter()
.filter(|t| now_ms - *t <= window_ms)
.count();
let expected = if kept >= RECONNECT_MIN_DELAY_ATTEMPTS {
input.max(RECONNECT_MIN_DELAY)
} else {
input
};
let gated = throttle.gated_delay(input);
assert_eq!(
gated, expected,
"Model mismatch at {now_ms}ms with {kept} attempts in window"
);
assert!(
gated >= input,
"Floor must never lower the backoff delay"
);
}
}
}
});
}
}
}
}