zebra_chain/diagnostic/task/
future.rs

Help
1//! Diagnostic types and functions for Zebra async future tasks:
2//! - task handles
3//! - errors and panics
4
5use std::{future, panic};
6
7use futures::future::{BoxFuture, FutureExt};
8use tokio::task::{JoinError, JoinHandle};
9
10use crate::shutdown::is_shutting_down;
11
12use super::{CheckForPanics, WaitForPanics};
13
14/// This is the return type of the [`JoinHandle`] future.
15impl<T> CheckForPanics for Result<T, JoinError> {
16    /// The [`JoinHandle`]'s task output, after resuming any panics,
17    /// and ignoring task cancellations on shutdown.
18    type Output = Result<T, JoinError>;
19
20    /// Returns the task result if the task finished normally.
21    /// Otherwise, resumes any panics, and ignores any expected errors.
22    /// Handles unexpected errors based on `panic_on_unexpected_termination`.
23    ///
24    /// If the task finished normally, returns `Some(T)`.
25    /// If the task was cancelled, returns `None`.
26    #[track_caller]
27    fn check_for_panics_with(self, panic_on_unexpected_termination: bool) -> Self::Output {
28        match self {
29            Ok(task_output) => Ok(task_output),
30            Err(join_error) => {
31                Err(join_error.check_for_panics_with(panic_on_unexpected_termination))
32            }
33        }
34    }
35}
36
37impl CheckForPanics for JoinError {
38    /// The [`JoinError`] after resuming any panics, and logging any unexpected task cancellations.
39    type Output = JoinError;
40
41    /// Resume any panics and panic on unexpected task cancellations.
42    /// Always returns [`JoinError::Cancelled`](JoinError::is_cancelled).
43    #[track_caller]
44    fn check_for_panics_with(self, panic_on_unexpected_termination: bool) -> Self::Output {
45        match self.try_into_panic() {
46            Ok(panic_payload) => panic::resume_unwind(panic_payload),
47
48            // We could ignore this error, but then we'd have to change the return type.
49            Err(task_cancelled) => {
50                if !panic_on_unexpected_termination {
51                    debug!(?task_cancelled, "ignoring expected task termination");
52
53                    task_cancelled
54                } else if is_shutting_down() {
55                    debug!(
56                        ?task_cancelled,
57                        "ignoring task termination because Zebra is shutting down"
58                    );
59
60                    task_cancelled
61                } else {
62                    panic!("task unexpectedly exited with: {:?}", task_cancelled)
63                }
64            }
65        }
66    }
67}
68
69impl<T> WaitForPanics for JoinHandle<T>
70where
71    T: Send + 'static,
72{
73    type Output = BoxFuture<'static, T>;
74
75    /// Returns a future which waits for `self` to finish, then checks if its output is:
76    /// - a panic payload: resume that panic,
77    /// - an unexpected termination:
78    ///   - if `panic_on_unexpected_termination` is true, panic with that error,
79    ///   - otherwise, hang waiting for shutdown,
80    /// - an expected termination: hang waiting for shutdown.
81    ///
82    /// Otherwise, returns the task return value of `self`.
83    ///
84    /// # Panics
85    ///
86    /// If `self` contains a panic payload, or [`JoinHandle::abort()`] has been called on `self`.
87    ///
88    /// # Hangs
89    ///
90    /// If `self` contains an expected termination, and we're shutting down anyway.
91    /// If we're ignoring terminations because `panic_on_unexpected_termination` is `false`.
92    /// Futures hang by returning `Pending` and not setting a waker, so this uses minimal resources.
93    #[track_caller]
94    fn wait_for_panics_with(self, panic_on_unexpected_termination: bool) -> Self::Output {
95        async move {
96            match self
97                .await
98                .check_for_panics_with(panic_on_unexpected_termination)
99            {
100                Ok(task_output) => task_output,
101                Err(_expected_cancel_error) => future::pending().await,
102            }
103        }
104        .boxed()
105    }
106}
zebra_chain/diagnostic/task/future.rs

zebra_chain/diagnostic/task/
future.rs
