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 replacement attempts for one book.
17//!
18//! - [`BookRecoveryState`] admits one running recovery per book and cancels obsolete work.
19//! - [`BookRecovery`] runs replacement attempts until the adapter accepts a snapshot or cancels the
20//!   episode. A book never ends in a terminal failure state.
21//! - [`BookRecoveryOutcome`] carries pending, accepted, or rejected results from the adapter's
22//!   frame handling to the recovery runner.
23//!
24//! # Recovery Lifecycle
25//!
26//! The adapter claims an episode before starting work and supplies the replacement operation,
27//! retry classification, and error construction. Each attempt receives a child cancellation token
28//! and a closed snapshot gate. The adapter opens the gate after the intended connection confirms
29//! the write, then accepts a valid snapshot through [`BookRecovery::accept`]. Write completion
30//! alone never completes recovery.
31//!
32//! Attempts back off within a retry budget of eight attempts in 180 seconds. An exhausted budget,
33//! or an error the adapter classifies as not retryable, moves the episode to attempts at a jittered
34//! interval that doubles from one minute to fifteen minutes, which bounds subscription traffic for
35//! a book that keeps failing. Each such attempt is bounded by `max(snapshot_timeout, 60 seconds)`,
36//! so a stalled write or a disabled snapshot deadline cannot end the retries. The runner logs the
37//! transition as an error, each later failed attempt as a warning, and completion as info.
38//!
39//! # Adapters
40//!
41//! The adapter serializes claims and snapshot acceptance under its state lock or owning task. A
42//! reconnect keeps a running episode, whose replacement may be in flight, so the reconnect can
43//! neither replenish its budget nor abandon a write halfway; it wakes only a wait between attempts
44//! after the budget, as [`BookRecoveryState::reset_on_reconnect`] describes. Removing or resetting
45//! a [`BookRecoveryState`] cancels its episode; dropping an attempt cancels that attempt's child
46//! token.
47//!
48//! Task spawning, subscription correlation, and book cache updates remain adapter-owned. Claim an
49//! episode only once the replacement has what it needs, and cancel the episode if its task cannot
50//! start: a running episode without a task keeps its book owned with nothing running.
51
52use 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
63// Includes the initial attempt
64const 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
74// Attempts after the retry budget, until acceptance or cancellation
75const 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/// Outcome published by an adapter's snapshot and rejection handling.
80#[derive(Debug, Clone)]
81pub enum BookRecoveryOutcome<E> {
82    Pending,
83    Accepted,
84    Rejected(E),
85}
86
87/// One recovery episode, running until a snapshot is accepted or the episode is cancelled.
88#[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    /// Closes snapshot acceptance, returning `false` if this episode has already ended.
109    #[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    /// Returns whether the adapter has accepted a snapshot for this episode.
122    #[must_use]
123    pub fn is_accepted(&self) -> bool {
124        matches!(*self.outcome.borrow(), BookRecoveryOutcome::Accepted)
125    }
126
127    /// Returns whether this episode still owns its book: not accepted and not cancelled.
128    #[must_use]
129    pub fn is_running(&self) -> bool {
130        !self.is_accepted() && !self.cancellation.is_cancelled()
131    }
132
133    /// Returns `true` when the gate permits accepting a snapshot for this episode.
134    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    /// Replaces the subscription and waits for a snapshot until one is accepted or the episode is
147    /// cancelled.
148    ///
149    /// Each send receives a child cancellation token. Dropping an attempt cancels queued
150    /// transport work. Keeping this future alive across reconnects preserves the retry budget.
151    /// A zero snapshot timeout disables only the individual snapshot deadline: the budget still
152    /// ends a pending attempt, and each attempt after it is bounded by
153    /// `max(snapshot_timeout, 60 seconds)`. Errors that `should_retry` rejects skip the rest of
154    /// the budget.
155    ///
156    /// Whole-attempt deadlines use `create_error`; `timeout_error` handles snapshot deadlines.
157    #[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            // Bounds a stalled write and a disabled snapshot deadline alike
207            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                    // A reconnect ends the wait so the next attempt uses the new connection
222                    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            // A snapshot accepted between attempts ends the wait; no replacement write can be in
247            // flight then, since acceptance requires an open gate.
248            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/// Owns a book's recovery episode under its adapter's state lock.
317#[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    /// Claims one recovery episode, refusing while another is running.
330    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    /// Returns the current episode.
342    #[must_use]
343    pub fn current(&self) -> Option<&Arc<BookRecovery<E>>> {
344        self.recovery.as_ref()
345    }
346
347    /// Returns whether a running episode owns the book.
348    #[must_use]
349    pub fn is_running(&self) -> bool {
350        self.recovery.as_ref().is_some_and(|r| r.is_running())
351    }
352
353    /// Keeps a running episode, whose replacement may be in flight, and cancels any other so the
354    /// next claim starts fresh.
355    ///
356    /// The kept episode keeps its budget. The reconnect never interrupts an attempt or shortens a
357    /// backoff inside the budget. An episode waiting between attempts after its budget attempts
358    /// again at once; a reconnect during an attempt or before the budget runs out ends the next
359    /// such wait instead.
360    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    /// Cancels the current episode.
373    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(), &current));
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                    // The tenth attempt is the second at the ceiling interval
491                    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    // Neither a disabled snapshot deadline nor a write that never completes may end the retries
567    #[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        // The budget ends at 180s; two bounded ceiling attempts finish before the 700s timeout
624        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    // A reconnect ends the wait at the ceiling so the kept episode retries on the new connection
662    #[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}