add in method for building a TpuClient for LocalCluster tests

anza-xyz · Mar 14, 2024 · 664a003 · 664a003
1 parent b3fd87f
commit 664a003
Show file tree

Hide file tree

Showing 5 changed files with 61 additions and 51 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/bench-tps/tests/bench_tps.rs b/bench-tps/tests/bench_tps.rs
@@ -7,14 +7,11 @@ use {
         cli::{Config, InstructionPaddingConfig},
         send_batch::generate_durable_nonce_accounts,
     },
-    solana_client::{
-        connection_cache::ConnectionCache,
-        tpu_client::{TpuClient, TpuClientConfig},
-    },
+    solana_client::tpu_client::{TpuClient, TpuClientConfig},
     solana_core::validator::ValidatorConfig,
     solana_faucet::faucet::run_local_faucet,
     solana_local_cluster::{
-        local_cluster::{ClusterConfig, LocalCluster},
+        local_cluster::{build_tpu_quic_client, ClusterConfig, LocalCluster},
         validator_configs::make_identical_validator_configs,
     },
     solana_rpc::rpc::JsonRpcConfig,
@@ -78,24 +75,9 @@ fn test_bench_tps_local_cluster(config: Config) {
 
     cluster.transfer(&cluster.funding_keypair, &faucet_pubkey, 100_000_000);
 
-    let ConnectionCache::Quic(cache) = &*cluster.connection_cache else {
-        panic!("Expected a Quic ConnectionCache.");
-    };
-
-    let rpc_pubsub_url = format!("ws://{}/", cluster.entry_point_info.rpc_pubsub().unwrap());
-    let rpc_url = format!("http://{}", cluster.entry_point_info.rpc().unwrap());
-
-    let client = Arc::new(
-        TpuClient::new_with_connection_cache(
-            Arc::new(RpcClient::new(rpc_url)),
-            rpc_pubsub_url.as_str(),
-            TpuClientConfig::default(),
-            cache.clone(),
-        )
-        .unwrap_or_else(|err| {
-            panic!("Could not create TpuClient {err:?}");
-        }),
-    );
+    let client = Arc::new(build_tpu_quic_client(&cluster).unwrap_or_else(|err| {
+        panic!("Could not create TpuClient with Quic Cache {err:?}");
+    }));
 
     let lamports_per_account = 100;
 

diff --git a/dos/src/main.rs b/dos/src/main.rs
@@ -824,13 +824,12 @@ pub mod test {
         solana_gossip::contact_info::LegacyContactInfo,
         solana_local_cluster::{
             cluster::Cluster,
-            local_cluster::{ClusterConfig, LocalCluster},
+            local_cluster::{build_tpu_quic_client, ClusterConfig, LocalCluster},
             validator_configs::make_identical_validator_configs,
         },
         solana_quic_client::{QuicConfig, QuicConnectionManager, QuicPool},
         solana_rpc::rpc::JsonRpcConfig,
         solana_sdk::timing::timestamp,
-        solana_tpu_client::tpu_client::TpuClientConfig,
     };
 
     const TEST_SEND_BATCH_SIZE: usize = 1;
@@ -843,29 +842,6 @@ pub mod test {
         );
     }
 
