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
169 #[cfg(all(feature = "simulation", madsim))]
170 let jitter = if madsim::runtime::Handle::try_current().is_ok() {
171 madsim::rand::thread_rng().gen_range(0..=self.jitter_ms)
172 } else {
173 rand::rng().random_range(0..=self.jitter_ms) };
175
176 let base = std::cmp::min(
178 self.delay_current,
179 self.delay_max
180 .saturating_sub(Duration::from_millis(self.jitter_ms)),
181 );
182 let delay_with_jitter = base + Duration::from_millis(jitter);
183
184 let clamped_delay = delay_with_jitter.clamp(self.delay_initial, self.delay_max);
186
187 let current_nanos = self.delay_current.as_nanos() as u64;
190 let max_nanos = self.delay_max.as_nanos() as u64;
191 let next_nanos = (current_nanos as f64 * self.factor) as u64;
192 self.delay_current = Duration::from_nanos(next_nanos.min(max_nanos));
193
194 clamped_delay
195 }
196
197 pub const fn reset(&mut self) {
199 self.delay_current = self.delay_initial;
200 self.immediate_reconnect = self.immediate_reconnect_original;
201 }
202
203 #[must_use]
207 pub const fn current_delay(&self) -> Duration {
208 self.delay_current
209 }
210}
211
212pub(crate) async fn wait_reconnect_delay(
213 duration: Duration,
214 connection_mode: &AtomicU8,
215 state_notify: &tokio::sync::Notify,
216) -> bool {
217 if duration.is_zero() {
218 return true;
219 }
220
221 tokio::select! {
222 biased;
223 () = dst::time::sleep(duration) => true,
224 () = async {
225 loop {
226 let mut notified = pin!(state_notify.notified());
227 notified.as_mut().enable();
228
229 if !ConnectionMode::from_atomic(connection_mode).is_reconnect() {
230 break;
231 }
232 notified.await;
233 }
234 } => false,
235 }
236}
237
238#[cfg(test)]
239mod tests {
240 use std::time::Duration;
241
242 use rstest::rstest;
243
244 use super::*;
245
246 #[rstest]
247 fn test_no_jitter_exponential_growth() {
248 let initial = Duration::from_millis(100);
249 let max = Duration::from_millis(1600);
250 let factor = 2.0;
251 let jitter = 0;
252 let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, false).unwrap();
253
254 let d1 = backoff.next_duration();
256 assert_eq!(d1, Duration::from_millis(100));
257
258 let d2 = backoff.next_duration();
260 assert_eq!(d2, Duration::from_millis(200));
261
262 let d3 = backoff.next_duration();
264 assert_eq!(d3, Duration::from_millis(400));
265
266 let d4 = backoff.next_duration();
268 assert_eq!(d4, Duration::from_millis(800));
269
270 let d5 = backoff.next_duration();
272 assert_eq!(d5, Duration::from_millis(1600));
273
274 let d6 = backoff.next_duration();
276 assert_eq!(d6, Duration::from_millis(1600));
277 }
278
279 #[rstest]
280 fn test_reset() {
281 let initial = Duration::from_millis(100);
282 let max = Duration::from_millis(1600);
283 let factor = 2.0;
284 let jitter = 0;
285 let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, false).unwrap();
286
287 let _ = backoff.next_duration(); backoff.reset();
290 let d = backoff.next_duration();
291 assert_eq!(d, Duration::from_millis(100));
293 }
294
295 #[rstest]
296 fn test_jitter_within_bounds() {
297 let initial = Duration::from_millis(100);
298 let max = Duration::from_secs(1);
299 let factor = 2.0;
300 let jitter = 50;
301
302 for _ in 0..10 {
304 let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, false).unwrap();
305 let base = backoff.delay_current;
307 let delay = backoff.next_duration();
308 let min_expected = base;
310 let max_expected = base + Duration::from_millis(jitter);
311 assert!(
312 delay >= min_expected,
313 "Delay {delay:?} is less than expected minimum {min_expected:?}"
314 );
315 assert!(
316 delay <= max_expected,
317 "Delay {delay:?} exceeds expected maximum {max_expected:?}"
318 );
319 }
320 }
321
322 #[rstest]
323 fn test_factor_less_than_two() {
324 let initial = Duration::from_millis(100);
325 let max = Duration::from_millis(200);
326 let factor = 1.5;
327 let jitter = 0;
328 let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, false).unwrap();
329
330 let d1 = backoff.next_duration();
332 assert_eq!(d1, Duration::from_millis(100));
333
334 let d2 = backoff.next_duration();
336 assert_eq!(d2, Duration::from_millis(150));
337
338 let d3 = backoff.next_duration();
340 assert_eq!(d3, Duration::from_millis(200));
341
342 let d4 = backoff.next_duration();
344 assert_eq!(d4, Duration::from_millis(200));
345 }
346
347 #[rstest]
348 fn test_max_delay_is_respected() {
349 let initial = Duration::from_millis(500);
350 let max = Duration::from_secs(1);
351 let factor = 3.0;
352 let jitter = 0;
353 let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, false).unwrap();
354
355 let d1 = backoff.next_duration();
357 assert_eq!(d1, Duration::from_millis(500));
358
359 let d2 = backoff.next_duration();
361 assert_eq!(d2, Duration::from_secs(1));
362
363 let d3 = backoff.next_duration();
365 assert_eq!(d3, Duration::from_secs(1));
366 }
367
368 #[rstest]
369 fn test_current_delay_getter() {
370 let initial = Duration::from_millis(100);
371 let max = Duration::from_millis(1600);
372 let factor = 2.0;
373 let jitter = 0;
374 let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, false).unwrap();
375
376 assert_eq!(backoff.current_delay(), initial);
377
378 let _ = backoff.next_duration();
379 assert_eq!(backoff.current_delay(), Duration::from_millis(200));
380
381 let _ = backoff.next_duration();
382 assert_eq!(backoff.current_delay(), Duration::from_millis(400));
383
384 backoff.reset();
385 assert_eq!(backoff.current_delay(), initial);
386 }
387
388 #[rstest]
389 fn test_validation_zero_initial_delay() {
390 let result = ExponentialBackoff::new(Duration::ZERO, Duration::from_secs(1), 2.0, 0, false);
391 assert!(result.is_err());
392 assert!(
393 result
394 .unwrap_err()
395 .to_string()
396 .contains("delay_initial must be non-zero")
397 );
398 }
399
400 #[rstest]
401 fn test_validation_max_less_than_initial() {
402 let result = ExponentialBackoff::new(
403 Duration::from_secs(1),
404 Duration::from_millis(500),
405 2.0,
406 0,
407 false,
408 );
409 assert!(result.is_err());
410 assert!(
411 result
412 .unwrap_err()
413 .to_string()
414 .contains("delay_max must be >= delay_initial")
415 );
416 }
417
418 #[rstest]
419 fn test_validation_factor_too_small() {
420 let result = ExponentialBackoff::new(
421 Duration::from_millis(100),
422 Duration::from_secs(1),
423 0.5,
424 0,
425 false,
426 );
427 assert!(result.is_err());
428 assert!(result.unwrap_err().to_string().contains("factor"));
429 }
430
431 #[rstest]
432 fn test_validation_factor_too_large() {
433 let result = ExponentialBackoff::new(
434 Duration::from_millis(100),
435 Duration::from_secs(1),
436 150.0,
437 0,
438 false,
439 );
440 assert!(result.is_err());
441 assert!(result.unwrap_err().to_string().contains("factor"));
442 }
443
444 #[rstest]
445 fn test_validation_delay_max_exceeds_u64_max_nanos() {
446 let max_valid = Duration::from_nanos(u64::MAX);
449 let too_large = max_valid + Duration::from_nanos(1);
450
451 let result = ExponentialBackoff::new(Duration::from_millis(100), too_large, 2.0, 0, false);
452 assert!(result.is_err());
453 assert!(
454 result
455 .unwrap_err()
456 .to_string()
457 .contains("delay_max exceeds maximum representable duration")
458 );
459 }
460
461 #[rstest]
462 fn test_immediate_first() {
463 let initial = Duration::from_millis(100);
464 let max = Duration::from_millis(1600);
465 let factor = 2.0;
466 let jitter = 0;
467 let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, true).unwrap();
468
469 let d1 = backoff.next_duration();
471 assert_eq!(
472 d1,
473 Duration::ZERO,
474 "Expected immediate reconnect (zero delay) on first call"
475 );
476
477 let d2 = backoff.next_duration();
479 assert_eq!(
480 d2, initial,
481 "Expected the delay to be the initial delay after immediate reconnect"
482 );
483
484 let d3 = backoff.next_duration();
486 let expected = initial * 2; assert_eq!(
488 d3, expected,
489 "Expected exponential growth from the initial delay"
490 );
491 }
492
493 #[rstest]
494 fn test_reset_restores_immediate_first() {
495 let initial = Duration::from_millis(100);
496 let max = Duration::from_millis(1600);
497 let factor = 2.0;
498 let jitter = 0;
499 let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, true).unwrap();
500
501 let d1 = backoff.next_duration();
503 assert_eq!(d1, Duration::ZERO);
504
505 let d2 = backoff.next_duration();
507 assert_eq!(d2, initial);
508
509 backoff.reset();
511 let d3 = backoff.next_duration();
512 assert_eq!(
513 d3,
514 Duration::ZERO,
515 "Reset should restore immediate_first behavior"
516 );
517 }
518
519 #[rstest]
520 fn test_jitter_never_exceeds_max_delay() {
521 let initial = Duration::from_millis(100);
522 let max = Duration::from_secs(1);
523 let factor = 2.0;
524 let jitter = 500;
525
526 let mut backoff = ExponentialBackoff::new(initial, max, factor, jitter, false).unwrap();
527
528 while backoff.current_delay() < max {
530 backoff.next_duration();
531 }
532
533 for _ in 0..100 {
535 let delay = backoff.next_duration();
536 assert!(
537 delay <= max,
538 "Delay with jitter {delay:?} exceeded max {max:?}"
539 );
540 }
541 }
542
543 #[rstest]
544 fn test_jitter_spreads_delays_at_cap() {
545 let initial = Duration::from_millis(100);
549 let max = Duration::from_secs(1);
550 let mut backoff = ExponentialBackoff::new(initial, max, 2.0, 500, false).unwrap();
551 while backoff.current_delay() < max {
552 backoff.next_duration();
553 }
554
555 let mut distinct = std::collections::HashSet::new();
556 for _ in 0..100 {
557 distinct.insert(backoff.next_duration());
558 }
559
560 assert!(
561 distinct.len() >= 2,
562 "Jitter must keep spreading delays once the backoff saturates at the cap"
563 );
564 }
565
566 #[rstest]
567 fn test_jitter_wider_than_max_never_returns_zero_delay() {
568 let max = Duration::from_millis(50);
571 let mut backoff =
572 ExponentialBackoff::new(Duration::from_millis(10), max, 2.0, 100, false).unwrap();
573
574 for _ in 0..200 {
575 let delay = backoff.next_duration();
576 assert!(
577 !delay.is_zero(),
578 "Non-immediate backoff delay must be positive"
579 );
580 assert!(delay <= max, "Delay {delay:?} exceeded max {max:?}");
581 }
582 }
583
584 #[cfg(not(all(feature = "simulation", madsim)))]
585 #[tokio::test(start_paused = true)]
586 async fn test_reconnect_delay_ignores_notifications_until_elapsed() {
587 let mode = AtomicU8::new(ConnectionMode::Reconnect.as_u8());
588 let notify = tokio::sync::Notify::new();
589 let mut wait = pin!(wait_reconnect_delay(Duration::from_secs(5), &mode, ¬ify));
590
591 assert!(futures_util::poll!(&mut wait).is_pending());
592 notify.notify_waiters();
593 assert!(futures_util::poll!(&mut wait).is_pending());
594 tokio::time::advance(Duration::from_secs(4)).await;
595 assert!(futures_util::poll!(&mut wait).is_pending());
596 tokio::time::advance(Duration::from_secs(1)).await;
597 assert_eq!(futures_util::poll!(&mut wait), std::task::Poll::Ready(true));
598 }
599
600 #[cfg(not(all(feature = "simulation", madsim)))]
601 #[rstest]
602 #[case::disconnect_before_wait(ConnectionMode::Disconnect, false)]
603 #[case::closed_before_wait(ConnectionMode::Closed, false)]
604 #[case::disconnect_during_wait(ConnectionMode::Disconnect, true)]
605 #[case::closed_during_wait(ConnectionMode::Closed, true)]
606 #[tokio::test(start_paused = true)]
607 async fn test_reconnect_delay_stops_on_terminal_state(
608 #[case] terminal: ConnectionMode,
609 #[case] during_wait: bool,
610 ) {
611 let mode = AtomicU8::new(if during_wait {
612 ConnectionMode::Reconnect.as_u8()
613 } else {
614 terminal.as_u8()
615 });
616
617 let notify = tokio::sync::Notify::new();
618 let mut wait = pin!(wait_reconnect_delay(Duration::from_secs(5), &mode, ¬ify));
619
620 if during_wait {
621 assert!(futures_util::poll!(&mut wait).is_pending());
622 mode.store(terminal.as_u8(), std::sync::atomic::Ordering::SeqCst);
623 notify.notify_waiters();
624 }
625
626 assert_eq!(
627 futures_util::poll!(&mut wait),
628 std::task::Poll::Ready(false)
629 );
630 }
631
632 #[cfg(not(all(feature = "simulation", madsim)))]
636 #[tokio::test(flavor = "current_thread", start_paused = true)]
637 async fn test_throttle_passes_delay_through_below_attempt_threshold() {
638 let mut throttle = ReconnectThrottle::default();
639
640 for _ in 0..RECONNECT_MIN_DELAY_ATTEMPTS {
641 assert_eq!(
642 throttle.gated_delay(Duration::ZERO),
643 Duration::ZERO,
644 "Cold window must not floor an immediate reconnect"
645 );
646 throttle.record_attempt();
647 }
648 }
649
650 #[cfg(not(all(feature = "simulation", madsim)))]
651 #[tokio::test(flavor = "current_thread", start_paused = true)]
652 async fn test_throttle_floors_delay_once_threshold_trips() {
653 let mut throttle = ReconnectThrottle::default();
654 for _ in 0..RECONNECT_MIN_DELAY_ATTEMPTS {
655 throttle.record_attempt();
656 }
657
658 assert_eq!(
659 throttle.gated_delay(Duration::ZERO),
660 RECONNECT_MIN_DELAY,
661 "Hot window must floor an immediate reconnect"
662 );
663 assert_eq!(
664 throttle.gated_delay(Duration::from_millis(25)),
665 RECONNECT_MIN_DELAY,
666 "Hot window must raise a sub-floor backoff delay"
667 );
668 assert_eq!(
669 throttle.gated_delay(Duration::from_secs(5)),
670 Duration::from_secs(5),
671 "Hot window must not lower a backoff delay above the floor"
672 );
673 }
674
675 #[cfg(not(all(feature = "simulation", madsim)))]
676 #[tokio::test(flavor = "current_thread", start_paused = true)]
677 async fn test_throttle_lifts_floor_after_window_expires() {
678 let mut throttle = ReconnectThrottle::default();
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 for _ in 0..RECONNECT_MIN_DELAY_ATTEMPTS {
699 throttle.record_attempt();
700 }
701
702 dst::time::sleep(Duration::from_mins(1)).await;
705 assert_eq!(
706 throttle.gated_delay(Duration::ZERO),
707 RECONNECT_MIN_DELAY,
708 "Stable uptime inside the window must not lift the floor"
709 );
710
711 dst::time::sleep(RECONNECT_MIN_DELAY_WINDOW).await;
712 assert_eq!(
713 throttle.gated_delay(Duration::ZERO),
714 Duration::ZERO,
715 "Floor must lift once attempts drain from the window"
716 );
717 }
718
719 #[cfg(not(all(feature = "simulation", madsim)))]
720 mod throttle_proptests {
721 use proptest::prelude::*;
722 use rstest::rstest;
723
724 use super::*;
725
726 fn build_paused_runtime() -> tokio::runtime::Runtime {
727 tokio::runtime::Builder::new_current_thread()
728 .enable_time()
729 .start_paused(true)
730 .build()
731 .unwrap()
732 }
733
734 #[derive(Clone, Debug)]
735 enum ThrottleOp {
736 Attempt,
737 AdvanceMs(u64),
738 Gate(u64),
739 }
740
741 fn throttle_op_strategy() -> impl Strategy<Value = ThrottleOp> {
742 prop_oneof![
743 3 => Just(ThrottleOp::Attempt),
744 2 => (0u64..=180_000).prop_map(ThrottleOp::AdvanceMs),
745 3 => (0u64..=10_000).prop_map(ThrottleOp::Gate),
746 ]
747 }
748
749 proptest! {
750 #![proptest_config(ProptestConfig {
752 failure_persistence: Some(Box::new(
753 proptest::test_runner::FileFailurePersistence::Direct(
754 concat!(env!("CARGO_MANIFEST_DIR"), "/proptest-regressions/backoff.txt")
755 )
756 )),
757 ..ProptestConfig::default()
758 })]
759
760 #[rstest]
761 fn test_throttle_threshold_boundary(
762 attempt_count in 0usize..=8,
763 input_ms in 0u64..=5_000,
764 ) {
765 let runtime = build_paused_runtime();
766 runtime.block_on(async move {
767 let mut throttle = ReconnectThrottle::default();
768
769 for _ in 0..attempt_count {
770 throttle.record_attempt();
771 }
772
773 let input = Duration::from_millis(input_ms);
774 let expected = if attempt_count >= RECONNECT_MIN_DELAY_ATTEMPTS {
775 input.max(RECONNECT_MIN_DELAY)
776 } else {
777 input
778 };
779 assert_eq!(
780 throttle.gated_delay(input),
781 expected,
782 "Threshold mismatch at {attempt_count} attempts in window"
783 );
784 });
785 }
786
787 #[rstest]
788 fn test_throttle_matches_rolling_window_model(
789 ops in prop::collection::vec(throttle_op_strategy(), 1..=500),
790 ) {
791 let runtime = build_paused_runtime();
792 runtime.block_on(async move {
793 let mut throttle = ReconnectThrottle::default();
794 let mut attempt_times_ms: Vec<u64> = Vec::new();
795 let mut now_ms = 0u64;
796 let window_ms = RECONNECT_MIN_DELAY_WINDOW.as_millis() as u64;
797
798 for op in ops {
799 match op {
800 ThrottleOp::Attempt => {
801 throttle.record_attempt();
802 attempt_times_ms.push(now_ms);
803 }
804 ThrottleOp::AdvanceMs(ms) => {
805 dst::time::sleep(Duration::from_millis(ms)).await;
806 now_ms += ms;
807 }
808 ThrottleOp::Gate(input_ms) => {
809 let input = Duration::from_millis(input_ms);
810 let kept = attempt_times_ms
811 .iter()
812 .filter(|t| now_ms - *t <= window_ms)
813 .count();
814 let expected = if kept >= RECONNECT_MIN_DELAY_ATTEMPTS {
815 input.max(RECONNECT_MIN_DELAY)
816 } else {
817 input
818 };
819 let gated = throttle.gated_delay(input);
820 assert_eq!(
821 gated, expected,
822 "Model mismatch at {now_ms}ms with {kept} attempts in window"
823 );
824 assert!(
825 gated >= input,
826 "Floor must never lower the backoff delay"
827 );
828 }
829 }
830 }
831 });
832 }
833 }
834 }
835}