Add get_blocks and get_slot methods to bench-tps-client (anza-xyz#94)

* add get_block(s)/slot methods to BenchTpsClient * Update Cargo.lock * add commitment level for get_slot/blocks
codebender828 · Mar 9, 2024 · ca1af7b · ca1af7b
1 parent 264f4df
commit ca1af7b
Show file tree

Hide file tree

Showing 7 changed files with 132 additions and 5 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/bench-tps/Cargo.toml b/bench-tps/Cargo.toml
@@ -37,6 +37,7 @@ solana-sdk = { workspace = true }
 solana-streamer = { workspace = true }
 solana-thin-client = { workspace = true }
 solana-tpu-client = { workspace = true }
+solana-transaction-status = { workspace = true }
 solana-version = { workspace = true }
 spl-instruction-padding = { workspace = true }
 thiserror = { workspace = true }

diff --git a/bench-tps/src/bench_tps_client.rs b/bench-tps/src/bench_tps_client.rs
@@ -1,11 +1,12 @@
 use {
-    solana_rpc_client_api::client_error::Error as ClientError,
+    solana_rpc_client_api::{client_error::Error as ClientError, config::RpcBlockConfig},
     solana_sdk::{
         account::Account, commitment_config::CommitmentConfig, epoch_info::EpochInfo, hash::Hash,
-        message::Message, pubkey::Pubkey, signature::Signature, transaction::Transaction,
-        transport::TransportError,
+        message::Message, pubkey::Pubkey, signature::Signature, slot_history::Slot,
+        transaction::Transaction, transport::TransportError,
     },
     solana_tpu_client::tpu_client::TpuSenderError,
+    solana_transaction_status::UiConfirmedBlock,
     thiserror::Error,
 };
 
@@ -93,6 +94,21 @@ pub trait BenchTpsClient {
     ) -> Result<Account>;
 
     fn get_multiple_accounts(&self, pubkeys: &[Pubkey]) -> Result<Vec<Option<Account>>>;
+
+    fn get_slot_with_commitment(&self, commitment_config: CommitmentConfig) -> Result<Slot>;
+
+    fn get_blocks_with_commitment(
+        &self,
+        start_slot: Slot,
+        end_slot: Option<Slot>,
+        commitment_config: CommitmentConfig,
+    ) -> Result<Vec<Slot>>;
+
+    fn get_block_with_config(
+        &self,
+        slot: Slot,
+        rpc_block_config: RpcBlockConfig,
+    ) -> Result<UiConfirmedBlock>;
 }
 
 mod bank_client;