-    fn build_tpu_quic_client(
-        cluster: &LocalCluster,
-    ) -> Arc<TpuClient<QuicPool, QuicConnectionManager, QuicConfig>> {
-        let rpc_pubsub_url = format!("ws://{}/", cluster.entry_point_info.rpc_pubsub().unwrap());
-        let rpc_url = format!("http://{}", cluster.entry_point_info.rpc().unwrap());
-
-        let ConnectionCache::Quic(cache) = &*cluster.connection_cache else {
-            panic!("Expected a Quic ConnectionCache.");
-        };
-
-        Arc::new(
-            TpuClient::new_with_connection_cache(
-                Arc::new(RpcClient::new(rpc_url)),
-                rpc_pubsub_url.as_str(),
-                TpuClientConfig::default(),
-                cache.clone(),
-            )
-            .unwrap_or_else(|err| {
-                panic!("Could not create TpuClient with Quic Cache {err:?}");
-            }),
-        )
-    }
-
     #[test]
     fn test_dos() {
         let nodes = [ContactInfo::new_localhost(
@@ -1003,7 +979,9 @@ pub mod test {
             .unwrap();
         let nodes_slice = [node];
 
-        let client = build_tpu_quic_client(&cluster);
+        let client = Arc::new(build_tpu_quic_client(&cluster).unwrap_or_else(|err| {
+            panic!("Could not create TpuClient with Quic Cache {err:?}");
+        }));
 
         // creates one transaction with 8 valid signatures and sends it 10 times
         run_dos(
@@ -1135,7 +1113,9 @@ pub mod test {
             .unwrap();
         let nodes_slice = [node];
 
-        let client = build_tpu_quic_client(&cluster);
+        let client = Arc::new(build_tpu_quic_client(&cluster).unwrap_or_else(|err| {
+            panic!("Could not create TpuClient with Quic Cache {err:?}");
+        }));
 
         // creates one transaction and sends it 10 times
         // this is done in single thread

diff --git a/local-cluster/Cargo.toml b/local-cluster/Cargo.toml
@@ -24,6 +24,7 @@ solana-gossip = { workspace = true }
 solana-ledger = { workspace = true }
 solana-logger = { workspace = true }
 solana-pubsub-client = { workspace = true }
+solana-quic-client = { workspace = true }
 solana-rpc-client = { workspace = true }
 solana-rpc-client-api = { workspace = true }
 solana-runtime = { workspace = true }

diff --git a/local-cluster/src/local_cluster.rs b/local-cluster/src/local_cluster.rs
@@ -7,7 +7,12 @@ use {
     itertools::izip,
     log::*,
     solana_accounts_db::utils::create_accounts_run_and_snapshot_dirs,
-    solana_client::{connection_cache::ConnectionCache, thin_client::ThinClient},
+    solana_client::{
+        connection_cache::ConnectionCache,
+        rpc_client::RpcClient,
+        thin_client::ThinClient,
+        tpu_client::{TpuClient, TpuClientConfig},
+    },
     solana_core::{
         consensus::tower_storage::FileTowerStorage,
         validator::{Validator, ValidatorConfig, ValidatorStartProgress},
@@ -18,6 +23,7 @@ use {
         gossip_service::discover_cluster,
     },
     solana_ledger::{create_new_tmp_ledger, shred::Shred},
+    solana_quic_client::{QuicConfig, QuicConnectionManager, QuicPool},
     solana_runtime::{
         genesis_utils::{
             create_genesis_config_with_vote_accounts_and_cluster_type, GenesisConfigInfo,
@@ -63,6 +69,46 @@ use {
     },
 };
 
+pub fn build_tpu_quic_client(
+    cluster: &LocalCluster,
+) -> Result<TpuClient<QuicPool, QuicConnectionManager, QuicConfig>> {
+    build_tpu_client(cluster, |rpc_url| Arc::new(RpcClient::new(rpc_url)))
+}
+
+pub fn build_tpu_quic_client_with_commitment(
+    cluster: &LocalCluster,
+    commitment_config: CommitmentConfig,
+) -> Result<TpuClient<QuicPool, QuicConnectionManager, QuicConfig>> {
+    build_tpu_client(cluster, |rpc_url| {
+        Arc::new(RpcClient::new_with_commitment(rpc_url, commitment_config))
+    })
+}
+
+fn build_tpu_client<F>(
+    cluster: &LocalCluster,
+    rpc_client_builder: F,
+) -> Result<TpuClient<QuicPool, QuicConnectionManager, QuicConfig>>
+where
+    F: FnOnce(String) -> Arc<RpcClient>,
+{
+    let rpc_pubsub_url = format!("ws://{}/", cluster.entry_point_info.rpc_pubsub().unwrap());
+    let rpc_url = format!("http://{}", cluster.entry_point_info.rpc().unwrap());
+
+    let ConnectionCache::Quic(cache) = &*cluster.connection_cache else {
+        panic!("Expected a Quic ConnectionCache.");
+    };
+
+    let tpu_client = TpuClient::new_with_connection_cache(
+        rpc_client_builder(rpc_url),
+        rpc_pubsub_url.as_str(),
+        TpuClientConfig::default(),
+        cache.clone(),
+    )
+    .map_err(|err| Error::new(ErrorKind::Other, format!("TpuSenderError: {}", err)))?;
+
+    Ok(tpu_client)
+}
+
 const DUMMY_SNAPSHOT_CONFIG_PATH_MARKER: &str = "dummy";
 
 pub struct ClusterConfig {