From ee425389542fb70b88b65150ad2a02653da5d762 Mon Sep 17 00:00:00 2001
From: Ryo Onodera <ryoqun@gmail.com>
Date: Mon, 3 Mar 2025 21:25:41 +0900
Subject: [PATCH] more clean up and enable local cluster test

---
 local-cluster/tests/local_cluster.rs |  1 -
 unified-scheduler-logic/src/lib.rs   |  8 --------
 unified-scheduler-pool/src/lib.rs    | 24 +++++++++++++-----------
 3 files changed, 13 insertions(+), 20 deletions(-)

diff --git a/local-cluster/tests/local_cluster.rs b/local-cluster/tests/local_cluster.rs
index d2be55188232a3..7a0a2ce30f294f 100644
--- a/local-cluster/tests/local_cluster.rs
+++ b/local-cluster/tests/local_cluster.rs
@@ -5850,7 +5850,6 @@ fn test_randomly_mixed_block_verification_methods_between_bootstrap_and_not() {
 
 #[test]
 #[serial]
-#[ignore]
 fn test_randomly_mixed_block_production_methods_between_bootstrap_and_not() {
     // tailored logging just to see two block production methods are working correctly
     solana_logger::setup_with_default(
diff --git a/unified-scheduler-logic/src/lib.rs b/unified-scheduler-logic/src/lib.rs
index 6f252650449702..79e764fe2976a5 100644
--- a/unified-scheduler-logic/src/lib.rs
+++ b/unified-scheduler-logic/src/lib.rs
@@ -839,14 +839,6 @@ impl SchedulingStateMachine {
         transaction: RuntimeTransaction<SanitizedTransaction>,
         index: usize,
         usage_queue_loader: &mut impl FnMut(Pubkey) -> UsageQueue,
-    ) -> Task {
-        Self::do_create_task(transaction, index, usage_queue_loader)
-    }
-
-    pub fn do_create_task(
-        transaction: RuntimeTransaction<SanitizedTransaction>,
-        index: usize,
-        usage_queue_loader: &mut impl FnMut(Pubkey) -> UsageQueue,
     ) -> Task {
         // It's crucial for tasks to be validated with
         // `account_locks::validate_account_locks()` prior to the creation.
diff --git a/unified-scheduler-pool/src/lib.rs b/unified-scheduler-pool/src/lib.rs
index 574ec245c77f84..ec217c8d0f2c9c 100644
--- a/unified-scheduler-pool/src/lib.rs
+++ b/unified-scheduler-pool/src/lib.rs
@@ -240,6 +240,10 @@ impl HandlerContext {
             },
         }
     }
+
+    fn banking_stage_helper(&self) -> &BankingStageHelper {
+        self.banking_stage_helper.as_ref().unwrap()
+    }
 }
 
 #[derive(Debug, Clone)]
@@ -322,7 +326,7 @@ impl BankingStageHelper {
         transaction: RuntimeTransaction<SanitizedTransaction>,
         index: usize,
     ) -> Task {
-        SchedulingStateMachine::do_create_task(transaction, index, &mut |pubkey| {
+        SchedulingStateMachine::create_task(transaction, index, &mut |pubkey| {
             self.usage_queue_loader.load(pubkey)
         })
     }
@@ -335,9 +339,9 @@ impl BankingStageHelper {
         self.do_create_task(transaction, index)
     }
 
-    fn recreate_task(&self, task: Task) -> Task {
+    fn recreate_task(&self, executed_task: Box<ExecutedTask>) -> Task {
         let new_index = self.generate_task_ids(1);
-        let transaction = task.into_transaction();
+        let transaction = executed_task.into_task().into_transaction();
         self.do_create_task(transaction, new_index)
     }
 
@@ -1008,7 +1012,7 @@ impl ExecutedTask {
         })
     }
 
-    fn into_inner(self) -> Task {
+    fn into_task(self) -> Task {
         self.task
     }
 }
@@ -1717,9 +1721,7 @@ impl<S: SpawnableScheduler<TH>, TH: TaskHandler> ThreadManager<S, TH> {
                                 if trigger_ending {
                                     assert_matches!(scheduling_mode, BlockProduction);
                                     session_ending = true;
-                                    let task = handler_context.banking_stage_helper.as_ref().unwrap().recreate_task(
-                                        executed_task.into_inner(),
-                                    );
+                                    let task = handler_context.banking_stage_helper().recreate_task(executed_task);
                                     state_machine.buffer_task(task);
                                 }
                             },
@@ -1773,9 +1775,7 @@ impl<S: SpawnableScheduler<TH>, TH: TaskHandler> ThreadManager<S, TH> {
                                 if trigger_ending {
                                     assert_matches!(scheduling_mode, BlockProduction);
                                     session_ending = true;
-                                    let task = handler_context.banking_stage_helper.as_ref().unwrap().recreate_task(
-                                        executed_task.into_inner(),
-                                    );
+                                    let task = handler_context.banking_stage_helper().recreate_task(executed_task);
                                     state_machine.buffer_task(task);
                                 }
                             },
@@ -1831,7 +1831,7 @@ impl<S: SpawnableScheduler<TH>, TH: TaskHandler> ThreadManager<S, TH> {
                                 runnable_task_sender
                                     .send_chained_channel(&new_context, handler_count)
                                     .unwrap();
-                                if scheduling_mode == BlockVerification {
+                                if matches!(scheduling_mode, BlockVerification) {
                                     result_with_timings = new_result_with_timings;
                                     break;
                                 }
@@ -1849,6 +1849,8 @@ impl<S: SpawnableScheduler<TH>, TH: TaskHandler> ThreadManager<S, TH> {
                             }
                             Ok(NewTaskPayload::Unblock) => {
                                 assert_matches!(scheduling_mode, BlockProduction);
+                                // borrow checker won't allow us to do this inside OpenSubchannel
+                                // to remove duplicate assignments...
                                 result_with_timings = new_result_with_timings;
                                 break;
                             }