From 88dcdc2baacf03f83e454c42dd5c701afe165725 Mon Sep 17 00:00:00 2001
From: Ryo Onodera <ryoqun@gmail.com>
Date: Sat, 28 Oct 2023 22:13:35 +0900
Subject: [PATCH] Clean up wait_for_* fns

---
 runtime/src/bank.rs                     |  2 +-
 runtime/src/installed_scheduler_pool.rs | 41 +++++++++++++------------
 scheduler-pool/src/lib.rs               |  2 +-
 3 files changed, 23 insertions(+), 22 deletions(-)

diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
index b9dac4cb814e0b..fb2f0ce6953199 100644
--- a/runtime/src/bank.rs
+++ b/runtime/src/bank.rs
@@ -4193,7 +4193,7 @@ impl Bank {
     ) {
         // This is needed because recent_blockhash updates necessitate synchronizations for
         // consistent tx check_age handling.
-        BankWithScheduler::wait_for_reusable_scheduler(self, scheduler);
+        BankWithScheduler::wait_for_paused_scheduler(self, scheduler);
 
         // Only acquire the write lock for the blockhash queue on block boundaries because
         // readers can starve this write lock acquisition and ticks would be slowed down too
diff --git a/runtime/src/installed_scheduler_pool.rs b/runtime/src/installed_scheduler_pool.rs
index 78af30c6c30420..c5ed48000824ca 100644
--- a/runtime/src/installed_scheduler_pool.rs
+++ b/runtime/src/installed_scheduler_pool.rs
@@ -299,12 +299,23 @@ impl BankWithScheduler {
 
     #[must_use]
     pub fn wait_for_completed_scheduler(&self) -> Option<ResultWithTimings> {
-        self.inner
-            .wait_for_scheduler(WaitReason::TerminatedToFreeze)
+        BankWithSchedulerInner::wait_for_scheduler(
+            &self.inner.bank,
+            &self.inner.scheduler,
+            WaitReason::TerminatedToFreeze,
+        )
     }
 
-    pub(crate) fn wait_for_reusable_scheduler(bank: &Bank, scheduler: &InstalledSchedulerRwLock) {
-        BankWithSchedulerInner::wait_for_reusable_scheduler(bank, scheduler);
+    pub(crate) fn wait_for_paused_scheduler(bank: &Bank, scheduler: &InstalledSchedulerRwLock) {
+        let maybe_result_with_timings = BankWithSchedulerInner::wait_for_scheduler(
+            bank,
+            scheduler,
+            WaitReason::PausedForRecentBlockhash,
+        );
+        assert!(
+            maybe_result_with_timings.is_none(),
+            "Premature result was returned from scheduler after paused"
+        );
     }
 
     // take needless &mut only to communicate its semantic mutability to humans...
@@ -338,12 +349,7 @@ impl BankWithScheduler {
 
 impl BankWithSchedulerInner {
     #[must_use]
-    fn wait_for_scheduler(&self, reason: WaitReason) -> Option<ResultWithTimings> {
-        Self::do_wait_for_scheduler(&self.bank, &self.scheduler, reason)
-    }
-
-    #[must_use]
-    fn do_wait_for_scheduler(
+    fn wait_for_scheduler(
         bank: &Bank,
         scheduler: &InstalledSchedulerRwLock,
         reason: WaitReason,
@@ -379,17 +385,12 @@ impl BankWithSchedulerInner {
 
     #[must_use]
     fn wait_for_completed_scheduler_from_drop(&self) -> Option<Result<()>> {
-        let maybe_result_with_timings = self.wait_for_scheduler(WaitReason::DroppedFromBankForks);
-        maybe_result_with_timings.map(|(result, _timings)| result)
-    }
-
-    fn wait_for_reusable_scheduler(bank: &Bank, scheduler: &InstalledSchedulerRwLock) {
-        let maybe_result_with_timings =
-            Self::do_wait_for_scheduler(bank, scheduler, WaitReason::PausedForRecentBlockhash);
-        assert!(
-            maybe_result_with_timings.is_none(),
-            "Premature result was returned from scheduler after paused"
+        let maybe_result_with_timings = Self::wait_for_scheduler(
+            &self.bank,
+            &self.scheduler,
+            WaitReason::DroppedFromBankForks,
         );
+        maybe_result_with_timings.map(|(result, _timings)| result)
     }
 
     fn drop_scheduler(&self) {
diff --git a/scheduler-pool/src/lib.rs b/scheduler-pool/src/lib.rs
index a0d320c94bc61c..139512e8013381 100644
--- a/scheduler-pool/src/lib.rs
+++ b/scheduler-pool/src/lib.rs
@@ -638,7 +638,7 @@ mod tests {
 
         fn wait_for_termination(&mut self, reason: &WaitReason) -> Option<ResultWithTimings> {
             if TRIGGER_RACE_CONDITION && matches!(reason, WaitReason::PausedForRecentBlockhash) {
-                // this is equivalent to NOT calling wait_for_reusable_scheduler() in
+                // this is equivalent to NOT calling wait_for_paused_scheduler() in
                 // register_recent_blockhash().
                 return None;
             }