Consolidate IP and node/config management in tests

crates/mysten-network/src/multiaddr.rs

@@ -84,6 +84,22 @@ impl Multiaddr {
         }
     }
 
+    // Converts a /ip{4,6}/-/tcp/-[/-] Multiaddr to SocketAddr.
+    // Useful when an external library only accepts SocketAddr, e.g. to start a local server.
+    // See `client::endpoint_from_multiaddr()` for converting to Endpoint for clients.
+    pub fn to_socket_addr(&self) -> Result<SocketAddr> {
+        let mut iter = self.iter();
+        let ip = match iter.next().ok_or_else(|| {
+            eyre!("failed to convert to SocketAddr: Multiaddr does not contain IP")
+        })? {
+            Protocol::Ip4(ip4_addr) => IpAddr::V4(ip4_addr),
+            Protocol::Ip6(ip6_addr) => IpAddr::V6(ip6_addr),
+            unsupported => return Err(eyre!("unsupported protocol {unsupported}")),
+        };
+        let tcp_port = parse_tcp(&mut iter)?;
+        Ok(SocketAddr::new(ip, tcp_port))
+    }
+
     /// Set the ip address to `0.0.0.0`. For instance, it converts the following address
     /// `/ip4/155.138.174.208/tcp/1500/http` into `/ip4/0.0.0.0/tcp/1500/http`.
     pub fn zero_ip_multi_address(&self) -> Self {
@@ -151,23 +167,6 @@ impl<'de> serde::Deserialize<'de> for Multiaddr {
     }
 }
 
