Skip to main content

nautilus_live/book/
recovery.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! Recovery ownership and bounded replacement attempts for one book.
17//!
18//! - [`BookRecoveryState`] admits one recovery owner, rejects stale failure reports, and cancels
19//!   obsolete work. Terminal failure suppresses new claims until the adapter resets the state.
20//! - [`BookRecovery`] runs replacement attempts, waits for an accepted snapshot, and applies
21//!   backoff, attempt limits, and a total elapsed-time budget.
22//! - [`BookRecoveryOutcome`] carries pending, accepted, or rejected results from the adapter's
23//!   frame handling to the recovery runner.
24//!
25//! # Recovery Lifecycle
26//!
27//! The adapter claims an episode before starting work and supplies the replacement operation,
28//! retry classification, and error construction. Each attempt receives a child cancellation token
29//! and a closed snapshot gate. The adapter opens the gate after the intended connection confirms
30//! the write, then accepts a valid snapshot through [`BookRecovery::accept`]. Write completion
31//! alone never completes recovery.
32//!
33//! # Adapters
34//!
35//! The adapter serializes claims, snapshot acceptance, and failure reporting under its state lock
36//! or owning task. It retains the same episode across reconnects to preserve the remaining budget.
37//! Removing or resetting its [`BookRecoveryState`] cancels the episode; dropping an attempt cancels
38//! that attempt's child token.
39//!
40//! Task spawning, subscription correlation, and book cache updates remain adapter-owned.
41
42use 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
50// Includes the initial attempt
51const 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/// Outcome published by an adapter's snapshot and rejection handling.
62#[derive(Debug, Clone)]
63pub enum BookRecoveryOutcome<E> {
64    Pending,
65    Accepted,
66    Rejected(E),
67}
68
69/// One recovery episode, retained across reconnects until a snapshot or terminal failure.
70#[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    /// Closes snapshot acceptance, returning `false` if this episode has already ended.
89    #[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    /// Returns whether the adapter has accepted a snapshot for this episode.
102    #[must_use]
103    pub fn is_accepted(&self) -> bool {
104        matches!(*self.outcome.borrow(), BookRecoveryOutcome::Accepted)
105    }
106
107    /// Returns `true` when the gate permits accepting a snapshot for this episode.
108    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    /// Replaces the subscription and waits for a snapshot, with bounded retries.
121    ///
122    /// Each send receives a child cancellation token. Dropping an attempt cancels queued
123    /// transport work. Keeping this future alive across reconnects preserves the attempt limit
124    /// and total elapsed-time budget.
125    /// A zero snapshot timeout disables only the individual snapshot deadline.
126    ///
127    /// # Errors
128    ///
129    /// Returns the terminal adapter error, cancellation, or exhausted retry budget.
130    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/// Owns a book's recovery and terminal suppression under its adapter's state lock.
209#[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    /// Claims one recovery episode, refusing duplicate or failed work.
226    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    /// Returns the current episode.
238    #[must_use]
239    pub fn current(&self) -> Option<&Arc<BookRecovery<E>>> {
240        self.recovery.as_ref()
241    }
242
243    /// Returns whether exhausted or permanent failure suppresses book output.
244    #[must_use]
245    pub fn is_failed(&self) -> bool {
246        self.failed
247    }
248
249    /// Returns `true` when failure is recorded for the current, incomplete owner.
250    ///
251    /// A stale or accepted owner returns `false`. Passing `None` records failure unconditionally.
252    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    /// Cancels obsolete work and clears terminal suppression for an explicit restart.
269    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(), &current));
308        assert!(current.accept());
309        assert!(!state.fail(Some(&current)));
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}