fix: return TaskManager in spawn_node (#6101)

This commit is contained in:
Dan Cline
2024-01-17 05:34:34 -05:00
committed by GitHub
parent afa860757d
commit 5637eed19e

View File

@ -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 =