diff --git a/bench-tps/src/bench_tps_client/bank_client.rs b/bench-tps/src/bench_tps_client/bank_client.rs
@@ -1,5 +1,6 @@
 use {
     crate::bench_tps_client::{BenchTpsClient, BenchTpsError, Result},
+    solana_rpc_client_api::config::RpcBlockConfig,
     solana_runtime::bank_client::BankClient,
     solana_sdk::{
         account::Account,
@@ -10,8 +11,10 @@ use {
         message::Message,
         pubkey::Pubkey,
         signature::Signature,
+        slot_history::Slot,
         transaction::Transaction,
     },
+    solana_transaction_status::UiConfirmedBlock,
 };
 
 impl BenchTpsClient for BankClient {
@@ -111,4 +114,25 @@ impl BenchTpsClient for BankClient {
     fn get_multiple_accounts(&self, _pubkeys: &[Pubkey]) -> Result<Vec<Option<Account>>> {
         unimplemented!("BankClient doesn't support get_multiple_accounts");
     }
+
+    fn get_slot_with_commitment(&self, commitment_config: CommitmentConfig) -> Result<Slot> {
+        SyncClient::get_slot_with_commitment(self, commitment_config).map_err(|err| err.into())
+    }
+
+    fn get_blocks_with_commitment(
+        &self,
+        _start_slot: Slot,
+        _end_slot: Option<Slot>,
+        _commitment_config: CommitmentConfig,
+    ) -> Result<Vec<Slot>> {
+        unimplemented!("BankClient doesn't support get_blocks");
+    }
+
+    fn get_block_with_config(
+        &self,
+        _slot: Slot,
+        _rpc_block_config: RpcBlockConfig,
+    ) -> Result<UiConfirmedBlock> {
+        unimplemented!("BankClient doesn't support get_block_with_config");
+    }
 }
diff --git a/bench-tps/src/bench_tps_client/rpc_client.rs b/bench-tps/src/bench_tps_client/rpc_client.rs
@@ -1,10 +1,13 @@
 use {
     crate::bench_tps_client::{BenchTpsClient, BenchTpsError, Result},
     solana_rpc_client::rpc_client::RpcClient,
+    solana_rpc_client_api::config::RpcBlockConfig,
     solana_sdk::{
         account::Account, commitment_config::CommitmentConfig, epoch_info::EpochInfo, hash::Hash,
-        message::Message, pubkey::Pubkey, signature::Signature, transaction::Transaction,
+        message::Message, pubkey::Pubkey, signature::Signature, slot_history::Slot,
+        transaction::Transaction,
     },
+    solana_transaction_status::UiConfirmedBlock,
 };
 
 impl BenchTpsClient for RpcClient {
@@ -104,4 +107,26 @@ impl BenchTpsClient for RpcClient {
     fn get_multiple_accounts(&self, pubkeys: &[Pubkey]) -> Result<Vec<Option<Account>>> {
         RpcClient::get_multiple_accounts(self, pubkeys).map_err(|err| err.into())
     }
+
+    fn get_slot_with_commitment(&self, commitment_config: CommitmentConfig) -> Result<Slot> {
+        RpcClient::get_slot_with_commitment(self, commitment_config).map_err(|err| err.into())
+    }
+
+    fn get_blocks_with_commitment(
+        &self,
+        start_slot: Slot,
+        end_slot: Option<Slot>,
+        commitment_config: CommitmentConfig,
+    ) -> Result<Vec<Slot>> {
+        RpcClient::get_blocks_with_commitment(self, start_slot, end_slot, commitment_config)
+            .map_err(|err| err.into())
+    }
+
+    fn get_block_with_config(
+        &self,
+        slot: Slot,
+        rpc_block_config: RpcBlockConfig,
+    ) -> Result<UiConfirmedBlock> {
+        RpcClient::get_block_with_config(self, slot, rpc_block_config).map_err(|err| err.into())
+    }
 }
diff --git a/bench-tps/src/bench_tps_client/thin_client.rs b/bench-tps/src/bench_tps_client/thin_client.rs
@@ -1,6 +1,7 @@
 use {
     crate::bench_tps_client::{BenchTpsClient, BenchTpsError, Result},
     solana_client::thin_client::ThinClient,
+    solana_rpc_client_api::config::RpcBlockConfig,
     solana_sdk::{
         account::Account,
         client::{AsyncClient, Client, SyncClient},
@@ -10,8 +11,10 @@ use {
         message::Message,
         pubkey::Pubkey,
         signature::Signature,
+        slot_history::Slot,
         transaction::Transaction,
     },
+    solana_transaction_status::UiConfirmedBlock,
 };
 
 impl BenchTpsClient for ThinClient {
@@ -110,4 +113,31 @@ impl BenchTpsClient for ThinClient {
             .get_multiple_accounts(pubkeys)
             .map_err(|err| err.into())
     }
+
+    fn get_slot_with_commitment(&self, commitment_config: CommitmentConfig) -> Result<Slot> {
+        self.rpc_client()
+            .get_slot_with_commitment(commitment_config)
+            .map_err(|err| err.into())
+    }
+
+    fn get_blocks_with_commitment(
+        &self,
+        start_slot: Slot,
+        end_slot: Option<Slot>,
+        commitment_config: CommitmentConfig,
+    ) -> Result<Vec<Slot>> {
+        self.rpc_client()
+            .get_blocks_with_commitment(start_slot, end_slot, commitment_config)
+            .map_err(|err| err.into())
+    }
+
+    fn get_block_with_config(
+        &self,
+        slot: Slot,
+        rpc_block_config: RpcBlockConfig,
+    ) -> Result<UiConfirmedBlock> {
+        self.rpc_client()
+            .get_block_with_config(slot, rpc_block_config)
+            .map_err(|err| err.into())
+    }
 }
diff --git a/bench-tps/src/bench_tps_client/tpu_client.rs b/bench-tps/src/bench_tps_client/tpu_client.rs
@@ -4,10 +4,13 @@ use {
     solana_connection_cache::connection_cache::{
         ConnectionManager, ConnectionPool, NewConnectionConfig,
     },
+    solana_rpc_client_api::config::RpcBlockConfig,
     solana_sdk::{
         account::Account, commitment_config::CommitmentConfig, epoch_info::EpochInfo, hash::Hash,
-        message::Message, pubkey::Pubkey, signature::Signature, transaction::Transaction,
+        message::Message, pubkey::Pubkey, signature::Signature, slot_history::Slot,
+        transaction::Transaction,
     },
+    solana_transaction_status::UiConfirmedBlock,
 };
 
 impl<P, M, C> BenchTpsClient for TpuClient<P, M, C>
@@ -130,4 +133,31 @@ where
             .get_multiple_accounts(pubkeys)
             .map_err(|err| err.into())
     }
+
+    fn get_slot_with_commitment(&self, commitment_config: CommitmentConfig) -> Result<Slot> {
+        self.rpc_client()
+            .get_slot_with_commitment(commitment_config)
+            .map_err(|err| err.into())
+    }
+
+    fn get_blocks_with_commitment(
+        &self,
+        start_slot: Slot,
+        end_slot: Option<Slot>,
+        commitment_config: CommitmentConfig,
+    ) -> Result<Vec<Slot>> {
+        self.rpc_client()
+            .get_blocks_with_commitment(start_slot, end_slot, commitment_config)
+            .map_err(|err| err.into())
+    }
+
+    fn get_block_with_config(
+        &self,
+        slot: Slot,
+        rpc_block_config: RpcBlockConfig,
+    ) -> Result<UiConfirmedBlock> {
+        self.rpc_client()
+            .get_block_with_config(slot, rpc_block_config)
+            .map_err(|err| err.into())
+    }
 }