-// Converts a /ip{4,6}/-/tcp/-[/-] Multiaddr to SocketAddr.
-// Useful when an external library only accepts SocketAddr, e.g. to start a local server.
-// See `client::endpoint_from_multiaddr()` for converting to Endpoint for clients.
-pub fn to_socket_addr(addr: &Multiaddr) -> Result<SocketAddr> {
-    let mut iter = addr.iter();
-    let ip = match iter
-        .next()
-        .ok_or_else(|| eyre!("failed to convert to SocketAddr: Multiaddr does not contain IP"))?
-    {
-        Protocol::Ip4(ip4_addr) => IpAddr::V4(ip4_addr),
-        Protocol::Ip6(ip6_addr) => IpAddr::V6(ip6_addr),
-        unsupported => return Err(eyre!("unsupported protocol {unsupported}")),
-    };
-    let tcp_port = parse_tcp(&mut iter)?;
-    Ok(SocketAddr::new(ip, tcp_port))
-}
-
 pub(crate) fn parse_tcp<'a, T: Iterator<Item = Protocol<'a>>>(protocols: &mut T) -> Result<u16> {
     if let Protocol::Tcp(port) = protocols
         .next()
@@ -275,25 +274,29 @@ pub(crate) fn parse_unix(address: &Multiaddr) -> Result<(Cow<'_, str>, &'static
 
 #[cfg(test)]
 mod test {
-    use super::{to_socket_addr, Multiaddr};
+    use super::Multiaddr;
     use multiaddr::multiaddr;
 
     #[test]
     fn test_to_socket_addr_basic() {
         let multi_addr_ipv4 = Multiaddr(multiaddr!(Ip4([127, 0, 0, 1]), Tcp(10500u16)));
-        let socket_addr_ipv4 =
-            to_socket_addr(&multi_addr_ipv4).expect("Couldn't convert to socket addr");
+        let socket_addr_ipv4 = multi_addr_ipv4
+            .to_socket_addr()
+            .expect("Couldn't convert to socket addr");
         assert_eq!(socket_addr_ipv4.to_string(), "127.0.0.1:10500");
 
         let multi_addr_ipv6 = Multiaddr(multiaddr!(Ip6([172, 0, 0, 1, 1, 1, 1, 1]), Tcp(10500u16)));
-        let socket_addr_ipv6 =
-            to_socket_addr(&multi_addr_ipv6).expect("Couldn't convert to socket addr");
+        let socket_addr_ipv6 = multi_addr_ipv6
+            .to_socket_addr()
+            .expect("Couldn't convert to socket addr");
         assert_eq!(socket_addr_ipv6.to_string(), "[ac::1:1:1:1:1]:10500");
     }
 
     #[test]
     fn test_to_socket_addr_unsupported_protocol() {
         let multi_addr_dns = Multiaddr(multiaddr!(Dnsaddr("mysten.sui"), Tcp(10500u16)));
-        let _ = to_socket_addr(&multi_addr_dns).expect_err("DNS is unsupported");
+        let _ = multi_addr_dns
+            .to_socket_addr()
+            .expect_err("DNS is unsupported");
     }
 }

crates/sui-benchmark/Cargo.toml

@@ -38,6 +38,7 @@ sui-json-rpc-types = { path = "../sui-json-rpc-types" }
 sui-protocol-config = { path = "../sui-protocol-config" }
 sui-test-transaction-builder = { path = "../sui-test-transaction-builder" }
 sui-swarm-config = { path = "../sui-swarm-config" }
+sui-swarm = { path = "../sui-swarm" }
 telemetry-subscribers.workspace = true
 roaring = "0.10.1"
 regex = "1.7.1"

crates/sui-benchmark/src/benchmark_setup.rs

@@ -1,24 +1,28 @@
 // Copyright (c) Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
+
 use crate::bank::BenchmarkBank;
 use crate::options::Opts;
 use crate::util::get_ed25519_keypair_from_keystore;
 use crate::{FullNodeProxy, LocalValidatorAggregatorProxy, ValidatorProxy};
 use anyhow::{anyhow, bail, Context, Result};
 use prometheus::Registry;
+use rand::rngs::OsRng;
 use rand::seq::SliceRandom;
 use std::path::PathBuf;
 use std::sync::Arc;
 use std::thread::JoinHandle;
 use std::time::Duration;
-use sui_config::utils;
+use sui_config::local_ip_utils;
+use sui_config::node::ExpensiveSafetyCheckConfig;
+use sui_swarm_config::node_config_builder::FullnodeConfigBuilder;
 use sui_types::base_types::ObjectID;
 use sui_types::base_types::SuiAddress;
 use sui_types::crypto::{deterministic_random_account_key, AccountKeyPair};
 use sui_types::object::generate_max_test_gas_objects_with_owner;
 use sui_types::object::Owner;
+use test_utils::authority::spawn_test_authorities;
 use test_utils::authority::test_and_configure_authority_configs_with_objects;
-use test_utils::authority::{spawn_fullnode, spawn_test_authorities};
 use tokio::runtime::Builder;
 use tokio::sync::{oneshot, Barrier};
 use tokio::time::sleep;
@@ -112,8 +116,8 @@ impl Env {
             .clone();
         // Make the client runtime wait until we are done creating genesis objects
         let cloned_config = config.clone();
-        let fullnode_ip = format!("{}", utils::get_local_ip_for_tests());
-        let fullnode_rpc_port = utils::get_available_port(&fullnode_ip);
+        let fullnode_ip = local_ip_utils::localhost_for_testing();
+        let fullnode_rpc_port = local_ip_utils::get_available_port(&fullnode_ip);
         let fullnode_barrier = Arc::new(Barrier::new(2));
         let fullnode_barrier_clone = fullnode_barrier.clone();
         // spawn a thread to spin up sui nodes on the multi-threaded server runtime.
@@ -130,7 +134,17 @@ impl Env {
             server_runtime.block_on(async move {
                 // Setup the network
                 let _validators: Vec<_> = spawn_test_authorities(&cloned_config).await;
-                let _fullnode = spawn_fullnode(&cloned_config, Some(fullnode_rpc_port)).await;
+
+                let node_config = FullnodeConfigBuilder::new()
+                    .with_rpc_port(fullnode_rpc_port)
+                    .with_expensive_safety_check_config(
+                        ExpensiveSafetyCheckConfig::new_disable_all(),
+                    )
+                    .build(&mut OsRng, &cloned_config);
+                let node = sui_swarm::memory::Node::new(node_config);
+                node.start().await.unwrap();
+                let _fullnode = node.get_node_handle().unwrap();
+
                 fullnode_barrier_clone.wait().await;
                 barrier.wait().await;
                 recv.await.expect("Unable to wait for terminate signal");

crates/sui-benchmark/tests/simtest.rs

@@ -4,6 +4,7 @@
 #[cfg(msim)]
 mod test {
     use rand::{distributions::uniform::SampleRange, thread_rng, Rng};
+    use std::collections::HashSet;
     use std::path::PathBuf;
     use std::str::FromStr;
     use std::sync::atomic::{AtomicBool, Ordering};
@@ -102,16 +103,35 @@ mod test {
         test_simulated_load(TestInitData::new(&test_cluster).await, 120).await;
     }
 
+    /// Get a list of nodes that we don't want to kill in the crash recovery tests.
+    /// This includes the client node which is the node that is running the test, as well as
+    /// rpc fullnode which are needed to run the benchmark.
+    fn get_keep_alive_nodes(cluster: &TestCluster) -> HashSet<sui_simulator::task::NodeId> {
+        let mut keep_alive_nodes = HashSet::new();
+        // The first fullnode in the swarm ins the rpc fullnode.
+        keep_alive_nodes.insert(
+            cluster
+                .swarm
+                .fullnodes()
+                .next()
+                .unwrap()
+                .get_node_handle()
+                .unwrap()
+                .with(|n| n.get_sim_node_id()),
+        );
+        keep_alive_nodes.insert(sui_simulator::current_simnode_id());
+        keep_alive_nodes
+    }
+
     fn handle_failpoint(
         dead_validator: Arc<Mutex<Option<DeadValidator>>>,
-        client_node: sui_simulator::task::NodeId,
+        keep_alive_nodes: HashSet<sui_simulator::task::NodeId>,
         probability: f64,
     ) {
         let mut dead_validator = dead_validator.lock().unwrap();
         let cur_node = sui_simulator::current_simnode_id();
 
-        // never kill the client node (which is running the test)
-        if cur_node == client_node {
+        if keep_alive_nodes.contains(&cur_node) {
             return;
         }
 
@@ -167,7 +187,8 @@ mod test {
         let dead_validator_orig: Arc<Mutex<Option<DeadValidator>>> = Default::default();
 
         let dead_validator = dead_validator_orig.clone();
-        let client_node = sui_simulator::current_simnode_id();
+        let keep_alive_nodes = get_keep_alive_nodes(&test_cluster);
+        let keep_alive_nodes_clone = keep_alive_nodes.clone();
         register_fail_points(
             &[
                 "batch-write-before",
@@ -180,28 +201,35 @@ mod test {
                 "highest-executed-checkpoint",
             ],
             move || {
-                handle_failpoint(dead_validator.clone(), client_node, 0.02);
+                handle_failpoint(dead_validator.clone(), keep_alive_nodes_clone.clone(), 0.02);
             },
         );
 
         let dead_validator = dead_validator_orig.clone();
+        let keep_alive_nodes_clone = keep_alive_nodes.clone();
         register_fail_point_async("crash", move || {
             let dead_validator = dead_validator.clone();
+            let keep_alive_nodes_clone = keep_alive_nodes_clone.clone();
             async move {
-                handle_failpoint(dead_validator.clone(), client_node, 0.01);
+                handle_failpoint(dead_validator.clone(), keep_alive_nodes_clone.clone(), 0.01);
             }
         });
 
         // Narwhal fail points.
         let dead_validator = dead_validator_orig.clone();
+        let keep_alive_nodes_clone = keep_alive_nodes.clone();
         register_fail_points(
             &[
                 "narwhal-rpc-response",
                 "narwhal-store-before-write",
                 "narwhal-store-after-write",
             ],
             move || {
-                handle_failpoint(dead_validator.clone(), client_node, 0.001);
+                handle_failpoint(
+                    dead_validator.clone(),
+                    keep_alive_nodes_clone.clone(),
+                    0.001,
+                );
             },
         );
         register_fail_point_async("narwhal-delay", || delay_failpoint(10..20, 0.001));
@@ -215,9 +243,9 @@ mod test {
         let test_cluster = build_test_cluster(4, 10000).await;
 
         let dead_validator: Arc<Mutex<Option<DeadValidator>>> = Default::default();
-        let client_node = sui_simulator::current_simnode_id();
+        let keep_alive_nodes = get_keep_alive_nodes(&test_cluster);
         register_fail_points(&["before-open-new-epoch-store"], move || {
-            handle_failpoint(dead_validator.clone(), client_node, 1.0);
+            handle_failpoint(dead_validator.clone(), keep_alive_nodes.clone(), 1.0);
         });
         test_simulated_load(TestInitData::new(&test_cluster).await, 120).await;
     }

crates/sui-cluster-test/src/cluster.rs

@@ -186,7 +186,7 @@ impl Cluster for LocalNewCluster {
             cluster_builder = cluster_builder.with_epoch_duration_ms(epoch_duration_ms);
         }
         if let Some(rpc_port) = fullnode_port {
-            cluster_builder = cluster_builder.set_fullnode_rpc_port(rpc_port);
+            cluster_builder = cluster_builder.with_fullnode_rpc_port(rpc_port);
         }
 
         let mut test_cluster = cluster_builder.build().await?;

crates/sui-config/src/lib.rs

@@ -11,11 +11,11 @@ use tracing::trace;
 
 pub mod certificate_deny_config;
 pub mod genesis;
+pub mod local_ip_utils;
 pub mod node;
 pub mod node_config_metrics;
 pub mod p2p;
 pub mod transaction_deny_config;
-pub mod utils;
 
 pub use node::{ConsensusConfig, NodeConfig};