stratum_apps/
fallback_coordinator.rs

1use std::future::Future;
2use tokio_util::{sync::CancellationToken, task::TaskTracker};
3
4/// Coordinates fallback across the components of a runtime generation.
5///
6/// Fallback is signalled through a [`CancellationToken`] and acknowledged by lifetime: a future
7/// wrapped with [`FallbackCoordinator::track`] counts as pending until it is dropped, whether it
8/// returned, was cancelled or panicked, so there is nothing a component has to remember to call
9/// on its way out. [`FallbackCoordinator::trigger_fallback_and_wait`] cancels the token and
10/// resolves once no tracked future is left, which it checks rather than infers from a wakeup: a
11/// future tracked after the trigger, or after every other one already finished, is still waited
12/// for.
13///
14/// This is meant to be used together with [`crate::task_manager::TaskManager`]:
15///
16/// ```ignore
17/// task_manager.spawn(fallback_coordinator.track(async move {
18///     tokio::select! {
19///         _ = fallback_token.cancelled() => { /* clean up */ }
20///         // ...
21///     }
22/// }));
23/// ```
24///
25/// A coordinator serves one fallback cycle: once the wait returned, the runtime instantiates a
26/// fresh one for the next generation.
27#[derive(Debug, Clone, Default)]
28pub struct FallbackCoordinator {
29    signal: CancellationToken,
30    tasks: TaskTracker,
31}
32
33impl FallbackCoordinator {
34    pub fn new() -> Self {
35        Self::default()
36    }
37
38    /// Get the cancellation token that signals fallback.
39    pub fn token(&self) -> CancellationToken {
40        self.signal.clone()
41    }
42
43    /// Wrap a fallback-relevant future so that fallback waits for it until it is dropped.
44    ///
45    /// Tracking starts here, not when the future is first polled, so a tracked future handed to a
46    /// spawner is pending from the moment it exists. Track before fallback is triggered or from
47    /// inside a future that is itself tracked: a future tracked after the wait has observed the
48    /// tracker empty belongs to no generation and is not waited for.
49    pub fn track<F: Future>(&self, fut: F) -> impl Future<Output = F::Output> + use<F> {
50        self.tasks.track_future(fut)
51    }
52
53    /// Trigger fallback and wait until every tracked future has been dropped.
54    pub async fn trigger_fallback_and_wait(&self) {
55        tracing::debug!(
56            pending = self.tasks.len(),
57            "FallbackCoordinator: triggering fallback"
58        );
59        self.signal.cancel();
60        self.tasks.close();
61        // `wait` resolves once the tracker has been empty at some point since the call, so work
62        // tracked after that point would not be covered: look again, and wait again if there is
63        // any. `wait` re-reads the closed-and-empty state when it is created and on every poll,
64        // so a drop between the check and the next `wait` is seen by that `wait`, not missed.
65        loop {
66            self.tasks.wait().await;
67            if self.tasks.is_empty() {
68                break;
69            }
70            tracing::debug!(
71                pending = self.tasks.len(),
72                "FallbackCoordinator: work tracked after the tracker emptied, waiting again"
73            );
74        }
75        tracing::debug!("FallbackCoordinator: all components finished fallback cleanup");
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82    use futures::FutureExt as _;
83    use std::pin::pin;
84
85    #[tokio::test]
86    async fn waits_for_work_tracked_after_an_earlier_one_finished() {
87        let coordinator = FallbackCoordinator::new();
88        let first = coordinator.track(std::future::pending::<()>());
89        let mut wait = pin!(coordinator.trigger_fallback_and_wait());
90        assert!(wait.as_mut().now_or_never().is_none());
91
92        drop(first);
93        let second = coordinator.track(std::future::pending::<()>());
94        assert!(
95            wait.as_mut().now_or_never().is_none(),
96            "fallback completed while a component was still pending"
97        );
98
99        drop(second);
100        assert!(wait.as_mut().now_or_never().is_some());
101    }
102
103    #[tokio::test]
104    async fn tracked_future_counts_from_creation_until_dropped() {
105        let coordinator = FallbackCoordinator::new();
106        let never_polled = coordinator.track(std::future::pending::<()>());
107        let mut wait = pin!(coordinator.trigger_fallback_and_wait());
108        assert!(wait.as_mut().now_or_never().is_none());
109
110        drop(never_polled);
111        assert!(wait.as_mut().now_or_never().is_some());
112    }
113
114    #[tokio::test]
115    async fn panicking_tracked_task_is_released() {
116        let coordinator = FallbackCoordinator::new();
117        let task = tokio::spawn(coordinator.track(async { panic!("task failed") }));
118        assert!(task.await.unwrap_err().is_panic());
119
120        coordinator.trigger_fallback_and_wait().await;
121    }
122}
stratum_apps/fallback_coordinator.rs

stratum_apps/
fallback_coordinator.rs
