zebra_chain/diagnostic/task/
thread.rs

1//! Diagnostic types and functions for Zebra OS thread tasks:
2//! - task handles
3//! - errors and panics
4
5use std::{
6    panic,
7    sync::Arc,
8    thread::{self, JoinHandle},
9};
10
11use crate::shutdown::is_shutting_down;
12
13use super::{CheckForPanics, WaitForPanics};
14
15impl<T> CheckForPanics for thread::Result<T>
16where
17    T: std::fmt::Debug,
18{
19    type Output = T;
20
21    /// # Panics
22    ///
23    /// - if the thread panicked.
24    /// - if the thread is cancelled, `panic_on_unexpected_termination` is true, and
25    ///   Zebra is not shutting down.
26    ///
27    /// Threads can't be cancelled except by using a panic, so there are no thread errors here.
28    /// `panic_on_unexpected_termination` is
29    #[track_caller]
30    fn check_for_panics_with(self, panic_on_unexpected_termination: bool) -> Self::Output {
31        match self {
32            // The value returned by the thread when it finished.
33            Ok(thread_output) => {
34                if !panic_on_unexpected_termination {
35                    debug!(?thread_output, "ignoring expected thread exit");
36
37                    thread_output
38                } else if is_shutting_down() {
39                    debug!(
40                        ?thread_output,
41                        "ignoring thread exit because Zebra is shutting down"
42                    );
43
44                    thread_output
45                } else {
46                    panic!("thread unexpectedly exited with: {:?}", thread_output)
47                }
48            }
49
50            // A thread error is always a panic.
51            Err(panic_payload) => panic::resume_unwind(panic_payload),
52        }
53    }
54}
55
56impl<T> WaitForPanics for JoinHandle<T>
57where
58    T: std::fmt::Debug,
59{
60    type Output = T;
61
62    /// Waits for the thread to finish, then panics if the thread panicked.
63    #[track_caller]
64    fn wait_for_panics_with(self, panic_on_unexpected_termination: bool) -> Self::Output {
65        self.join()
66            .check_for_panics_with(panic_on_unexpected_termination)
67    }
68}
69
70impl<T> WaitForPanics for Arc<JoinHandle<T>>
71where
72    T: std::fmt::Debug,
73{
74    type Output = Option<T>;
75
76    /// If this is the final `Arc`, waits for the thread to finish, then panics if the thread
77    /// panicked. Otherwise, returns the thread's return value.
78    ///
79    /// If this is not the final `Arc`, drops the handle and immediately returns `None`.
80    #[track_caller]
81    fn wait_for_panics_with(self, panic_on_unexpected_termination: bool) -> Self::Output {
82        // If we are the last Arc with a reference to this handle,
83        // we can wait for it and propagate any panics.
84        //
85        // We use into_inner() because it guarantees that exactly one of the tasks gets the
86        // JoinHandle. try_unwrap() lets us keep the JoinHandle, but it can also miss panics.
87        //
88        // This is more readable as an expanded statement.
89        #[allow(clippy::manual_map)]
90        if let Some(handle) = Arc::into_inner(self) {
91            Some(handle.wait_for_panics_with(panic_on_unexpected_termination))
92        } else {
93            None
94        }
95    }
96}
97
98impl<T> CheckForPanics for &mut Option<Arc<JoinHandle<T>>>
99where
100    T: std::fmt::Debug,
101{
102    type Output = Option<T>;
103
104    /// If this is the final `Arc`, checks if the thread has finished, then panics if the thread
105    /// panicked. Otherwise, returns the thread's return value.
106    ///
107    /// If the thread has not finished, or this is not the final `Arc`, returns `None`.
108    #[track_caller]
109    fn check_for_panics_with(self, panic_on_unexpected_termination: bool) -> Self::Output {
110        let handle = self.take()?;
111
112        if handle.is_finished() {
113            // This is the same as calling `self.wait_for_panics()`, but we can't do that,
114            // because we've taken `self`.
115            #[allow(clippy::manual_map)]
116            return handle.wait_for_panics_with(panic_on_unexpected_termination);
117        }
118
119        *self = Some(handle);
120
121        None
122    }
123}
124
125impl<T> WaitForPanics for &mut Option<Arc<JoinHandle<T>>>
126where
127    T: std::fmt::Debug,
128{
129    type Output = Option<T>;
130
131    /// If this is the final `Arc`, waits for the thread to finish, then panics if the thread
132    /// panicked. Otherwise, returns the thread's return value.
133    ///
134    /// If this is not the final `Arc`, drops the handle and returns `None`.
135    #[track_caller]
136    fn wait_for_panics_with(self, panic_on_unexpected_termination: bool) -> Self::Output {
137        // This is more readable as an expanded statement.
138        #[allow(clippy::manual_map)]
139        if let Some(output) = self
140            .take()?
141            .wait_for_panics_with(panic_on_unexpected_termination)
142        {
143            Some(output)
144        } else {
145            // Some other task has a reference, so we should give up ours to let them use it.
146            None
147        }
148    }
149}