From 5637eed19e1f69e38e6df41eb1032589cbbb4fdd Mon Sep 17 00:00:00 2001
From: Dan Cline <6798349+Rjected@users.noreply.github.com>
Date: Wed, 17 Jan 2024 05:34:34 -0500
Subject: [PATCH] fix: return TaskManager in spawn_node (#6101)

---
 bin/reth/src/builder/mod.rs | 25 +++++++++++--------------
 1 file changed, 11 insertions(+), 14 deletions(-)

diff --git a/bin/reth/src/builder/mod.rs b/bin/reth/src/builder/mod.rs
index 2a7efea41..bbab5e91a 100644
--- a/bin/reth/src/builder/mod.rs
+++ b/bin/reth/src/builder/mod.rs
@@ -100,12 +100,9 @@ use std::{
     path::PathBuf,
     sync::Arc,
 };
-use tokio::{
-    runtime::Handle,
-    sync::{
-        mpsc::{unbounded_channel, Receiver, UnboundedSender},
-        oneshot, watch,
-    },
+use tokio::sync::{
+    mpsc::{unbounded_channel, Receiver, UnboundedSender},
+    oneshot, watch,
 };
 use tracing::*;
 
@@ -1386,7 +1383,7 @@ impl NodeHandle {
 }
 
 /// A simple function to launch a node with the specified [NodeConfig], spawning tasks on the
-/// [TaskExecutor] constructed from [Handle::current].
+/// [TaskExecutor] constructed from [TaskManager::current].
 ///
 /// # Example
 /// ```
@@ -1402,14 +1399,13 @@ impl NodeHandle {
 ///     let builder = NodeConfig::test().with_rpc(rpc_args).with_instance(2);
 ///
 ///     // Spawn the builder, returning a handle to the node
-///     let _handle = spawn_node(builder).await.unwrap();
+///     let (_handle, _manager) = spawn_node(builder).await.unwrap();
 /// }
 /// ```
-pub async fn spawn_node(config: NodeConfig) -> eyre::Result<NodeHandle> {
-    let handle = Handle::current();
-    let task_manager = TaskManager::new(handle);
+pub async fn spawn_node(config: NodeConfig) -> eyre::Result<(NodeHandle, TaskManager)> {
+    let task_manager = TaskManager::current();
     let ext = DefaultRethNodeCommandConfig;
-    config.launch::<()>(ext, task_manager.executor()).await
+    Ok((config.launch::<()>(ext, task_manager.executor()).await?, task_manager))
 }
 
 #[cfg(test)]
@@ -1426,7 +1422,7 @@ mod tests {
         // NOTE: tests here manually set an instance number. The alternative would be to use an
         // atomic counter. This works for `cargo test` but if tests would be run in `nextest` then
         // they would become flaky. So new tests should manually set a unique instance number.
-        let handle =
+        let (handle, _manager) =
             spawn_node(NodeConfig::test().with_rpc(rpc_args).with_instance(1)).await.unwrap();
 
         // call a function on the node
@@ -1440,7 +1436,7 @@ mod tests {
     #[tokio::test]
     async fn rpc_handles_none_without_http() {
         // this launches a test node _without_ http
-        let handle = spawn_node(NodeConfig::test().with_instance(2)).await.unwrap();
+        let (handle, _manager) = spawn_node(NodeConfig::test().with_instance(2)).await.unwrap();
 
         // ensure that the `http_client` is none
         let maybe_client = handle.rpc_server_handles().rpc.http_client();
@@ -1454,6 +1450,7 @@ mod tests {
 
         // this reserves instances 3-6
         let starting_instance = 3;
+        // contains handles and managers
         let mut handles = Vec::new();
         for i in 0..num_nodes {
             let handle =