Simplify trait to a single method

zebra-chain/src/diagnostic/task.rs

@@ -24,15 +24,9 @@ pub trait CheckForPanics {
 }
 
 /// A trait that waits for a task to finish, then handles panics and cancellations.
-pub trait WaitForTermination {
+pub trait WaitForPanics {
     /// The underlying task output, after removing panics and unwrapping termination results.
-    type UnwrappedOutput;
-
-    /// The underlying task output, after removing panics,
-    /// and replacing expected terminations with a default value.
-    ///
-    /// This is typically a `Result` containing the underlying task output type.
-    type ResultOutput;
+    type Output;
 
     /// Waits for `self` to finish, then check if its output is:
     /// - a panic payload: resume that panic,
@@ -49,20 +43,5 @@ pub trait WaitForTermination {
     ///
     /// If `self` contains an expected termination, and we're shutting down anyway.
     #[track_caller]
-    fn panic_on_early_termination(self) -> Self::UnwrappedOutput;
-
-    /// Waits for `self` to finish, then check if its output is:
-    /// - a panic payload: resume that panic,
-    /// - an unexpected termination: return that error,
-    /// - an expected termination: return the provided `expected_termination_value`.
-    ///
-    /// Otherwise, returns the task return value of `self`.
-    ///
-    /// # Panics
-    ///
-    /// If `self` contains a panic payload.
-    #[track_caller]
-    fn error_on_early_termination<D>(self, expected_termination_value: D) -> Self::ResultOutput
-    where
-        D: Into<Self::ResultOutput> + Send + 'static;
+    fn wait_for_panics(self) -> Self::Output;
 }

zebra-chain/src/diagnostic/task/future.rs

@@ -9,7 +9,7 @@ use tokio::task::{JoinError, JoinHandle};
 
 use crate::shutdown::is_shutting_down;
 
-use super::{CheckForPanics, WaitForTermination};
+use super::{CheckForPanics, WaitForPanics};
 
 /// This is the return type of the [`JoinHandle`] future.
 impl<T> CheckForPanics for Result<T, JoinError> {
@@ -59,13 +59,11 @@ impl CheckForPanics for JoinError {
     }
 }
 
-impl<T> WaitForTermination for JoinHandle<T>
+impl<T> WaitForPanics for JoinHandle<T>
 where
     T: Send + 'static,
 {
-    type UnwrappedOutput = BoxFuture<'static, T>;
-
-    type ResultOutput = BoxFuture<'static, Result<T, JoinError>>;
+    type Output = BoxFuture<'static, T>;
 
     /// Returns a future which waits for `self` to finish, then checks if its output is:
     /// - a panic payload: resume that panic,
@@ -83,7 +81,7 @@ where
     /// If `self` contains an expected termination, and we're shutting down anyway.
     /// Futures hang by returning `Pending` and not setting a waker, so this uses minimal resources.
     #[track_caller]
-    fn panic_on_early_termination(self) -> Self::UnwrappedOutput {
+    fn wait_for_panics(self) -> Self::Output {
         async move {
             match self.await.check_for_panics() {
                 Ok(task_output) => task_output,
@@ -92,46 +90,4 @@ where
         }
         .boxed()
     }
-
-    /// Returns a future which waits for `self` to finish, then checks if its output is:
-    /// - a panic payload: resume that panic,
-    /// - an unexpected termination: return that error,
-    /// - an expected termination: return the provided `expected_termination_value`.
-    ///
-    /// Otherwise, returns the task return value of `self`.
-    ///
-    /// # Panics
-    ///
-    /// If `self` contains a panic payload.
-    #[track_caller]
-    fn error_on_early_termination<D>(self, expected_termination_value: D) -> Self::ResultOutput
-    where
-        D: Into<Self::ResultOutput> + Send + 'static,
-    {
-        async move {
-            let task_result = self.await;
-
-            let Err(join_error) = task_result else {
-                // Always `Ok`
-                return task_result;
-            };
-
-            match join_error.try_into_panic() {
-                Ok(panic_payload) => panic::resume_unwind(panic_payload),
-
-                // Substitute the default value
-                Err(task_cancelled) if is_shutting_down() => {
-                    debug!(
-                        ?task_cancelled,
-                        "ignoring cancelled task because Zebra is shutting down"
-                    );
-
-                    expected_termination_value.into().await
-                }
-
-                Err(task_cancelled) => Err(task_cancelled),
-            }
-        }
-        .boxed()
-    }
 }

zebra-chain/src/diagnostic/task/thread.rs

@@ -2,9 +2,13 @@
 //! - task handles
 //! - errors and panics
 
-use std::{panic, thread};
+use std::{
+    panic,
+    sync::Arc,
+    thread::{self, JoinHandle},
+};
 
-use super::{CheckForPanics, WaitForTermination};
+use super::{CheckForPanics, WaitForPanics};
 
 impl<T> CheckForPanics for thread::Result<T> {
     type Output = T;
@@ -24,25 +28,12 @@ impl<T> CheckForPanics for thread::Result<T> {
     }
 }
 
-impl<T> WaitForTermination for thread::JoinHandle<T> {
-    type UnwrappedOutput = T;
-
-    type ResultOutput = thread::Result<T>;
+impl<T> WaitForPanics for JoinHandle<T> {
+    type Output = T;
 
     /// Waits for the thread to finish, then panics if the thread panicked.
     #[track_caller]
-    fn panic_on_early_termination(self) -> Self::UnwrappedOutput {
+    fn wait_for_panics(self) -> Self::Output {
         self.join().check_for_panics()
     }
-
-    /// Waits for the thread to finish, then panics if the thread panicked.
-    ///
-    /// Threads can't be cancelled except by using a panic, so the returned result is always `Ok`.
-    #[track_caller]
-    fn error_on_early_termination<D>(self, _expected_termination_value: D) -> Self::ResultOutput
-    where
-        D: Into<Self::ResultOutput> + 'static,
-    {
-        Ok(self.panic_on_early_termination())
-    }
 }

zebra-network/src/peer_set/candidate_set.rs

@@ -8,7 +8,7 @@ use tokio::time::{sleep_until, timeout, Instant};
 use tower::{Service, ServiceExt};
 use tracing::Span;
 
-use zebra_chain::{diagnostic::task::WaitForTermination, serialization::DateTime32};
+use zebra_chain::{diagnostic::task::WaitForPanics, serialization::DateTime32};
 
 use crate::{
     constants, meta_addr::MetaAddrChange, peer_set::set::MorePeers, types::MetaAddr, AddressBook,
@@ -348,7 +348,7 @@ where
         tokio::task::spawn_blocking(move || {
             span.in_scope(|| address_book.lock().unwrap().extend(addrs))
         })
-        .panic_on_early_termination()
+        .wait_for_panics()
         .await
     }
 
@@ -403,7 +403,7 @@ where
         // Correctness: Spawn address book accesses on a blocking thread, to avoid deadlocks (see #1976).
         let span = Span::current();
         let next_peer = tokio::task::spawn_blocking(move || span.in_scope(next_peer))
-            .panic_on_early_termination()
+            .wait_for_panics()
             .await?;
 
         // Security: rate-limit new outbound peer connections

zebra-network/src/peer_set/initialize.rs

@@ -32,7 +32,7 @@ use tower::{
 use tracing::Span;
 use tracing_futures::Instrument;
 
-use zebra_chain::{chain_tip::ChainTip, diagnostic::task::WaitForTermination};
+use zebra_chain::{chain_tip::ChainTip, diagnostic::task::WaitForPanics};
 
 use crate::{
     address_book_updater::AddressBookUpdater,
@@ -206,7 +206,7 @@ where
 
     // Wait for the initial seed peer count
     let mut active_outbound_connections = initial_peers_join
-        .panic_on_early_termination()
+        .wait_for_panics()
         .await
         .expect("unexpected error connecting to initial peers");
     let active_initial_peer_count = active_outbound_connections.update_count();
@@ -343,7 +343,7 @@ where
                 }
                 .in_current_span(),
             )
-            .panic_on_early_termination()
+            .wait_for_panics()
         })
         .collect();
 
@@ -615,7 +615,7 @@ where
                 peerset_tx.clone(),
             )
             .await?
-            .panic_on_early_termination();
+            .wait_for_panics();
 
             handshakes.push(handshake_task);
 
@@ -903,7 +903,7 @@ where
                     }
                     .in_current_span(),
                 )
-                .panic_on_early_termination();
+                .wait_for_panics();
 
                 handshakes.push(handshake_or_crawl_handle);
             }
@@ -924,7 +924,7 @@ where
                     }
                     .in_current_span(),
                 )
-                .panic_on_early_termination();
+                .wait_for_panics();
 
                 handshakes.push(crawl_handle);
             }
@@ -1100,7 +1100,7 @@ async fn report_failed(address_book: Arc<std::sync::Mutex<AddressBook>>, addr: M
     let updated_addr = tokio::task::spawn_blocking(move || {
         span.in_scope(|| address_book.lock().unwrap().update(addr))
     })
-    .panic_on_early_termination()
+    .wait_for_panics()
     .await;
 
     assert_eq!(