Avoid panics on cancelled tasks in JoinSet.

crates/sui-network/src/discovery/mod.rs

@@ -105,7 +105,19 @@ impl DiscoveryEventLoop {
                     self.handle_trusted_peer_change_event(event);
                 }
                 Some(task_result) = self.tasks.join_next() => {
-                    task_result.unwrap();
+                    match task_result {
+                        Ok(()) => {},
+                        Err(e) => {
+                            if e.is_cancelled() {
+                                // avoid crashing on ungraceful shutdown
+                            } else if e.is_panic() {
+                                // propagate panics.
+                                std::panic::resume_unwind(e.into_panic());
+                            } else {
+                                panic!("task failed: {e}");
+                            }
+                        },
+                    };
                 },
                 // Once the shutdown notification resolves we can terminate the event loop
                 _ = &mut self.shutdown_handle => {

crates/sui-network/src/state_sync/mod.rs

@@ -323,7 +323,19 @@ where
                     self.handle_peer_event(peer_event);
                 },
                 Some(task_result) = self.tasks.join_next() => {
-                    task_result.unwrap();
+                    match task_result {
+                        Ok(()) => {},
+                        Err(e) => {
+                            if e.is_cancelled() {
+                                // avoid crashing on ungraceful shutdown
+                            } else if e.is_panic() {
+                                // propagate panics.
+                                std::panic::resume_unwind(e.into_panic());
+                            } else {
+                                panic!("task failed: {e}");
+                            }
+                        },
+                    };
 
                     if matches!(&self.sync_checkpoint_contents_task, Some(t) if t.is_finished()) {
                         self.sync_checkpoint_contents_task = None;

narwhal/primary/src/certificate_fetcher.rs

@@ -182,9 +182,19 @@ impl CertificateFetcher {
                     }
                 },
                 Some(result) = self.fetch_certificates_task.join_next(), if !self.fetch_certificates_task.is_empty() => {
-                    // propagate any panics. We don't expect for cancellations to get propagated as
-                    // we gracefully shutdown the component by exiting the loop first
-                    result.unwrap();
+                    match result {
+                        Ok(()) => {},
+                        Err(e) => {
+                            if e.is_cancelled() {
+                                // avoid crashing on ungraceful shutdown
+                            } else if e.is_panic() {
+                                // propagate panics.
+                                std::panic::resume_unwind(e.into_panic());
+                            } else {
+                                panic!("fetch certificates task failed: {e}");
+                            }
+                        },
+                    };
 
                     // Kick start another fetch task after the previous one terminates.
                     // If all targets have been fetched, the new task will clean up the targets and exit.

narwhal/primary/src/certifier.rs

@@ -407,7 +407,17 @@ impl Certifier {
                             self.synchronizer.accept_own_certificate(certificate).await
                         },
                         Ok(Err(e)) => Err(e),
-                        Err(_) => Err(DagError::ShuttingDown),
+                        Err(e) => {
+                            if e.is_cancelled() {
+                                // Ungraceful shutdown.
+                                Err(DagError::ShuttingDown)
+                            } else if e.is_panic() {
+                                // propagate panics.
+                                std::panic::resume_unwind(e.into_panic());
+                            } else {
+                                panic!("propose header task failed: {e}");
+                            }
+                        },
                     }
                 },