From 5abb627b3c387674bc0a3b1375ea7e263d8cd632 Mon Sep 17 00:00:00 2001
From: Wen <113942165+wen-coding@users.noreply.github.com>
Date: Thu, 7 Nov 2024 12:20:48 -0800
Subject: [PATCH] Wen restart fix flaky test (#3525)

* Fix flaky test, make sure bank is frozen before reading bankhash.

* Add more comments here.
---
 wen-restart/src/wen_restart.rs | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/wen-restart/src/wen_restart.rs b/wen-restart/src/wen_restart.rs
index 0c79f32eea8bbe..4557cab0cc155d 100644
--- a/wen-restart/src/wen_restart.rs
+++ b/wen-restart/src/wen_restart.rs
@@ -1817,8 +1817,11 @@ mod tests {
                 .unwrap()
                 .get(my_heaviest_fork_slot)
             {
-                my_heaviest_fork_bankhash = bank.hash();
-                break;
+                // When deciding the local heaviest fork, we will freeze the bank.
+                if bank.is_frozen() {
+                    my_heaviest_fork_bankhash = bank.hash();
+                    break;
+                }
             }
             sleep(Duration::from_millis(100));
         }