1use std::{future::Future, sync::Arc};
43
44use nautilus_common::live::dst::time::{self, Duration};
45use nautilus_network::retry::{RetryConfig, RetryManager};
46use tokio_util::sync::CancellationToken;
47
48use super::snapshot::SnapshotGate;
49
50const ATTEMPTS_MAX: u32 = 8;
52const ELAPSED_MAX_MS: u64 = 180_000;
53const OPERATION_TIMEOUT_MS: Option<u64> = None;
54
55const RETRY_DELAY_INITIAL_MS: u64 = 1_000;
56const RETRY_DELAY_MAX_MS: u64 = 10_000;
57const RETRY_BACKOFF_FACTOR: f64 = 2.0;
58const RETRY_JITTER_MAX_MS: u64 = 1_000;
59const RETRY_FIRST_IMMEDIATE: bool = true;
60
61#[derive(Debug, Clone)]
63pub enum BookRecoveryOutcome<E> {
64 Pending,
65 Accepted,
66 Rejected(E),
67}
68
69#[derive(Debug)]
71pub struct BookRecovery<E> {
72 pub cancellation: CancellationToken,
73 pub outcome: tokio::sync::watch::Sender<BookRecoveryOutcome<E>>,
74 pub gate: SnapshotGate,
75}
76
77impl<E: Clone> Default for BookRecovery<E> {
78 fn default() -> Self {
79 Self {
80 cancellation: CancellationToken::new(),
81 outcome: tokio::sync::watch::channel(BookRecoveryOutcome::Pending).0,
82 gate: SnapshotGate::default(),
83 }
84 }
85}
86
87impl<E: Clone> BookRecovery<E> {
88 #[must_use]
90 pub fn begin_replacement(&self) -> bool {
91 let mut gate = self.gate.lock();
92
93 if self.cancellation.is_cancelled() || self.is_accepted() {
94 return false;
95 }
96
97 gate.close();
98 true
99 }
100
101 #[must_use]
103 pub fn is_accepted(&self) -> bool {
104 matches!(*self.outcome.borrow(), BookRecoveryOutcome::Accepted)
105 }
106
107 pub fn accept(&self) -> bool {
109 let gate = self.gate.lock();
110 if gate.is_closed() || (self.cancellation.is_cancelled() && !self.is_accepted()) {
111 return false;
112 }
113
114 self.outcome.send_replace(BookRecoveryOutcome::Accepted);
115 true
116 }
117}
118
119impl<E: Clone + std::error::Error> BookRecovery<E> {
120 pub async fn run<F, Fut>(
131 &self,
132 snapshot_timeout: Duration,
133 replace: F,
134 should_retry: impl Fn(&E) -> bool,
135 create_error: impl Fn(String) -> E,
136 timeout_error: impl Fn() -> E,
137 ) -> Result<(), E>
138 where
139 F: Fn(CancellationToken, SnapshotGate) -> Fut,
140 Fut: Future<Output = Result<(), E>>,
141 {
142 let manager = RetryManager::<E>::new(RetryConfig {
143 max_retries: ATTEMPTS_MAX - 1,
144 initial_delay_ms: RETRY_DELAY_INITIAL_MS,
145 max_delay_ms: RETRY_DELAY_MAX_MS,
146 backoff_factor: RETRY_BACKOFF_FACTOR,
147 jitter_ms: RETRY_JITTER_MAX_MS,
148 immediate_first: RETRY_FIRST_IMMEDIATE,
149 operation_timeout_ms: OPERATION_TIMEOUT_MS,
150 max_elapsed_ms: Some(ELAPSED_MAX_MS),
151 });
152
153 manager
154 .invocation(
155 "book recovery",
156 || async {
157 let mut outcome = self.outcome.subscribe();
158
159 if !self.begin_replacement() {
160 return Ok(());
161 }
162
163 self.outcome.send_if_modified(|outcome| {
164 if matches!(outcome, BookRecoveryOutcome::Rejected(_)) {
165 *outcome = BookRecoveryOutcome::Pending;
166 true
167 } else {
168 false
169 }
170 });
171
172 let cancel = self.cancellation.child_token();
173 let _guard = cancel.clone().drop_guard();
174 replace(cancel, self.gate.clone()).await?;
175
176 let wait = async {
177 loop {
178 match outcome.borrow_and_update().clone() {
179 BookRecoveryOutcome::Accepted => return Ok(()),
180 BookRecoveryOutcome::Rejected(e) => return Err(e),
181 BookRecoveryOutcome::Pending => {}
182 }
183
184 outcome
185 .changed()
186 .await
187 .map_err(|e| create_error(e.to_string()))?;
188 }
189 };
190
191 if snapshot_timeout.is_zero() {
192 wait.await
193 } else {
194 time::timeout(snapshot_timeout, wait)
195 .await
196 .unwrap_or_else(|_| Err(timeout_error()))
197 }
198 },
199 should_retry,
200 |e| create_error(e.to_string()),
201 )
202 .cancellation_token(&self.cancellation)
203 .execute()
204 .await
205 }
206}
207
208#[derive(Debug)]
210pub struct BookRecoveryState<E> {
211 recovery: Option<Arc<BookRecovery<E>>>,
212 failed: bool,
213}
214
215impl<E> Default for BookRecoveryState<E> {
216 fn default() -> Self {
217 Self {
218 recovery: None,
219 failed: false,
220 }
221 }
222}
223
224impl<E: Clone> BookRecoveryState<E> {
225 pub fn claim(&mut self) -> Option<Arc<BookRecovery<E>>> {
227 if self.failed || self.recovery.as_ref().is_some_and(|r| !r.is_accepted()) {
228 return None;
229 }
230
231 self.reset();
232 let recovery = Arc::new(BookRecovery::default());
233 self.recovery = Some(Arc::clone(&recovery));
234 Some(recovery)
235 }
236
237 #[must_use]
239 pub fn current(&self) -> Option<&Arc<BookRecovery<E>>> {
240 self.recovery.as_ref()
241 }
242
243 #[must_use]
245 pub fn is_failed(&self) -> bool {
246 self.failed
247 }
248
249 pub fn fail(&mut self, owner: Option<&Arc<BookRecovery<E>>>) -> bool {
253 if let Some(owner) = owner
254 && (!self
255 .recovery
256 .as_ref()
257 .is_some_and(|current| Arc::ptr_eq(current, owner))
258 || owner.is_accepted())
259 {
260 return false;
261 }
262
263 self.reset();
264 self.failed = true;
265 true
266 }
267
268 pub fn reset(&mut self) {
270 if let Some(recovery) = self.recovery.take() {
271 recovery.cancellation.cancel();
272 }
273
274 self.failed = false;
275 }
276}
277
278impl<E> Drop for BookRecoveryState<E> {
279 fn drop(&mut self) {
280 if let Some(recovery) = &self.recovery {
281 recovery.cancellation.cancel();
282 }
283 }
284}
285
286#[cfg(test)]
287mod tests {
288 use std::sync::atomic::{AtomicUsize, Ordering};
289
290 use nautilus_network::error::SendError;
291 use rstest::rstest;
292
293 use super::*;
294
295 #[rstest]
296 fn replacement_owner_cannot_fail_or_accept_after_reset() {
297 let mut state = BookRecoveryState::<SendError>::default();
298 let old = state.claim().unwrap();
299 assert!(state.claim().is_none());
300 state.reset();
301 let current = state.claim().unwrap();
302
303 assert!(old.cancellation.is_cancelled());
304 assert!(!old.accept());
305 assert!(!state.fail(Some(&old)));
306 assert!(!state.is_failed());
307 assert!(Arc::ptr_eq(state.current().unwrap(), ¤t));
308 assert!(current.accept());
309 assert!(!state.fail(Some(¤t)));
310 }
311
312 #[rstest]
313 fn gate_and_failure_suppress_snapshots_until_explicit_restart() {
314 let mut state = BookRecoveryState::<SendError>::default();
315 let recovery = state.claim().unwrap();
316 assert!(recovery.begin_replacement());
317 assert!(!recovery.accept());
318 assert!(state.fail(Some(&recovery)));
319 recovery.gate.open();
320
321 assert!(!recovery.accept());
322 assert!(state.is_failed());
323 assert!(state.claim().is_none());
324 state.reset();
325 assert!(state.claim().is_some());
326 }
327
328 #[tokio::test(start_paused = true)]
329 async fn missing_snapshots_exhaust_exactly_eight_attempts() {
330 let recovery = BookRecovery::<SendError>::default();
331 let attempts = AtomicUsize::new(0);
332
333 let result = recovery
334 .run(
335 Duration::from_secs(1),
336 |_, gate| {
337 attempts.fetch_add(1, Ordering::SeqCst);
338 gate.open();
339 async { Ok(()) }
340 },
341 |_| true,
342 SendError::BrokenPipe,
343 || SendError::Timeout,
344 )
345 .await;
346
347 assert!(result.is_err());
348 assert_eq!(attempts.load(Ordering::SeqCst), 8);
349 assert!(!recovery.is_accepted());
350 }
351
352 #[tokio::test(start_paused = true)]
353 async fn elapsed_budget_cancels_pending_send() {
354 let recovery = BookRecovery::<SendError>::default();
355 let attempts = AtomicUsize::new(0);
356 let child = parking_lot::Mutex::new(None);
357 let started = time::Instant::now();
358
359 let result = recovery
360 .run(
361 Duration::from_secs(1),
362 |cancel, _| {
363 attempts.fetch_add(1, Ordering::SeqCst);
364 *child.lock() = Some(cancel);
365 std::future::pending::<Result<(), SendError>>()
366 },
367 |_| true,
368 SendError::BrokenPipe,
369 || SendError::Timeout,
370 )
371 .await;
372
373 assert_eq!(
374 result.unwrap_err().to_string(),
375 "send failed: broken pipe (Retry budget exceeded (1/8))"
376 );
377 assert_eq!(started.elapsed(), Duration::from_secs(180));
378 assert_eq!(attempts.load(Ordering::SeqCst), 1);
379 assert!(child.lock().as_ref().unwrap().is_cancelled());
380 }
381
382 #[tokio::test(start_paused = true)]
383 async fn snapshot_during_send_completes_without_retry() {
384 let recovery = BookRecovery::<SendError>::default();
385 let attempts = AtomicUsize::new(0);
386
387 let result = recovery
388 .run(
389 Duration::from_secs(1),
390 |_, gate| {
391 attempts.fetch_add(1, Ordering::SeqCst);
392 assert!(!recovery.accept());
393 gate.open();
394 assert!(recovery.accept());
395 async { Ok(()) }
396 },
397 |_| true,
398 SendError::BrokenPipe,
399 || SendError::Timeout,
400 )
401 .await;
402
403 assert!(result.is_ok());
404 assert_eq!(attempts.load(Ordering::SeqCst), 1);
405 assert!(recovery.is_accepted());
406 }
407
408 #[tokio::test(start_paused = true)]
409 async fn dropping_owner_cancels_pending_attempt() {
410 let mut state = BookRecoveryState::<SendError>::default();
411 let recovery = state.claim().unwrap();
412 let child = parking_lot::Mutex::new(None);
413
414 let operation = recovery.run(
415 Duration::ZERO,
416 |cancel, _| {
417 *child.lock() = Some(cancel);
418 std::future::pending::<Result<(), SendError>>()
419 },
420 |_| true,
421 SendError::BrokenPipe,
422 || SendError::Timeout,
423 );
424
425 tokio::pin!(operation);
426 tokio::select! {
427 result = &mut operation => panic!("unexpected completion: {result:?}"),
428 () = tokio::task::yield_now() => {},
429 }
430 drop(state);
431 assert!(operation.await.is_err());
432 assert!(child.lock().as_ref().unwrap().is_cancelled());
433 }
434}