1use std::{fmt::Display, future::Future, sync::Arc};
53
54use nautilus_common::live::dst::time::{self, Duration};
55use nautilus_network::{
56 backoff::ExponentialBackoff,
57 retry::{RetryConfig, RetryManager},
58};
59use tokio_util::sync::CancellationToken;
60
61use super::snapshot::SnapshotGate;
62
63const ATTEMPTS_MAX: u32 = 8;
65const ELAPSED_MAX_MS: u64 = 180_000;
66const OPERATION_TIMEOUT_MS: Option<u64> = None;
67
68const RETRY_DELAY_INITIAL_MS: u64 = 1_000;
69const RETRY_DELAY_MAX_MS: u64 = 10_000;
70const RETRY_BACKOFF_FACTOR: f64 = 2.0;
71const RETRY_JITTER_MAX_MS: u64 = 1_000;
72const RETRY_FIRST_IMMEDIATE: bool = true;
73
74const RETRY_DELAY_CEILING_INITIAL_MS: u64 = 60_000;
76const RETRY_DELAY_CEILING_MAX_MS: u64 = 900_000;
77const RETRY_JITTER_CEILING_MAX_MS: u64 = 5_000;
78
79#[derive(Debug, Clone)]
81pub enum BookRecoveryOutcome<E> {
82 Pending,
83 Accepted,
84 Rejected(E),
85}
86
87#[derive(Debug)]
89pub struct BookRecovery<E> {
90 pub cancellation: CancellationToken,
91 pub outcome: tokio::sync::watch::Sender<BookRecoveryOutcome<E>>,
92 pub gate: SnapshotGate,
93 reconnected: tokio::sync::Notify,
94}
95
96impl<E: Clone> Default for BookRecovery<E> {
97 fn default() -> Self {
98 Self {
99 cancellation: CancellationToken::new(),
100 outcome: tokio::sync::watch::channel(BookRecoveryOutcome::Pending).0,
101 gate: SnapshotGate::default(),
102 reconnected: tokio::sync::Notify::new(),
103 }
104 }
105}
106
107impl<E: Clone> BookRecovery<E> {
108 #[must_use]
110 pub fn begin_replacement(&self) -> bool {
111 let mut gate = self.gate.lock();
112
113 if self.cancellation.is_cancelled() || self.is_accepted() {
114 return false;
115 }
116
117 gate.close();
118 true
119 }
120
121 #[must_use]
123 pub fn is_accepted(&self) -> bool {
124 matches!(*self.outcome.borrow(), BookRecoveryOutcome::Accepted)
125 }
126
127 #[must_use]
129 pub fn is_running(&self) -> bool {
130 !self.is_accepted() && !self.cancellation.is_cancelled()
131 }
132
133 pub fn accept(&self) -> bool {
135 let gate = self.gate.lock();
136 if gate.is_closed() || (self.cancellation.is_cancelled() && !self.is_accepted()) {
137 return false;
138 }
139
140 self.outcome.send_replace(BookRecoveryOutcome::Accepted);
141 true
142 }
143}
144
145impl<E: Clone + std::error::Error> BookRecovery<E> {
146 #[allow(
158 clippy::missing_panics_doc,
159 reason = "the ceiling backoff configuration is static and valid"
160 )]
161 pub async fn run<F, Fut>(
162 &self,
163 book: impl Display,
164 snapshot_timeout: Duration,
165 replace: F,
166 should_retry: impl Fn(&E) -> bool,
167 create_error: impl Fn(String) -> E,
168 timeout_error: impl Fn() -> E,
169 ) where
170 F: Fn(CancellationToken, SnapshotGate) -> Fut,
171 Fut: Future<Output = Result<(), E>>,
172 {
173 let attempt = || self.attempt(snapshot_timeout, &replace, &create_error, &timeout_error);
174
175 let budget = RetryManager::<E>::new(RetryConfig {
176 max_retries: ATTEMPTS_MAX - 1,
177 initial_delay_ms: RETRY_DELAY_INITIAL_MS,
178 max_delay_ms: RETRY_DELAY_MAX_MS,
179 backoff_factor: RETRY_BACKOFF_FACTOR,
180 jitter_ms: RETRY_JITTER_MAX_MS,
181 immediate_first: RETRY_FIRST_IMMEDIATE,
182 operation_timeout_ms: OPERATION_TIMEOUT_MS,
183 max_elapsed_ms: Some(ELAPSED_MAX_MS),
184 });
185
186 let result = budget
187 .invocation("book recovery", attempt, should_retry, |e| {
188 create_error(e.to_string())
189 })
190 .cancellation_token(&self.cancellation)
191 .execute()
192 .await;
193
194 if let Err(e) = result {
195 if self.cancellation.is_cancelled() {
196 return;
197 }
198
199 log::error!(
200 "Book recovery for {book} did not complete within its retry budget; retrying at \
201 intervals growing from {}s to {}s until a snapshot is accepted: {e}",
202 RETRY_DELAY_CEILING_INITIAL_MS / 1_000,
203 RETRY_DELAY_CEILING_MAX_MS / 1_000,
204 );
205
206 let attempt_max =
208 snapshot_timeout.max(Duration::from_millis(RETRY_DELAY_CEILING_INITIAL_MS));
209
210 let attempts = async {
211 let mut backoff = ExponentialBackoff::new(
212 Duration::from_millis(RETRY_DELAY_CEILING_INITIAL_MS),
213 Duration::from_millis(RETRY_DELAY_CEILING_MAX_MS),
214 RETRY_BACKOFF_FACTOR,
215 RETRY_JITTER_CEILING_MAX_MS,
216 false,
217 )
218 .expect("static ceiling backoff configuration is valid");
219
220 loop {
221 tokio::select! {
223 biased;
224 () = self.reconnected.notified() => {}
225 () = time::sleep(backoff.next_duration()) => {}
226 }
227
228 let result = time::timeout(attempt_max, attempt())
229 .await
230 .unwrap_or_else(|_| {
231 Err(create_error(format!(
232 "Book recovery attempt timed out after {}ms (whole-attempt deadline)",
233 attempt_max.as_millis(),
234 )))
235 });
236
237 match result {
238 Ok(()) => return,
239 Err(e) => log::warn!("Book recovery attempt for {book} failed: {e}"),
240 }
241 }
242 };
243
244 let mut outcome = self.outcome.subscribe();
245
246 tokio::select! {
249 biased;
250 () = self.cancellation.cancelled() => return,
251 _ = outcome.wait_for(|o| matches!(o, BookRecoveryOutcome::Accepted)) => {}
252 () = attempts => {}
253 }
254 }
255
256 if self.is_accepted() {
257 log::info!("Book recovery completed for {book}");
258 }
259 }
260
261 async fn attempt<F, Fut>(
262 &self,
263 snapshot_timeout: Duration,
264 replace: &F,
265 create_error: &impl Fn(String) -> E,
266 timeout_error: &impl Fn() -> E,
267 ) -> Result<(), E>
268 where
269 F: Fn(CancellationToken, SnapshotGate) -> Fut,
270 Fut: Future<Output = Result<(), E>>,
271 {
272 let mut outcome = self.outcome.subscribe();
273
274 if !self.begin_replacement() {
275 return Ok(());
276 }
277
278 self.outcome.send_if_modified(|outcome| {
279 if matches!(outcome, BookRecoveryOutcome::Rejected(_)) {
280 *outcome = BookRecoveryOutcome::Pending;
281 true
282 } else {
283 false
284 }
285 });
286
287 let cancel = self.cancellation.child_token();
288 let _guard = cancel.clone().drop_guard();
289 replace(cancel, self.gate.clone()).await?;
290
291 let wait = async {
292 loop {
293 match outcome.borrow_and_update().clone() {
294 BookRecoveryOutcome::Accepted => return Ok(()),
295 BookRecoveryOutcome::Rejected(e) => return Err(e),
296 BookRecoveryOutcome::Pending => {}
297 }
298
299 outcome
300 .changed()
301 .await
302 .map_err(|e| create_error(e.to_string()))?;
303 }
304 };
305
306 if snapshot_timeout.is_zero() {
307 wait.await
308 } else {
309 time::timeout(snapshot_timeout, wait)
310 .await
311 .unwrap_or_else(|_| Err(timeout_error()))
312 }
313 }
314}
315
316#[derive(Debug)]
318pub struct BookRecoveryState<E> {
319 recovery: Option<Arc<BookRecovery<E>>>,
320}
321
322impl<E> Default for BookRecoveryState<E> {
323 fn default() -> Self {
324 Self { recovery: None }
325 }
326}
327
328impl<E: Clone> BookRecoveryState<E> {
329 pub fn claim(&mut self) -> Option<Arc<BookRecovery<E>>> {
331 if self.is_running() {
332 return None;
333 }
334
335 self.reset();
336 let recovery = Arc::new(BookRecovery::default());
337 self.recovery = Some(Arc::clone(&recovery));
338 Some(recovery)
339 }
340
341 #[must_use]
343 pub fn current(&self) -> Option<&Arc<BookRecovery<E>>> {
344 self.recovery.as_ref()
345 }
346
347 #[must_use]
349 pub fn is_running(&self) -> bool {
350 self.recovery.as_ref().is_some_and(|r| r.is_running())
351 }
352
353 pub fn reset_on_reconnect(&mut self) {
361 if let Some(recovery) = self
362 .recovery
363 .as_ref()
364 .filter(|recovery| recovery.is_running())
365 {
366 recovery.reconnected.notify_one();
367 } else {
368 self.reset();
369 }
370 }
371
372 pub fn reset(&mut self) {
374 if let Some(recovery) = self.recovery.take() {
375 recovery.cancellation.cancel();
376 }
377 }
378}
379
380impl<E> Drop for BookRecoveryState<E> {
381 fn drop(&mut self) {
382 if let Some(recovery) = &self.recovery {
383 recovery.cancellation.cancel();
384 }
385 }
386}
387
388#[cfg(test)]
389mod tests {
390 use std::sync::atomic::{AtomicUsize, Ordering};
391
392 use nautilus_network::error::SendError;
393 use rstest::rstest;
394
395 use super::*;
396
397 #[rstest]
398 fn replacement_owner_cannot_accept_after_reset() {
399 let mut state = BookRecoveryState::<SendError>::default();
400 let old = state.claim().unwrap();
401 let duplicate = state.claim();
402 state.reset();
403 let current = state.claim().unwrap();
404
405 assert!(duplicate.is_none());
406 assert!(old.cancellation.is_cancelled());
407 assert!(!old.accept());
408 assert!(Arc::ptr_eq(state.current().unwrap(), ¤t));
409 assert!(current.accept());
410 assert!(!state.is_running());
411 }
412
413 #[rstest]
414 fn closed_gate_suppresses_snapshot_until_opened() {
415 let mut state = BookRecoveryState::<SendError>::default();
416 let recovery = state.claim().unwrap();
417 assert!(recovery.begin_replacement());
418
419 let while_closed = recovery.accept();
420 recovery.gate.open();
421 let after_open = recovery.accept();
422
423 assert!(!while_closed);
424 assert!(after_open);
425 assert!(state.claim().is_some());
426 }
427
428 #[rstest]
429 fn claim_replaces_cancelled_episode() {
430 let mut state = BookRecoveryState::<SendError>::default();
431 let first = state.claim().unwrap();
432 first.cancellation.cancel();
433
434 let second = state.claim().unwrap();
435
436 assert!(!Arc::ptr_eq(&first, &second));
437 assert!(state.is_running());
438 }
439
440 #[rstest]
441 fn reconnect_keeps_running_episode() {
442 let mut state = BookRecoveryState::<SendError>::default();
443 let recovery = state.claim().unwrap();
444 assert!(recovery.begin_replacement());
445
446 state.reset_on_reconnect();
447
448 assert!(Arc::ptr_eq(state.current().unwrap(), &recovery));
449 assert!(!recovery.cancellation.is_cancelled());
450 }
451
452 #[rstest]
453 fn reconnect_resets_cancelled_episode() {
454 let mut state = BookRecoveryState::<SendError>::default();
455 let recovery = state.claim().unwrap();
456 recovery.cancellation.cancel();
457
458 state.reset_on_reconnect();
459
460 assert!(state.current().is_none());
461 }
462
463 #[rstest]
464 fn reconnect_resets_accepted_episode() {
465 let mut state = BookRecoveryState::<SendError>::default();
466 let recovery = state.claim().unwrap();
467 assert!(recovery.accept());
468
469 state.reset_on_reconnect();
470
471 assert!(state.current().is_none());
472 }
473
474 #[tokio::test(start_paused = true)]
475 async fn missing_snapshots_spend_budget_then_retry_at_ceiling() {
476 let recovery = BookRecovery::<SendError>::default();
477 let attempts = AtomicUsize::new(0);
478 let started = time::Instant::now();
479 let attempt_started = parking_lot::Mutex::new(started);
480
481 recovery
482 .run(
483 "BOOK",
484 Duration::from_secs(1),
485 |_, gate| {
486 let attempt = attempts.fetch_add(1, Ordering::SeqCst) + 1;
487 *attempt_started.lock() = time::Instant::now();
488 gate.open();
489
490 if attempt == 10 {
492 assert!(recovery.accept());
493 }
494
495 async { Ok(()) }
496 },
497 |_| true,
498 SendError::BrokenPipe,
499 || {
500 assert_eq!(attempt_started.lock().elapsed(), Duration::from_secs(1));
501 SendError::Timeout
502 },
503 )
504 .await;
505
506 assert_eq!(attempts.load(Ordering::SeqCst), 10);
507 assert!(recovery.is_accepted());
508 assert!(started.elapsed() >= Duration::from_secs(180));
509 }
510
511 #[tokio::test(start_paused = true)]
512 async fn elapsed_budget_cancels_pending_send_and_continues() {
513 let recovery = BookRecovery::<SendError>::default();
514 let attempts = AtomicUsize::new(0);
515 let child = parking_lot::Mutex::new(None);
516
517 let result = time::timeout(
518 Duration::from_secs(200),
519 recovery.run(
520 "BOOK",
521 Duration::from_secs(1),
522 |cancel, _| {
523 attempts.fetch_add(1, Ordering::SeqCst);
524 *child.lock() = Some(cancel);
525 std::future::pending::<Result<(), SendError>>()
526 },
527 |_| true,
528 SendError::BrokenPipe,
529 || SendError::Timeout,
530 ),
531 )
532 .await;
533
534 assert!(result.is_err());
535 assert_eq!(attempts.load(Ordering::SeqCst), 1);
536 assert!(recovery.is_running());
537 assert!(child.lock().as_ref().unwrap().is_cancelled());
538 }
539
540 #[tokio::test(start_paused = true)]
541 async fn non_retryable_error_moves_to_ceiling() {
542 let recovery = BookRecovery::<SendError>::default();
543 let attempts = AtomicUsize::new(0);
544
545 let result = time::timeout(
546 Duration::from_secs(59),
547 recovery.run(
548 "BOOK",
549 Duration::from_secs(1),
550 |_, _| {
551 attempts.fetch_add(1, Ordering::SeqCst);
552 async { Err(SendError::Closed) }
553 },
554 |_| false,
555 SendError::BrokenPipe,
556 || SendError::Timeout,
557 ),
558 )
559 .await;
560
561 assert!(result.is_err());
562 assert_eq!(attempts.load(Ordering::SeqCst), 1);
563 assert!(recovery.is_running());
564 }
565
566 #[rstest]
568 #[case::disabled_snapshot_deadline(Duration::ZERO, false, 60_000)]
569 #[case::stalled_write(Duration::from_secs(10), true, 60_000)]
570 #[case::long_snapshot_deadline(Duration::from_secs(120), true, 120_000)]
571 #[tokio::test(start_paused = true)]
572 async fn ceiling_attempts_are_bounded_and_continue(
573 #[case] snapshot_timeout: Duration,
574 #[case] stall_write: bool,
575 #[case] attempt_timeout_ms: u64,
576 ) {
577 let recovery = BookRecovery::<SendError>::default();
578 let attempts = AtomicUsize::new(0);
579 let attempt_started = parking_lot::Mutex::new(time::Instant::now());
580 let timeouts = AtomicUsize::new(0);
581
582 let result = time::timeout(
583 Duration::from_secs(700),
584 recovery.run(
585 "BOOK",
586 snapshot_timeout,
587 |_, gate| {
588 attempts.fetch_add(1, Ordering::SeqCst);
589 *attempt_started.lock() = time::Instant::now();
590 let stall = stall_write;
591 async move {
592 if stall {
593 std::future::pending::<()>().await;
594 }
595
596 gate.open();
597 Ok(())
598 }
599 },
600 |_| true,
601 |message| {
602 if attempts.load(Ordering::SeqCst) > 1 {
603 assert_eq!(
604 message,
605 format!(
606 "Book recovery attempt timed out after {attempt_timeout_ms}ms (whole-attempt deadline)"
607 ),
608 );
609 assert_eq!(
610 attempt_started.lock().elapsed(),
611 Duration::from_millis(attempt_timeout_ms),
612 );
613 timeouts.fetch_add(1, Ordering::SeqCst);
614 }
615
616 SendError::BrokenPipe(message)
617 },
618 || panic!("whole-attempt timeout must not use the snapshot error factory"),
619 ),
620 )
621 .await;
622
623 assert!(result.is_err());
625 assert_eq!(attempts.load(Ordering::SeqCst), 3);
626 assert_eq!(timeouts.load(Ordering::SeqCst), 2);
627 assert!(recovery.is_running());
628 }
629
630 #[tokio::test(start_paused = true)]
631 async fn snapshot_between_ceiling_attempts_completes_run() {
632 let recovery = BookRecovery::<SendError>::default();
633 let attempts = AtomicUsize::new(0);
634 let started = time::Instant::now();
635
636 let run = recovery.run(
637 "BOOK",
638 Duration::from_secs(1),
639 |_, gate| {
640 attempts.fetch_add(1, Ordering::SeqCst);
641 gate.open();
642 async { Err(SendError::Closed) }
643 },
644 |_| false,
645 SendError::BrokenPipe,
646 || SendError::Timeout,
647 );
648
649 let late_snapshot = async {
650 time::sleep(Duration::from_secs(10)).await;
651 assert!(recovery.accept());
652 };
653
654 tokio::join!(run, late_snapshot);
655
656 assert_eq!(attempts.load(Ordering::SeqCst), 1);
657 assert!(recovery.is_accepted());
658 assert_eq!(started.elapsed(), Duration::from_secs(10));
659 }
660
661 #[tokio::test(start_paused = true)]
663 async fn reconnect_wakes_episode_waiting_at_ceiling() {
664 let mut state = BookRecoveryState::<SendError>::default();
665 let recovery = state.claim().unwrap();
666 let attempts = AtomicUsize::new(0);
667
668 let run = recovery.run(
669 "BOOK",
670 Duration::from_secs(1),
671 |_, _| {
672 attempts.fetch_add(1, Ordering::SeqCst);
673 async { Err(SendError::Closed) }
674 },
675 |_| false,
676 SendError::BrokenPipe,
677 || SendError::Timeout,
678 );
679
680 let reconnect = async {
681 time::sleep(Duration::from_secs(10)).await;
682 let before = attempts.load(Ordering::SeqCst);
683 state.reset_on_reconnect();
684 time::sleep(Duration::from_secs(1)).await;
685 let after = attempts.load(Ordering::SeqCst);
686 recovery.cancellation.cancel();
687 (before, after)
688 };
689
690 let ((), (before, after)) = tokio::join!(run, reconnect);
691
692 assert_eq!(before, 1);
693 assert_eq!(after, 2);
694 assert!(Arc::ptr_eq(state.current().unwrap(), &recovery));
695 }
696
697 #[tokio::test(start_paused = true)]
698 async fn cancellation_at_ceiling_ends_run() {
699 let recovery = BookRecovery::<SendError>::default();
700 let attempts = AtomicUsize::new(0);
701
702 let run = recovery.run(
703 "BOOK",
704 Duration::from_secs(1),
705 |_, _| {
706 attempts.fetch_add(1, Ordering::SeqCst);
707 async { Err(SendError::Closed) }
708 },
709 |_| false,
710 SendError::BrokenPipe,
711 || SendError::Timeout,
712 );
713
714 let cancel = async {
715 time::sleep(Duration::from_secs(30)).await;
716 recovery.cancellation.cancel();
717 };
718
719 tokio::join!(run, cancel);
720
721 assert_eq!(attempts.load(Ordering::SeqCst), 1);
722 assert!(!recovery.is_running());
723 assert!(!recovery.is_accepted());
724 }
725
726 #[tokio::test(start_paused = true)]
727 async fn snapshot_during_send_completes_without_retry() {
728 let recovery = BookRecovery::<SendError>::default();
729 let attempts = AtomicUsize::new(0);
730
731 recovery
732 .run(
733 "BOOK",
734 Duration::from_secs(1),
735 |_, gate| {
736 attempts.fetch_add(1, Ordering::SeqCst);
737 assert!(!recovery.accept());
738 gate.open();
739 assert!(recovery.accept());
740 async { Ok(()) }
741 },
742 |_| true,
743 SendError::BrokenPipe,
744 || SendError::Timeout,
745 )
746 .await;
747
748 assert_eq!(attempts.load(Ordering::SeqCst), 1);
749 assert!(recovery.is_accepted());
750 }
751
752 #[tokio::test(start_paused = true)]
753 async fn dropping_owner_cancels_pending_attempt() {
754 let mut state = BookRecoveryState::<SendError>::default();
755 let recovery = state.claim().unwrap();
756 let child = parking_lot::Mutex::new(None);
757
758 let operation = recovery.run(
759 "BOOK",
760 Duration::ZERO,
761 |cancel, _| {
762 *child.lock() = Some(cancel);
763 std::future::pending::<Result<(), SendError>>()
764 },
765 |_| true,
766 SendError::BrokenPipe,
767 || SendError::Timeout,
768 );
769
770 tokio::pin!(operation);
771 tokio::select! {
772 () = &mut operation => panic!("unexpected completion"),
773 () = tokio::task::yield_now() => {},
774 }
775 drop(state);
776 operation.await;
777
778 assert!(child.lock().as_ref().unwrap().is_cancelled());
779 }
780}