Add tpu-client-next client to STS

anza-xyz · Feb 5, 2025 · f20cc08 · f20cc08
1 parent bd6e9f9
commit f20cc08
Show file tree

Hide file tree

Showing 8 changed files with 399 additions and 13 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/programs/sbf/Cargo.lock b/programs/sbf/Cargo.lock
diff --git a/send-transaction-service/Cargo.toml b/send-transaction-service/Cargo.toml
@@ -10,17 +10,20 @@ license = { workspace = true }
 edition = { workspace = true }
 
 [dependencies]
+async-trait = { workspace = true }
 crossbeam-channel = { workspace = true }
 itertools = { workspace = true }
 log = { workspace = true }
 solana-client = { workspace = true }
 solana-connection-cache = { workspace = true }
+solana-keypair = { workspace = true }
 solana-measure = { workspace = true }
 solana-metrics = { workspace = true }
 solana-runtime = { workspace = true }
 solana-sdk = { workspace = true }
-solana-tpu-client = { workspace = true }
+solana-tpu-client-next = { workspace = true }
 tokio = { workspace = true, features = ["full"] }
+tokio-util = { workspace = true }
 
 [dev-dependencies]
 solana-logger = { workspace = true }

diff --git a/send-transaction-service/src/send_transaction_service.rs b/send-transaction-service/src/send_transaction_service.rs
@@ -507,7 +507,10 @@ impl SendTransactionService {
 mod test {
     use {
         super::*,
-        crate::{test_utils::ClientWithCreator, tpu_info::NullTpuInfo},
+        crate::{
+            test_utils::ClientWithCreator, tpu_info::NullTpuInfo,
+            transaction_client::TpuClientNextClient,
+        },
         crossbeam_channel::{bounded, unbounded},
         solana_sdk::{
             account::AccountSharedData,
@@ -549,6 +552,12 @@ mod test {
         service_exit::<ConnectionCacheClient<NullTpuInfo>>(None);
     }
 
+    #[tokio::test(flavor = "multi_thread")]
+
+    async fn service_exit_with_tpu_client_next() {
+        service_exit::<TpuClientNextClient<NullTpuInfo>>(Some(Handle::current()));
+    }
+
     fn validator_exit<C: ClientWithCreator>(maybe_runtime: Option<Handle>) {
         let bank = Bank::default_for_tests();
         let bank_forks = BankForks::new_rw_arc(bank);
@@ -595,6 +604,11 @@ mod test {
         validator_exit::<ConnectionCacheClient<NullTpuInfo>>(None);
     }
 
+    #[tokio::test(flavor = "multi_thread")]
+    async fn validator_exit_with_tpu_client_next() {
+        validator_exit::<TpuClientNextClient<NullTpuInfo>>(Some(Handle::current()));
+    }
+
     fn process_transactions<C: ClientWithCreator>(maybe_runtime: Option<Handle>) {
         solana_logger::setup();
 
@@ -865,6 +879,11 @@ mod test {
         process_transactions::<ConnectionCacheClient<NullTpuInfo>>(None);
     }
 
+    #[tokio::test(flavor = "multi_thread")]
+    async fn process_transactions_with_tpu_client_next() {
+        process_transactions::<TpuClientNextClient<NullTpuInfo>>(Some(Handle::current()));
+    }
+
     fn retry_durable_nonce_transactions<C: ClientWithCreator>(maybe_runtime: Option<Handle>) {
         solana_logger::setup();
 
@@ -1169,4 +1188,11 @@ mod test {
     fn retry_durable_nonce_transactions_with_connection_cache() {
         retry_durable_nonce_transactions::<ConnectionCacheClient<NullTpuInfo>>(None);
     }
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn retry_durable_nonce_transactions_with_tpu_client_next() {
+        retry_durable_nonce_transactions::<TpuClientNextClient<NullTpuInfo>>(Some(
+            Handle::current(),
+        ));
+    }
 }
diff --git a/send-transaction-service/src/test_utils.rs b/send-transaction-service/src/test_utils.rs
@@ -4,7 +4,9 @@
 use {
     crate::{
         tpu_info::NullTpuInfo,
-        transaction_client::{ConnectionCacheClient, TpuInfoWithSendStatic, TransactionClient},
+        transaction_client::{
+            ConnectionCacheClient, TpuClientNextClient, TpuInfoWithSendStatic, TransactionClient,
+        },
     },
     solana_client::connection_cache::ConnectionCache,
     std::{net::SocketAddr, sync::Arc},
@@ -42,6 +44,26 @@ impl CreateClient for ConnectionCacheClient<NullTpuInfo> {
     }
 }
 
+impl CreateClient for TpuClientNextClient<NullTpuInfo> {
+    fn create_client(
+        maybe_runtime: Option<Handle>,
+        my_tpu_address: SocketAddr,
+        tpu_peers: Option<Vec<SocketAddr>>,
+        leader_forward_count: u64,
+    ) -> Self {
+        let runtime_handle =
+            maybe_runtime.expect("Runtime should be provided for the TpuClientNextClient.");
+        Self::new(
+            runtime_handle,
+            my_tpu_address,
+            tpu_peers,
+            None,
+            leader_forward_count,
+            None,
+        )
+    }
+}
+
 pub trait Cancelable {
     fn cancel(&self);
 }
@@ -53,6 +75,15 @@ where
     fn cancel(&self) {}
 }
 
+impl<T> Cancelable for TpuClientNextClient<T>
+where
+    T: TpuInfoWithSendStatic + Clone,
+{
+    fn cancel(&self) {
+        self.cancel().unwrap();
+    }
+}
+
 // Define type alias to simplify definition of test functions.
 pub trait ClientWithCreator:
     CreateClient + TransactionClient + Cancelable + Send + Clone + 'static

diff --git a/send-transaction-service/src/tpu_info.rs b/send-transaction-service/src/tpu_info.rs
@@ -23,10 +23,6 @@ pub trait TpuInfo {
     ///
     /// For example, if leader schedule was `[L1, L1, L1, L1, L2, L2, L2, L2,
     /// L1, ...]` it will return `[L1, L2, L1]`.
-    #[allow(
-        dead_code,
-        reason = "This function will be used when tpu-client-next will be added to this module."
-    )]
     fn get_not_unique_leader_tpus(&self, max_count: u64, protocol: Protocol) -> Vec<&SocketAddr>;
 
     /// In addition to the tpu address, also return the leader slot