From bac313a9986adff2ca77431f5276f9d90e463059 Mon Sep 17 00:00:00 2001
From: Ryo Onodera <ryoqun@gmail.com>
Date: Fri, 20 Dec 2024 09:38:29 +0900
Subject: [PATCH] fix ci?

---
 unified-scheduler-pool/src/lib.rs | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/unified-scheduler-pool/src/lib.rs b/unified-scheduler-pool/src/lib.rs
index c8b4eed328f17d..3791e3c0d4e21e 100644
--- a/unified-scheduler-pool/src/lib.rs
+++ b/unified-scheduler-pool/src/lib.rs
@@ -1468,6 +1468,7 @@ impl<S: SpawnableScheduler<TH>, TH: TaskHandler> ThreadManager<S, TH> {
                         .send(result_with_timings)
                         .expect("always outlived receiver");
 
+                    let mut new_result_with_timings = initialized_result_with_timings();
                     loop {
                         if session_resetting {
                             while let Some(task) = state_machine.schedule_next_unblocked_task() {
@@ -1482,7 +1483,8 @@ impl<S: SpawnableScheduler<TH>, TH: TaskHandler> ThreadManager<S, TH> {
                                 Ok(NewTaskPayload::OpenSubchannel(context_and_result_with_timings)),
                                 _,
                             ) => {
-                                let (new_context, new_result_with_timings) =
+                                let new_context;
+                                (new_context, new_result_with_timings) =
                                     *context_and_result_with_timings;
                                 // We just received subsequent (= not initial) session and about to
                                 // enter into the preceding `while(!is_finished) {...}` loop again.
@@ -1494,20 +1496,16 @@ impl<S: SpawnableScheduler<TH>, TH: TaskHandler> ThreadManager<S, TH> {
                                     session_ending = false;
                                 } else {
                                     session_pausing = false;
-                                    // Prevent processing transactions and thus touching poh until
-                                    // unblocked, which signals successful poh `set_bank()`-ing.
-                                    assert_matches!(
-                                        new_task_receiver.recv().unwrap(),
-                                        NewTaskPayload::Unblock
-                                    );
                                 }
 
                                 runnable_task_sender
                                     .send_chained_channel(&new_context, handler_count)
                                     .unwrap();
                                 context = new_context;
-                                result_with_timings = new_result_with_timings;
-                                break;
+                                if context.mode() == BlockVerification {
+                                    result_with_timings = new_result_with_timings;
+                                    break;
+                                }
                             }
                             (Ok(NewTaskPayload::CloseSubchannel), BlockProduction) => {
                                 info!("ignoring duplicate CloseSubchannel...");
@@ -1518,6 +1516,10 @@ impl<S: SpawnableScheduler<TH>, TH: TaskHandler> ThreadManager<S, TH> {
                             (Ok(NewTaskPayload::Payload(task)), BlockProduction) => {
                                 assert!(state_machine.do_schedule_task(task, true).is_none());
                             }
+                            (Ok(NewTaskPayload::Unblock), BlockProduction) => {
+                                result_with_timings = new_result_with_timings;
+                                break;
+                            }
                             (Ok(NewTaskPayload::Disconnect), BlockProduction) | (Err(_), _) => {
                                 // This unusual condition must be triggered by ThreadManager::drop().
                                 // Initialize result_with_timings with a harmless value...