From 318a67b39a98dfbe8659c8db88717aefbd69a445 Mon Sep 17 00:00:00 2001
From: Kirill Lykov <kirill.lykov@anza.xyz>
Date: Thu, 7 Mar 2024 17:01:08 +0000
Subject: [PATCH] add commitment level for get_slot/blocks

---
 bench-tps/src/bench_tps_client.rs             |  9 +++++++--
 bench-tps/src/bench_tps_client/bank_client.rs | 11 ++++++++---
 bench-tps/src/bench_tps_client/rpc_client.rs  | 14 ++++++++++----
 bench-tps/src/bench_tps_client/thin_client.rs | 15 +++++++++++----
 bench-tps/src/bench_tps_client/tpu_client.rs  | 15 +++++++++++----
 5 files changed, 47 insertions(+), 17 deletions(-)

diff --git a/bench-tps/src/bench_tps_client.rs b/bench-tps/src/bench_tps_client.rs
index 7bf2317d6a8751..0715d739879165 100644
--- a/bench-tps/src/bench_tps_client.rs
+++ b/bench-tps/src/bench_tps_client.rs
@@ -95,9 +95,14 @@ pub trait BenchTpsClient {
 
     fn get_multiple_accounts(&self, pubkeys: &[Pubkey]) -> Result<Vec<Option<Account>>>;
 
-    fn get_slot(&self) -> Result<Slot>;
+    fn get_slot_with_commitment(&self, commitment_config: CommitmentConfig) -> Result<Slot>;
 
-    fn get_blocks(&self, start_slot: Slot, end_slot: Option<Slot>) -> Result<Vec<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,
diff --git a/bench-tps/src/bench_tps_client/bank_client.rs b/bench-tps/src/bench_tps_client/bank_client.rs
index a2022110550fd6..3ea9080e51398a 100644
--- a/bench-tps/src/bench_tps_client/bank_client.rs
+++ b/bench-tps/src/bench_tps_client/bank_client.rs
@@ -115,11 +115,16 @@ impl BenchTpsClient for BankClient {
         unimplemented!("BankClient doesn't support get_multiple_accounts");
     }
 
-    fn get_slot(&self) -> Result<Slot> {
-        SyncClient::get_slot(self).map_err(|err| err.into())
+    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(&self, _start_slot: Slot, _end_slot: Option<Slot>) -> Result<Vec<Slot>> {
+    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");
     }
 
diff --git a/bench-tps/src/bench_tps_client/rpc_client.rs b/bench-tps/src/bench_tps_client/rpc_client.rs
index c13f837c813cca..87ec1b8690c417 100644
--- a/bench-tps/src/bench_tps_client/rpc_client.rs
+++ b/bench-tps/src/bench_tps_client/rpc_client.rs
@@ -108,12 +108,18 @@ impl BenchTpsClient for RpcClient {
         RpcClient::get_multiple_accounts(self, pubkeys).map_err(|err| err.into())
     }
 
-    fn get_slot(&self) -> Result<Slot> {
-        RpcClient::get_slot(self).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(&self, start_slot: Slot, end_slot: Option<Slot>) -> Result<Vec<Slot>> {
-        RpcClient::get_blocks(self, start_slot, end_slot).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(
diff --git a/bench-tps/src/bench_tps_client/thin_client.rs b/bench-tps/src/bench_tps_client/thin_client.rs
index b99d9bd02e68ba..22945c4494f453 100644
--- a/bench-tps/src/bench_tps_client/thin_client.rs
+++ b/bench-tps/src/bench_tps_client/thin_client.rs
@@ -114,13 +114,20 @@ impl BenchTpsClient for ThinClient {
             .map_err(|err| err.into())
     }
 
-    fn get_slot(&self) -> Result<Slot> {
-        self.rpc_client().get_slot().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(&self, start_slot: Slot, end_slot: Option<Slot>) -> Result<Vec<Slot>> {
+    fn get_blocks_with_commitment(
+        &self,
+        start_slot: Slot,
+        end_slot: Option<Slot>,
+        commitment_config: CommitmentConfig,
+    ) -> Result<Vec<Slot>> {
         self.rpc_client()
-            .get_blocks(start_slot, end_slot)
+            .get_blocks_with_commitment(start_slot, end_slot, commitment_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
index 02ebcb6c208ef3..6c053271ad3eec 100644
--- a/bench-tps/src/bench_tps_client/tpu_client.rs
+++ b/bench-tps/src/bench_tps_client/tpu_client.rs
@@ -134,13 +134,20 @@ where
             .map_err(|err| err.into())
     }
 
-    fn get_slot(&self) -> Result<Slot> {
-        self.rpc_client().get_slot().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(&self, start_slot: Slot, end_slot: Option<Slot>) -> Result<Vec<Slot>> {
+    fn get_blocks_with_commitment(
+        &self,
+        start_slot: Slot,
+        end_slot: Option<Slot>,
+        commitment_config: CommitmentConfig,
+    ) -> Result<Vec<Slot>> {
         self.rpc_client()
-            .get_blocks(start_slot, end_slot)
+            .get_blocks_with_commitment(start_slot, end_slot, commitment_config)
             .map_err(|err| err.into())
     }