1use std::{collections::VecDeque, pin::pin, sync::atomic::AtomicU8, time::Duration};
24
25#[cfg(all(feature = "simulation", madsim))]
26use madsim::rand::Rng;
27use nautilus_core::correctness::{check_in_range_inclusive_f64, check_predicate_true};
28use rand::RngExt;
29
30use crate::{dst, mode::ConnectionMode};
31
32pub(crate) const RECONNECT_STABILITY_THRESHOLD: Duration = Duration::from_secs(10);
34
35pub(crate) const RECONNECT_MIN_DELAY: Duration = Duration::from_secs(1);
40
41pub(crate) const RECONNECT_MIN_DELAY_ATTEMPTS: usize = 3;
44
45pub(crate) const RECONNECT_MIN_DELAY_WINDOW: Duration = Duration::from_mins(2);
50
51#[derive(Debug, Default)]
59pub(crate) struct ReconnectThrottle {
60 recent_attempts: VecDeque<dst::time::Instant>,
61}
62
63impl ReconnectThrottle {
64 pub(crate) fn gated_delay(&mut self, backoff_delay: Duration) -> Duration {
68 self.prune_expired();
69
70 if self.recent_attempts.len() >= RECONNECT_MIN_DELAY_ATTEMPTS {
71 backoff_delay.max(RECONNECT_MIN_DELAY)
72 } else {
73 backoff_delay
74 }
75 }
76
77 pub(crate) fn record_attempt(&mut self) {
79 self.prune_expired();
80 self.recent_attempts.push_back(dst::time::Instant::now());
81 }
82
83 fn prune_expired(&mut self) {
84 while self
85 .recent_attempts
86 .front()
87 .is_some_and(|oldest| oldest.elapsed() > RECONNECT_MIN_DELAY_WINDOW)
88 {
89 self.recent_attempts.pop_front();
90 }
91 }
92}
93
94#[derive(Clone, Debug)]
95pub struct ExponentialBackoff {
96 delay_initial: Duration,
97 delay_max: Duration,
98 delay_current: Duration,
99 factor: f64,
100 jitter_ms: u64,
101 immediate_reconnect: bool,
102 immediate_reconnect_original: bool,
103}
104
105impl ExponentialBackoff {
112 pub fn new(
122 delay_initial: Duration,
123 delay_max: Duration,
124 factor: f64,
125 jitter_ms: u64,
126 immediate_first: bool,
127 ) -> anyhow::Result<Self> {
128 check_predicate_true(!delay_initial.is_zero(), "delay_initial must be non-zero")?;
129 check_predicate_true(
130 delay_max >= delay_initial,
131 "delay_max must be >= delay_initial",
132 )?;
133 check_predicate_true(
134 delay_max.as_nanos() <= u128::from(u64::MAX),
135 "delay_max exceeds maximum representable duration (≈584 years)",
136 )?;
137 check_in_range_inclusive_f64(factor, 1.0, 100.0, "factor")?;
138
139 Ok(Self {
140 delay_initial,
141 delay_max,
142 delay_current: delay_initial,
143 factor,
144 jitter_ms,
145 immediate_reconnect: immediate_first,
146 immediate_reconnect_original: immediate_first,
147 })
148 }
149
150 pub fn next_duration(&mut self) -> Duration {
160 if self.immediate_reconnect && self.delay_current == self.delay_initial {
161 self.immediate_reconnect = false;
162 return Duration::ZERO;
163 }
164
165 #[cfg(not(all(feature = "simulation", madsim)))]
167 let jitter = rand::rng().random_range(0..=self.jitter_ms);
168 #[cfg(all(feature = "simulation", madsim))]
169 let jitter = if madsim::runtime::Handle::try_current().is_ok() {
170 madsim::rand::thread_rng().gen_range(0..=self.jitter_ms)
171 } else {
172 rand::rng().random_range(0..=self.jitter_ms) };
174
175 let base = std::cmp::min(
177 self.delay_current,
178 self.delay_max
179 .saturating_sub(Duration::from_millis(self.jitter_ms)),
180 );
181 let delay_with_jitter = base + Duration::from_millis(jitter);
182
183 let clamped_delay = delay_with_jitter.clamp(self.delay_initial, self.delay_max);
185
186 let current_nanos = self.delay_current.as_nanos() as u64;
189 let max_nanos = self.delay_max.as_nanos() as u64;
190 let next_nanos = (current_nanos as f64 * self.factor) as u64;
191 self.delay_current = Duration::from_nanos(next_nanos.min(max_nanos));
192
193 clamped_delay
194 }
195
196 pub const fn reset(&mut self) {
198 self.delay_current = self.delay_initial;
199 self.immediate_reconnect = self.immediate_reconnect_original;
200 }
201
202 #[must_use]
206 pub const fn current_delay(&self) -> Duration {
207 self.delay_current
208 }
209}
210
211pub(crate) async fn wait_reconnect_delay(
212 duration: Duration,
213 connection_mode: &AtomicU8,
214 state_notify: &tokio::sync::Notify,
215) -> bool {
216 if duration.is_zero() {
217 return true;
218 }
219
220 tokio::select! {
221 biased;
222 () = dst::time::sleep(duration) => true,
223 () = async {
224 loop {
225 let mut notified = pin!(state_notify.notified());
226 notified.as_mut().enable();
227
228 if !ConnectionMode::from_atomic(connection_mode).is_reconnect() {
229 break;
230 }
231 notified.await;
232 }
233 } => false,
234 }
235}
236
237#[cfg(test)]
238mod tests {
239 use std::time::Duration;
240
241 use rstest::rstest;
242
243 use super::*;
244
245 #[rstest]
246 fn test_no_jitter_exponential_growth() {
247 let initial = Duration::from_millis(100);
248 let max = Duration::from_millis(1600);
249 let factor = 2.0;
250 let jitter = 0;
251 let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, false).unwrap();
252
253 let d1 = backoff.next_duration();
255 assert_eq!(d1, Duration::from_millis(100));
256
257 let d2 = backoff.next_duration();
259 assert_eq!(d2, Duration::from_millis(200));
260
261 let d3 = backoff.next_duration();
263 assert_eq!(d3, Duration::from_millis(400));
264
265 let d4 = backoff.next_duration();
267 assert_eq!(d4, Duration::from_millis(800));
268
269 let d5 = backoff.next_duration();
271 assert_eq!(d5, Duration::from_millis(1600));
272
273 let d6 = backoff.next_duration();
275 assert_eq!(d6, Duration::from_millis(1600));
276 }
277
278 #[rstest]
279 fn test_reset() {
280 let initial = Duration::from_millis(100);
281 let max = Duration::from_millis(1600);
282 let factor = 2.0;
283 let jitter = 0;
284 let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, false).unwrap();
285
286 let _ = backoff.next_duration(); backoff.reset();
289 let d = backoff.next_duration();
290 assert_eq!(d, Duration::from_millis(100));
292 }
293
294 #[rstest]
295 fn test_jitter_within_bounds() {
296 let initial = Duration::from_millis(100);
297 let max = Duration::from_secs(1);
298 let factor = 2.0;
299 let jitter = 50;
300 for _ in 0..10 {
302 let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, false).unwrap();
303 let base = backoff.delay_current;
305 let delay = backoff.next_duration();
306 let min_expected = base;
308 let max_expected = base + Duration::from_millis(jitter);
309 assert!(
310 delay >= min_expected,
311 "Delay {delay:?} is less than expected minimum {min_expected:?}"
312 );
313 assert!(
314 delay <= max_expected,
315 "Delay {delay:?} exceeds expected maximum {max_expected:?}"
316 );
317 }
318 }
319
320 #[rstest]
321 fn test_factor_less_than_two() {
322 let initial = Duration::from_millis(100);
323 let max = Duration::from_millis(200);
324 let factor = 1.5;
325 let jitter = 0;
326 let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, false).unwrap();
327
328 let d1 = backoff.next_duration();
330 assert_eq!(d1, Duration::from_millis(100));
331
332 let d2 = backoff.next_duration();
334 assert_eq!(d2, Duration::from_millis(150));
335
336 let d3 = backoff.next_duration();
338 assert_eq!(d3, Duration::from_millis(200));
339
340 let d4 = backoff.next_duration();
342 assert_eq!(d4, Duration::from_millis(200));
343 }
344
345 #[rstest]
346 fn test_max_delay_is_respected() {
347 let initial = Duration::from_millis(500);
348 let max = Duration::from_secs(1);
349 let factor = 3.0;
350 let jitter = 0;
351 let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, false).unwrap();
352
353 let d1 = backoff.next_duration();
355 assert_eq!(d1, Duration::from_millis(500));
356
357 let d2 = backoff.next_duration();
359 assert_eq!(d2, Duration::from_secs(1));
360
361 let d3 = backoff.next_duration();
363 assert_eq!(d3, Duration::from_secs(1));
364 }
365
366 #[rstest]
367 fn test_current_delay_getter() {
368 let initial = Duration::from_millis(100);
369 let max = Duration::from_millis(1600);
370 let factor = 2.0;
371 let jitter = 0;
372 let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, false).unwrap();
373
374 assert_eq!(backoff.current_delay(), initial);
375
376 let _ = backoff.next_duration();
377 assert_eq!(backoff.current_delay(), Duration::from_millis(200));
378
379 let _ = backoff.next_duration();
380 assert_eq!(backoff.current_delay(), Duration::from_millis(400));
381
382 backoff.reset();
383 assert_eq!(backoff.current_delay(), initial);
384 }
385
386 #[rstest]
387 fn test_validation_zero_initial_delay() {
388 let result = ExponentialBackoff::new(Duration::ZERO, Duration::from_secs(1), 2.0, 0, false);
389 assert!(result.is_err());
390 assert!(
391 result
392 .unwrap_err()
393 .to_string()
394 .contains("delay_initial must be non-zero")
395 );
396 }
397
398 #[rstest]
399 fn test_validation_max_less_than_initial() {
400 let result = ExponentialBackoff::new(
401 Duration::from_secs(1),
402 Duration::from_millis(500),
403 2.0,
404 0,
405 false,
406 );
407 assert!(result.is_err());
408 assert!(
409 result
410 .unwrap_err()
411 .to_string()
412 .contains("delay_max must be >= delay_initial")
413 );
414 }
415
416 #[rstest]
417 fn test_validation_factor_too_small() {
418 let result = ExponentialBackoff::new(
419 Duration::from_millis(100),
420 Duration::from_secs(1),
421 0.5,
422 0,
423 false,
424 );
425 assert!(result.is_err());
426 assert!(result.unwrap_err().to_string().contains("factor"));
427 }
428
429 #[rstest]
430 fn test_validation_factor_too_large() {
431 let result = ExponentialBackoff::new(
432 Duration::from_millis(100),
433 Duration::from_secs(1),
434 150.0,
435 0,
436 false,
437 );
438 assert!(result.is_err());
439 assert!(result.unwrap_err().to_string().contains("factor"));
440 }
441
442 #[rstest]
443 fn test_validation_delay_max_exceeds_u64_max_nanos() {
444 let max_valid = Duration::from_nanos(u64::MAX);
447 let too_large = max_valid + Duration::from_nanos(1);
448
449 let result = ExponentialBackoff::new(Duration::from_millis(100), too_large, 2.0, 0, false);
450 assert!(result.is_err());
451 assert!(
452 result
453 .unwrap_err()
454 .to_string()
455 .contains("delay_max exceeds maximum representable duration")
456 );
457 }
458
459 #[rstest]
460 fn test_immediate_first() {
461 let initial = Duration::from_millis(100);
462 let max = Duration::from_millis(1600);
463 let factor = 2.0;
464 let jitter = 0;
465 let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, true).unwrap();
466
467 let d1 = backoff.next_duration();
469 assert_eq!(
470 d1,
471 Duration::ZERO,
472 "Expected immediate reconnect (zero delay) on first call"
473 );
474
475 let d2 = backoff.next_duration();
477 assert_eq!(
478 d2, initial,
479 "Expected the delay to be the initial delay after immediate reconnect"
480 );
481
482 let d3 = backoff.next_duration();
484 let expected = initial * 2; assert_eq!(
486 d3, expected,
487 "Expected exponential growth from the initial delay"
488 );
489 }
490
491 #[rstest]
492 fn test_reset_restores_immediate_first() {
493 let initial = Duration::from_millis(100);
494 let max = Duration::from_millis(1600);
495 let factor = 2.0;
496 let jitter = 0;
497 let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, true).unwrap();
498
499 let d1 = backoff.next_duration();
501 assert_eq!(d1, Duration::ZERO);
502
503 let d2 = backoff.next_duration();
505 assert_eq!(d2, initial);
506
507 backoff.reset();
509 let d3 = backoff.next_duration();
510 assert_eq!(
511 d3,
512 Duration::ZERO,
513 "Reset should restore immediate_first behavior"
514 );
515 }
516
517 #[rstest]
518 fn test_jitter_never_exceeds_max_delay() {
519 let initial = Duration::from_millis(100);
520 let max = Duration::from_secs(1);
521 let factor = 2.0;
522 let jitter = 500;
523
524 let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, false).unwrap();
525
526 while backoff.current_delay() < max {
528 backoff.next_duration();
529 }
530
531 for _ in 0..100 {
533 let delay = backoff.next_duration();
534 assert!(
535 delay <= max,
536 "Delay with jitter {delay:?} exceeded max {max:?}"
537 );
538 }
539 }
540
541 #[rstest]
542 fn test_jitter_spreads_delays_at_cap() {
543 let initial = Duration::from_millis(100);
547 let max = Duration::from_secs(1);
548 let mut backoff = ExponentialBackoff::new(initial, max, 2.0, 500, false).unwrap();
549
550 while backoff.current_delay() < max {
551 backoff.next_duration();
552 }
553
554 let mut distinct = std::collections::HashSet::new();
555 for _ in 0..100 {
556 distinct.insert(backoff.next_duration());
557 }
558
559 assert!(
560 distinct.len() >= 2,
561 "Jitter must keep spreading delays once the backoff saturates at the cap"
562 );
563 }
564
565 #[rstest]
566 fn test_jitter_wider_than_max_never_returns_zero_delay() {
567 let max = Duration::from_millis(50);
570 let mut backoff =
571 ExponentialBackoff::new(Duration::from_millis(10), max, 2.0, 100, false).unwrap();
572
573 for _ in 0..200 {
574 let delay = backoff.next_duration();
575 assert!(
576 !delay.is_zero(),
577 "Non-immediate backoff delay must be positive"
578 );
579 assert!(delay <= max, "Delay {delay:?} exceeded max {max:?}");
580 }
581 }
582
583 #[cfg(not(all(feature = "simulation", madsim)))]
584 #[tokio::test(start_paused = true)]
585 async fn test_reconnect_delay_ignores_notifications_until_elapsed() {
586 let mode = AtomicU8::new(ConnectionMode::Reconnect.as_u8());
587 let notify = tokio::sync::Notify::new();
588 let mut wait = pin!(wait_reconnect_delay(Duration::from_secs(5), &mode, ¬ify));
589
590 assert!(futures_util::poll!(&mut wait).is_pending());
591 notify.notify_waiters();
592 assert!(futures_util::poll!(&mut wait).is_pending());
593 tokio::time::advance(Duration::from_secs(4)).await;
594 assert!(futures_util::poll!(&mut wait).is_pending());
595 tokio::time::advance(Duration::from_secs(1)).await;
596 assert_eq!(futures_util::poll!(&mut wait), std::task::Poll::Ready(true));
597 }
598
599 #[cfg(not(all(feature = "simulation", madsim)))]
600 #[rstest]
601 #[case::disconnect_before_wait(ConnectionMode::Disconnect, false)]
602 #[case::closed_before_wait(ConnectionMode::Closed, false)]
603 #[case::disconnect_during_wait(ConnectionMode::Disconnect, true)]
604 #[case::closed_during_wait(ConnectionMode::Closed, true)]
605 #[tokio::test(start_paused = true)]
606 async fn test_reconnect_delay_stops_on_terminal_state(
607 #[case] terminal: ConnectionMode,
608 #[case] during_wait: bool,
609 ) {
610 let mode = AtomicU8::new(if during_wait {
611 ConnectionMode::Reconnect.as_u8()
612 } else {
613 terminal.as_u8()
614 });
615 let notify = tokio::sync::Notify::new();
616 let mut wait = pin!(wait_reconnect_delay(Duration::from_secs(5), &mode, ¬ify));
617
618 if during_wait {
619 assert!(futures_util::poll!(&mut wait).is_pending());
620 mode.store(terminal.as_u8(), std::sync::atomic::Ordering::SeqCst);
621 notify.notify_waiters();
622 }
623
624 assert_eq!(
625 futures_util::poll!(&mut wait),
626 std::task::Poll::Ready(false)
627 );
628 }
629
630 #[cfg(not(all(feature = "simulation", madsim)))]
634 #[tokio::test(flavor = "current_thread", start_paused = true)]
635 async fn test_throttle_passes_delay_through_below_attempt_threshold() {
636 let mut throttle = ReconnectThrottle::default();
637
638 for _ in 0..RECONNECT_MIN_DELAY_ATTEMPTS {
639 assert_eq!(
640 throttle.gated_delay(Duration::ZERO),
641 Duration::ZERO,
642 "Cold window must not floor an immediate reconnect"
643 );
644 throttle.record_attempt();
645 }
646 }
647
648 #[cfg(not(all(feature = "simulation", madsim)))]
649 #[tokio::test(flavor = "current_thread", start_paused = true)]
650 async fn test_throttle_floors_delay_once_threshold_trips() {
651 let mut throttle = ReconnectThrottle::default();
652
653 for _ in 0..RECONNECT_MIN_DELAY_ATTEMPTS {
654 throttle.record_attempt();
655 }
656
657 assert_eq!(
658 throttle.gated_delay(Duration::ZERO),
659 RECONNECT_MIN_DELAY,
660 "Hot window must floor an immediate reconnect"
661 );
662 assert_eq!(
663 throttle.gated_delay(Duration::from_millis(25)),
664 RECONNECT_MIN_DELAY,
665 "Hot window must raise a sub-floor backoff delay"
666 );
667 assert_eq!(
668 throttle.gated_delay(Duration::from_secs(5)),
669 Duration::from_secs(5),
670 "Hot window must not lower a backoff delay above the floor"
671 );
672 }
673
674 #[cfg(not(all(feature = "simulation", madsim)))]
675 #[tokio::test(flavor = "current_thread", start_paused = true)]
676 async fn test_throttle_lifts_floor_after_window_expires() {
677 let mut throttle = ReconnectThrottle::default();
678
679 for _ in 0..RECONNECT_MIN_DELAY_ATTEMPTS {
680 throttle.record_attempt();
681 }
682
683 assert_eq!(throttle.gated_delay(Duration::ZERO), RECONNECT_MIN_DELAY);
684
685 dst::time::sleep(RECONNECT_MIN_DELAY_WINDOW + Duration::from_secs(1)).await;
686
687 assert_eq!(
688 throttle.gated_delay(Duration::ZERO),
689 Duration::ZERO,
690 "Floor must lift once the rolling window drains"
691 );
692 }
693
694 #[cfg(not(all(feature = "simulation", madsim)))]
695 #[tokio::test(flavor = "current_thread", start_paused = true)]
696 async fn test_throttle_window_is_purely_time_based() {
697 let mut throttle = ReconnectThrottle::default();
698
699 for _ in 0..RECONNECT_MIN_DELAY_ATTEMPTS {
700 throttle.record_attempt();
701 }
702
703 dst::time::sleep(Duration::from_mins(1)).await;
706 assert_eq!(
707 throttle.gated_delay(Duration::ZERO),
708 RECONNECT_MIN_DELAY,
709 "Stable uptime inside the window must not lift the floor"
710 );
711
712 dst::time::sleep(RECONNECT_MIN_DELAY_WINDOW).await;
713 assert_eq!(
714 throttle.gated_delay(Duration::ZERO),
715 Duration::ZERO,
716 "Floor must lift once attempts drain from the window"
717 );
718 }
719
720 #[cfg(not(all(feature = "simulation", madsim)))]
721 mod throttle_proptests {
722 use proptest::prelude::*;
723 use rstest::rstest;
724
725 use super::*;
726
727 fn build_paused_runtime() -> tokio::runtime::Runtime {
728 tokio::runtime::Builder::new_current_thread()
729 .enable_time()
730 .start_paused(true)
731 .build()
732 .unwrap()
733 }
734
735 #[derive(Clone, Debug)]
736 enum ThrottleOp {
737 Attempt,
738 AdvanceMs(u64),
739 Gate(u64),
740 }
741
742 fn throttle_op_strategy() -> impl Strategy<Value = ThrottleOp> {
743 prop_oneof![
744 3 => Just(ThrottleOp::Attempt),
745 2 => (0u64..=180_000).prop_map(ThrottleOp::AdvanceMs),
746 3 => (0u64..=10_000).prop_map(ThrottleOp::Gate),
747 ]
748 }
749
750 proptest! {
751 #![proptest_config(ProptestConfig {
753 failure_persistence: Some(Box::new(
754 proptest::test_runner::FileFailurePersistence::Direct(
755 concat!(env!("CARGO_MANIFEST_DIR"), "/proptest-regressions/backoff.txt")
756 )
757 )),
758 ..ProptestConfig::default()
759 })]
760
761 #[rstest]
762 fn test_throttle_threshold_boundary(
763 attempt_count in 0usize..=8,
764 input_ms in 0u64..=5_000,
765 ) {
766 let runtime = build_paused_runtime();
767 runtime.block_on(async move {
768 let mut throttle = ReconnectThrottle::default();
769
770 for _ in 0..attempt_count {
771 throttle.record_attempt();
772 }
773
774 let input = Duration::from_millis(input_ms);
775 let expected = if attempt_count >= RECONNECT_MIN_DELAY_ATTEMPTS {
776 input.max(RECONNECT_MIN_DELAY)
777 } else {
778 input
779 };
780 assert_eq!(
781 throttle.gated_delay(input),
782 expected,
783 "Threshold mismatch at {attempt_count} attempts in window"
784 );
785 });
786 }
787
788 #[rstest]
789 fn test_throttle_matches_rolling_window_model(
790 ops in prop::collection::vec(throttle_op_strategy(), 1..=500),
791 ) {
792 let runtime = build_paused_runtime();
793 runtime.block_on(async move {
794 let mut throttle = ReconnectThrottle::default();
795 let mut attempt_times_ms: Vec<u64> = Vec::new();
796 let mut now_ms = 0u64;
797 let window_ms = RECONNECT_MIN_DELAY_WINDOW.as_millis() as u64;
798
799 for op in ops {
800 match op {
801 ThrottleOp::Attempt => {
802 throttle.record_attempt();
803 attempt_times_ms.push(now_ms);
804 }
805 ThrottleOp::AdvanceMs(ms) => {
806 dst::time::sleep(Duration::from_millis(ms)).await;
807 now_ms += ms;
808 }
809 ThrottleOp::Gate(input_ms) => {
810 let input = Duration::from_millis(input_ms);
811 let kept = attempt_times_ms
812 .iter()
813 .filter(|t| now_ms - *t <= window_ms)
814 .count();
815 let expected = if kept >= RECONNECT_MIN_DELAY_ATTEMPTS {
816 input.max(RECONNECT_MIN_DELAY)
817 } else {
818 input
819 };
820 let gated = throttle.gated_delay(input);
821 assert_eq!(
822 gated, expected,
823 "Model mismatch at {now_ms}ms with {kept} attempts in window"
824 );
825 assert!(
826 gated >= input,
827 "Floor must never lower the backoff delay"
828 );
829 }
830 }
831 }
832 });
833 }
834 }
835 }
836}