From 5b1b19b93f04514b78314a658db5c727ded4b291 Mon Sep 17 00:00:00 2001
From: Ryo Onodera <ryoqun@gmail.com>
Date: Sun, 15 Dec 2024 16:12:47 +0900
Subject: [PATCH] Call wait_for_completed_scheduler in banking simulator

---
 core/Cargo.toml                | 1 +
 core/src/banking_simulation.rs | 4 ++++
 programs/sbf/Cargo.lock        | 1 +
 svm/examples/Cargo.lock        | 1 +
 4 files changed, 7 insertions(+)

diff --git a/core/Cargo.toml b/core/Cargo.toml
index e1a6ddb1a6f928..5f2eb1af2f22ae 100644
--- a/core/Cargo.toml
+++ b/core/Cargo.toml
@@ -18,6 +18,7 @@ agave-banking-stage-ingress-types = { workspace = true }
 ahash = { workspace = true }
 anyhow = { workspace = true }
 arrayvec = { workspace = true }
+assert_matches = { workspace = true }
 base64 = { workspace = true }
 bincode = { workspace = true }
 bs58 = { workspace = true }
diff --git a/core/src/banking_simulation.rs b/core/src/banking_simulation.rs
index d1ed3144f9ccde..e811c9c6df9bd8 100644
--- a/core/src/banking_simulation.rs
+++ b/core/src/banking_simulation.rs
@@ -11,6 +11,7 @@ use {
         validator::BlockProductionMethod,
     },
     agave_banking_stage_ingress_types::BankingPacketBatch,
+    assert_matches::assert_matches,
     bincode::deserialize_from,
     crossbeam_channel::{unbounded, Sender},
     itertools::Itertools,
@@ -452,6 +453,9 @@ impl SimulatorLoop {
                 info!("Bank::new_from_parent()!");
 
                 logger.log_jitter(&bank);
+                if let Some((result, _execute_timings)) = bank.wait_for_completed_scheduler() {
+                    assert_matches!(result, Ok(()));
+                }
                 bank.freeze();
                 let new_slot = if bank.slot() == self.parent_slot {
                     info!("initial leader block!");
diff --git a/programs/sbf/Cargo.lock b/programs/sbf/Cargo.lock
index 59ddb6d61dbc5a..5e638153a072aa 100644
--- a/programs/sbf/Cargo.lock
+++ b/programs/sbf/Cargo.lock
@@ -5461,6 +5461,7 @@ dependencies = [
  "ahash 0.8.11",
  "anyhow",
  "arrayvec",
+ "assert_matches",
  "base64 0.22.1",
  "bincode",
  "bs58",
diff --git a/svm/examples/Cargo.lock b/svm/examples/Cargo.lock
index 10fe2bb241f742..7ec26911ca6477 100644
--- a/svm/examples/Cargo.lock
+++ b/svm/examples/Cargo.lock
@@ -5312,6 +5312,7 @@ dependencies = [
  "ahash 0.8.11",
  "anyhow",
  "arrayvec",
+ "assert_matches",
  "base64 0.22.1",
  "bincode",
  "bs58",