chore: use task manager graceful shutdown (#5591)

bin/reth/src/node/mod.rs

@@ -728,9 +728,10 @@ impl<Ext: RethCliExt> NodeCommand<Ext> {
         task_executor.spawn_critical("p2p eth request handler", eth);
 
         let known_peers_file = self.network.persistent_peers_file(default_peers_path);
-        task_executor.spawn_critical_with_shutdown_signal("p2p network task", |shutdown| {
-            run_network_until_shutdown(shutdown, network, known_peers_file)
-        });
+        task_executor
+            .spawn_critical_with_graceful_shutdown_signal("p2p network task", |shutdown| {
+                run_network_until_shutdown(shutdown, network, known_peers_file)
+            });
 
         handle
     }
@@ -1029,7 +1030,7 @@ impl<Ext: RethCliExt> NodeCommand<Ext> {
 /// Drives the [NetworkManager] future until a [Shutdown](reth_tasks::shutdown::Shutdown) signal is
 /// received. If configured, this writes known peers to `persistent_peers_file` afterwards.
 async fn run_network_until_shutdown<C>(
-    shutdown: reth_tasks::shutdown::Shutdown,
+    shutdown: reth_tasks::shutdown::GracefulShutdown,
     network: NetworkManager<C>,
     persistent_peers_file: Option<PathBuf>,
 ) where
@@ -1037,9 +1038,12 @@ async fn run_network_until_shutdown<C>(
 {
     pin_mut!(network, shutdown);
 
+    let mut graceful_guard = None;
     tokio::select! {
         _ = &mut network => {},
-        _ = shutdown => {},
+        guard = shutdown => {
+            graceful_guard = Some(guard);
+        },
     }
 
     if let Some(file_path) = persistent_peers_file {
@@ -1057,6 +1061,8 @@ async fn run_network_until_shutdown<C>(
             }
         }
     }
+
+    drop(graceful_guard)
 }
 
 #[cfg(test)]

bin/reth/src/runner.rs

@@ -33,9 +33,11 @@ impl CliRunner {
             task_manager,
             run_until_ctrl_c(command(context)),
         ))?;
-        // after the command has finished or exit signal was received we drop the task manager which
-        // fires the shutdown signal to all tasks spawned via the task executor
-        drop(task_manager);
+
+        // after the command has finished or exit signal was received we shutdown the task manager
+        // which fires the shutdown signal to all tasks spawned via the task executor and
+        // awaiting on tasks spawned with graceful shutdown
+        task_manager.graceful_shutdown_with_timeout(std::time::Duration::from_secs(10));
 
         // drop the tokio runtime on a separate thread because drop blocks until its pools
         // (including blocking pool) are shutdown. In other words `drop(tokio_runtime)` would block