use std::{future, panic};
use futures::future::{BoxFuture, FutureExt};
use tokio::task::{JoinError, JoinHandle};
use crate::shutdown::is_shutting_down;
use super::{CheckForPanics, WaitForPanics};
impl<T> CheckForPanics for Result<T, JoinError> {
type Output = Result<T, JoinError>;
#[track_caller]
fn check_for_panics_with(self, panic_on_unexpected_termination: bool) -> Self::Output {
match self {
Ok(task_output) => Ok(task_output),
Err(join_error) => {
Err(join_error.check_for_panics_with(panic_on_unexpected_termination))
}
}
}
}
impl CheckForPanics for JoinError {
type Output = JoinError;
#[track_caller]
fn check_for_panics_with(self, panic_on_unexpected_termination: bool) -> Self::Output {
match self.try_into_panic() {
Ok(panic_payload) => panic::resume_unwind(panic_payload),
Err(task_cancelled) => {
if !panic_on_unexpected_termination {
debug!(?task_cancelled, "ignoring expected task termination");
task_cancelled
} else if is_shutting_down() {
debug!(
?task_cancelled,
"ignoring task termination because Zebra is shutting down"
);
task_cancelled
} else {
panic!("task unexpectedly exited with: {task_cancelled:?}")
}
}
}
}
}
impl<T> WaitForPanics for JoinHandle<T>
where
T: Send + 'static,
{
type Output = BoxFuture<'static, T>;
#[track_caller]
fn wait_for_panics_with(self, panic_on_unexpected_termination: bool) -> Self::Output {
async move {
match self
.await
.check_for_panics_with(panic_on_unexpected_termination)
{
Ok(task_output) => task_output,
Err(_expected_cancel_error) => future::pending().await,
}
}
.boxed()
}
}