add get_account to BenchTpsClient (#26068)

* add get_account to BenchTpsClient * improve error reporting in get_account
solana-labs · Jun 21, 2022 · 4031a37 · 4031a37
1 parent 51f26dc
commit 4031a37
Show file tree

Hide file tree

Showing 5 changed files with 38 additions and 6 deletions.
diff --git a/bench-tps/src/bench_tps_client.rs b/bench-tps/src/bench_tps_client.rs
@@ -1,8 +1,9 @@
 use {
     solana_client::{client_error::ClientError, tpu_client::TpuSenderError},
     solana_sdk::{
-        commitment_config::CommitmentConfig, epoch_info::EpochInfo, hash::Hash, message::Message,
-        pubkey::Pubkey, signature::Signature, transaction::Transaction, transport::TransportError,
+        account::Account, commitment_config::CommitmentConfig, epoch_info::EpochInfo, hash::Hash,
+        message::Message, pubkey::Pubkey, signature::Signature, transaction::Transaction,
+        transport::TransportError,
     },
     thiserror::Error,
 };
@@ -79,6 +80,9 @@ pub trait BenchTpsClient {
         lamports: u64,
         recent_blockhash: &Hash,
     ) -> Result<Signature>;
+
+    /// Returns all information associated with the account of the provided pubkey
+    fn get_account(&self, pubkey: &Pubkey) -> Result<Account>;
 }
 
 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
@@ -2,6 +2,7 @@ use {
     crate::bench_tps_client::{BenchTpsClient, BenchTpsError, Result},
     solana_runtime::bank_client::BankClient,
     solana_sdk::{
+        account::Account,
         client::{AsyncClient, SyncClient},
         commitment_config::CommitmentConfig,
         epoch_info::EpochInfo,
@@ -82,4 +83,14 @@ impl BenchTpsClient for BankClient {
         // BankClient doesn't support airdrops
         Err(BenchTpsError::AirdropFailure)
     }
+
+    fn get_account(&self, pubkey: &Pubkey) -> Result<Account> {
+        SyncClient::get_account(self, pubkey)
+            .map_err(|err| err.into())
+            .and_then(|account| {
+                account.ok_or_else(|| {
+                    BenchTpsError::Custom(format!("AccountNotFound: pubkey={}", pubkey))
+                })
+            })
+    }
 }
diff --git a/bench-tps/src/bench_tps_client/rpc_client.rs b/bench-tps/src/bench_tps_client/rpc_client.rs
@@ -2,8 +2,8 @@ use {
     crate::bench_tps_client::{BenchTpsClient, Result},
     solana_client::rpc_client::RpcClient,
     solana_sdk::{
-        commitment_config::CommitmentConfig, epoch_info::EpochInfo, hash::Hash, message::Message,
-        pubkey::Pubkey, signature::Signature, transaction::Transaction,
+        account::Account, commitment_config::CommitmentConfig, epoch_info::EpochInfo, hash::Hash,
+        message::Message, pubkey::Pubkey, signature::Signature, transaction::Transaction,
     },
 };
 
@@ -80,4 +80,8 @@ impl BenchTpsClient for RpcClient {
         RpcClient::request_airdrop_with_blockhash(self, pubkey, lamports, recent_blockhash)
             .map_err(|err| err.into())
     }
+
+    fn get_account(&self, pubkey: &Pubkey) -> Result<Account> {
+        RpcClient::get_account(self, pubkey).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
@@ -2,6 +2,7 @@ use {
     crate::bench_tps_client::{BenchTpsClient, Result},
     solana_client::thin_client::ThinClient,
     solana_sdk::{
+        account::Account,
         client::{AsyncClient, Client, SyncClient},
         commitment_config::CommitmentConfig,
         epoch_info::EpochInfo,
@@ -83,4 +84,10 @@ impl BenchTpsClient for ThinClient {
             .request_airdrop_with_blockhash(pubkey, lamports, recent_blockhash)
             .map_err(|err| err.into())
     }
+
+    fn get_account(&self, pubkey: &Pubkey) -> Result<Account> {
+        self.rpc_client()
+            .get_account(pubkey)
+            .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
@@ -2,8 +2,8 @@ use {
     crate::bench_tps_client::{BenchTpsClient, Result},
     solana_client::tpu_client::TpuClient,
     solana_sdk::{
-        commitment_config::CommitmentConfig, epoch_info::EpochInfo, hash::Hash, message::Message,
-        pubkey::Pubkey, signature::Signature, transaction::Transaction,
+        account::Account, commitment_config::CommitmentConfig, epoch_info::EpochInfo, hash::Hash,
+        message::Message, pubkey::Pubkey, signature::Signature, transaction::Transaction,
     },
 };
 
@@ -96,4 +96,10 @@ impl BenchTpsClient for TpuClient {
             .request_airdrop_with_blockhash(pubkey, lamports, recent_blockhash)
             .map_err(|err| err.into())
     }
+
+    fn get_account(&self, pubkey: &Pubkey) -> Result<Account> {
+        self.rpc_client()
+            .get_account(pubkey)
+            .map_err(|err| err.into())
+    }
 }