From 91b237710236759d50927f8fc65f89b52dc8d6ae Mon Sep 17 00:00:00 2001
From: Ryo Onodera <ryoqun@gmail.com>
Date: Sun, 15 Dec 2024 15:33:17 +0900
Subject: [PATCH] Use BankWithScheduler in banking-bench for US

---
 Cargo.lock                |  1 +
 banking-bench/Cargo.toml  |  1 +
 banking-bench/src/main.rs | 10 +++++++---
 3 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index b010687a85f117..d12773560566f9 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -6053,6 +6053,7 @@ dependencies = [
 name = "solana-banking-bench"
 version = "2.2.0"
 dependencies = [
+ "assert_matches",
  "clap 3.2.23",
  "crossbeam-channel",
  "log",
diff --git a/banking-bench/Cargo.toml b/banking-bench/Cargo.toml
index 67ca53f88324c9..3bcaff24673478 100644
--- a/banking-bench/Cargo.toml
+++ b/banking-bench/Cargo.toml
@@ -9,6 +9,7 @@ license = { workspace = true }
 edition = { workspace = true }
 
 [dependencies]
+assert_matches = { workspace = true }
 clap = { version = "3.1.8", features = ["derive", "cargo"] }
 crossbeam-channel = { workspace = true }
 log = { workspace = true }
diff --git a/banking-bench/src/main.rs b/banking-bench/src/main.rs
index 99722c7cf6082e..97eacc07a807f5 100644
--- a/banking-bench/src/main.rs
+++ b/banking-bench/src/main.rs
@@ -1,5 +1,6 @@
 #![allow(clippy::arithmetic_side_effects)]
 use {
+    assert_matches::assert_matches,
     clap::{crate_description, crate_name, Arg, ArgEnum, Command},
     crossbeam_channel::{unbounded, Receiver},
     log::*,
@@ -349,7 +350,7 @@ fn main() {
     let (replay_vote_sender, _replay_vote_receiver) = unbounded();
     let bank0 = Bank::new_for_benches(&genesis_config);
     let bank_forks = BankForks::new_rw_arc(bank0);
-    let mut bank = bank_forks.read().unwrap().working_bank();
+    let mut bank = bank_forks.read().unwrap().working_bank_with_scheduler();
 
     // set cost tracker limits to MAX so it will not filter out TXs
     bank.write_cost_tracker()
@@ -552,13 +553,16 @@ fn main() {
             poh_time.stop();
 
             let mut new_bank_time = Measure::start("new_bank");
+            if let Some((result, _timings)) = bank.wait_for_completed_scheduler() {
+                assert_matches!(result, Ok(_));
+            }
             let new_slot = bank.slot() + 1;
-            let new_bank = Bank::new_from_parent(bank, &collector, new_slot);
+            let new_bank = Bank::new_from_parent(bank.clone(), &collector, new_slot);
             new_bank_time.stop();
 
             let mut insert_time = Measure::start("insert_time");
             bank_forks.write().unwrap().insert(new_bank);
-            bank = bank_forks.read().unwrap().working_bank();
+            bank = bank_forks.read().unwrap().working_bank_with_scheduler();
             insert_time.stop();
 
             assert!(poh_recorder.read().unwrap().bank().is_none());