Add get_account_with_commitment to BenchTpsClient (#27176)

solana-labs · Aug 18, 2022 · fda395a · fda395a
1 parent 7d765e3
commit fda395a
Show file tree

Hide file tree

Showing 5 changed files with 69 additions and 3 deletions.
diff --git a/bench-tps/src/bench_tps_client.rs b/bench-tps/src/bench_tps_client.rs
@@ -83,6 +83,13 @@ pub trait BenchTpsClient {
 
     /// Returns all information associated with the account of the provided pubkey
     fn get_account(&self, pubkey: &Pubkey) -> Result<Account>;
+
+    /// Returns all information associated with the account of the provided pubkey, using explicit commitment
+    fn get_account_with_commitment(
+        &self,
+        pubkey: &Pubkey,
+        commitment_config: CommitmentConfig,
+    ) -> 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
@@ -93,4 +93,18 @@ impl BenchTpsClient for BankClient {
                 })
             })
     }
+
+    fn get_account_with_commitment(
+        &self,
+        pubkey: &Pubkey,
+        commitment_config: CommitmentConfig,
+    ) -> Result<Account> {
+        SyncClient::get_account_with_commitment(self, pubkey, commitment_config)
+            .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
@@ -1,5 +1,5 @@
 use {
-    crate::bench_tps_client::{BenchTpsClient, Result},
+    crate::bench_tps_client::{BenchTpsClient, BenchTpsError, Result},
     solana_client::rpc_client::RpcClient,
     solana_sdk::{
         account::Account, commitment_config::CommitmentConfig, epoch_info::EpochInfo, hash::Hash,
@@ -84,4 +84,19 @@ impl BenchTpsClient for RpcClient {
     fn get_account(&self, pubkey: &Pubkey) -> Result<Account> {
         RpcClient::get_account(self, pubkey).map_err(|err| err.into())
     }
+
+    fn get_account_with_commitment(
+        &self,
+        pubkey: &Pubkey,
+        commitment_config: CommitmentConfig,
+    ) -> Result<Account> {
+        RpcClient::get_account_with_commitment(self, pubkey, commitment_config)
+            .map(|res| res.value)
+            .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/thin_client.rs b/bench-tps/src/bench_tps_client/thin_client.rs
@@ -1,5 +1,5 @@
 use {
-    crate::bench_tps_client::{BenchTpsClient, Result},
+    crate::bench_tps_client::{BenchTpsClient, BenchTpsError, Result},
     solana_client::thin_client::ThinClient,
     solana_sdk::{
         account::Account,
@@ -90,4 +90,18 @@ impl BenchTpsClient for ThinClient {
             .get_account(pubkey)
             .map_err(|err| err.into())
     }
+
+    fn get_account_with_commitment(
+        &self,
+        pubkey: &Pubkey,
+        commitment_config: CommitmentConfig,
+    ) -> Result<Account> {
+        SyncClient::get_account_with_commitment(self, pubkey, commitment_config)
+            .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/tpu_client.rs b/bench-tps/src/bench_tps_client/tpu_client.rs
@@ -1,5 +1,5 @@
 use {
-    crate::bench_tps_client::{BenchTpsClient, Result},
+    crate::bench_tps_client::{BenchTpsClient, BenchTpsError, Result},
     solana_client::tpu_client::TpuClient,
     solana_sdk::{
         account::Account, commitment_config::CommitmentConfig, epoch_info::EpochInfo, hash::Hash,
@@ -102,4 +102,20 @@ impl BenchTpsClient for TpuClient {
             .get_account(pubkey)
             .map_err(|err| err.into())
     }
+
+    fn get_account_with_commitment(
+        &self,
+        pubkey: &Pubkey,
+        commitment_config: CommitmentConfig,
+    ) -> Result<Account> {
+        self.rpc_client()
+            .get_account_with_commitment(pubkey, commitment_config)
+            .map(|res| res.value)
+            .map_err(|err| err.into())
+            .and_then(|account| {
+                account.ok_or_else(|| {
+                    BenchTpsError::Custom(format!("AccountNotFound: pubkey={}", pubkey))
+                })
+            })
+    }
 }