From 60ab42363e8b0235f2ba93ee1ebc1c8ad3d059ec Mon Sep 17 00:00:00 2001
From: Ryo Onodera <ryoqun@gmail.com>
Date: Sat, 14 Dec 2024 22:21:21 +0900
Subject: [PATCH] Clean do_take_resumed_scheduler() a bit

---
 unified-scheduler-pool/src/lib.rs | 61 +++++++++++++++++--------------
 1 file changed, 34 insertions(+), 27 deletions(-)

diff --git a/unified-scheduler-pool/src/lib.rs b/unified-scheduler-pool/src/lib.rs
index 7eb838d1e80595..daca95f7826791 100644
--- a/unified-scheduler-pool/src/lib.rs
+++ b/unified-scheduler-pool/src/lib.rs
@@ -486,35 +486,42 @@ where
     ) -> S {
         assert_matches!(result_with_timings, (Ok(_), _));
 
-        if matches!(context.mode(), BlockVerification) {
-            // pop is intentional for filo, expecting relatively warmed-up scheduler due to having been
-            // returned recently
-            if let Some((inner, _pooled_at)) =
-                self.scheduler_inners.lock().expect("not poisoned").pop()
-            {
+        match context.mode() {
+            BlockVerification => {
+                // pop is intentional for filo, expecting relatively warmed-up scheduler due to having been
+                // returned recently
+                if let Some((inner, _pooled_at)) =
+                    self.scheduler_inners.lock().expect("not poisoned").pop()
+                {
+                    S::from_inner(inner, context, result_with_timings)
+                } else {
+                    S::spawn(
+                        self.block_verification_handler_count,
+                        self.self_arc(),
+                        context,
+                        result_with_timings,
+                        None,
+                    )
+                }
+            }
+            BlockProduction => {
+                let id_and_inner = self
+                    .block_production_scheduler_inner
+                    .lock()
+                    .expect("not poisoned");
+                // Wait with condvar until scheduler is spawned elsewhere, which is guaranteed to
+                // be done eventually.
+                let (id, inner) = &mut *self
+                    .block_production_scheduler_condvar
+                    .wait_while(id_and_inner, |id_and_inner| id_and_inner.0.is_none())
+                    .unwrap();
+
+                let Some(inner) = inner.take() else {
+                    panic!("already taken: id: {id:?} slot: {:?}", context.slot());
+                };
+
                 S::from_inner(inner, context, result_with_timings)
-            } else {
-                S::spawn(
-                    self.block_verification_handler_count,
-                    self.self_arc(),
-                    context,
-                    result_with_timings,
-                    None,
-                )
             }
-        } else {
-            let mut id_and_inner = self
-                .block_production_scheduler_inner
-                .lock()
-                .expect("not poisoned");
-            id_and_inner = self
-                .block_production_scheduler_condvar
-                .wait_while(id_and_inner, |id_and_inner| id_and_inner.0.is_none())
-                .unwrap();
-            let Some(inner) = id_and_inner.1.take() else {
-                panic!("double take: {:?}, {:?}", context.slot(), context.mode());
-            };
-            S::from_inner(inner, context, result_with_timings)
         }
     }