Wen restart: play unfrozen block

wen-restart/src/wen_restart.rs

@@ -618,59 +618,49 @@ pub(crate) fn find_bankhash_of_heaviest_fork(
         if exit.load(Ordering::Relaxed) {
             return Err(WenRestartError::Exiting.into());
         }
-        let saved_bank;
-        {
-            saved_bank = bank_forks.read().unwrap().get(slot);
-        }
-        let bank = match saved_bank {
-            Some(cur_bank) => {
-                if !cur_bank.is_frozen() {
-                    return Err(WenRestartError::BlockNotFrozenAfterReplay(slot, None).into());
-                }
-                cur_bank
-            }
-            None => {
-                let new_bank = Bank::new_from_parent(
-                    parent_bank.clone(),
-                    &leader_schedule_cache
-                        .slot_leader_at(slot, Some(&parent_bank))
-                        .unwrap(),
-                    slot,
+        let saved_bank = bank_forks.read().unwrap().get_with_scheduler(slot);
+        let bank_with_scheduler = if let Some(bank) = saved_bank {
+            bank
+        } else {
+            let new_bank = Bank::new_from_parent(
+                parent_bank.clone(),
+                &leader_schedule_cache
+                    .slot_leader_at(slot, Some(&parent_bank))
+                    .unwrap(),
+                slot,
+            );
+            bank_forks.write().unwrap().insert_from_ledger(new_bank)
+        };
+        let bank = if bank_with_scheduler.is_frozen() {
+            bank_with_scheduler.clone_without_scheduler()
+        } else {
+            let mut progress = ConfirmationProgress::new(parent_bank.last_blockhash());
+            if let Err(e) = process_single_slot(
+                &blockstore,
+                &bank_with_scheduler,
+                &replay_tx_thread_pool,
+                &opts,
+                &recyclers,
+                &mut progress,
+                None,
+                None,
+                None,
+                None,
+                &mut timing,
+            ) {
+                return Err(
+                    WenRestartError::BlockNotFrozenAfterReplay(slot, Some(e.to_string())).into(),
                 );
-                let bank_with_scheduler;
-                {
-                    bank_with_scheduler = bank_forks.write().unwrap().insert_from_ledger(new_bank);
-                }
-                let mut progress = ConfirmationProgress::new(parent_bank.last_blockhash());
-                if let Err(e) = process_single_slot(
-                    &blockstore,
-                    &bank_with_scheduler,
-                    &replay_tx_thread_pool,
-                    &opts,
-                    &recyclers,
-                    &mut progress,
-                    None,
-                    None,
-                    None,
-                    None,
-                    &mut timing,
-                ) {
-                    return Err(WenRestartError::BlockNotFrozenAfterReplay(
-                        slot,
-                        Some(e.to_string()),
-                    )
-                    .into());
-                }
-                let cur_bank;
-                {
-                    cur_bank = bank_forks
-                        .read()
-                        .unwrap()
-                        .get(slot)
-                        .expect("bank should have been just inserted");
-                }
-                cur_bank
             }
+            let cur_bank;
+            {
+                cur_bank = bank_forks
+                    .read()
+                    .unwrap()
+                    .get(slot)
+                    .expect("bank should have been just inserted");
+            }
+            cur_bank
         };
         parent_bank = bank;
     }
@@ -3752,4 +3742,73 @@ mod tests {
         assert_eq!(pushed_slot, my_slot);
         assert_eq!(pushed_hash, my_hash);
     }
+
+    fn run_and_check_find_bankhash_of_heaviest_fork(
+        test_state: &WenRestartTestInitResult,
+        slots: &[Slot],
+        slot: Slot,
+    ) {
+        let exit = Arc::new(AtomicBool::new(false));
+        assert_eq!(
+            find_bankhash_of_heaviest_fork(
+                slot,
+                slots.to_vec(),
+                test_state.blockstore.clone(),
+                test_state.bank_forks.clone(),
+                test_state.bank_forks.read().unwrap().root_bank(),
+                &exit,
+            )
+            .unwrap(),
+            test_state
+                .bank_forks
+                .read()
+                .unwrap()
+                .get(slot)
+                .unwrap()
+                .hash()
+        );
+    }
+
+    #[test]
+    fn test_find_bankhash_of_heaviest_fork() {
+        let ledger_path = get_tmp_ledger_path_auto_delete!();
+        let test_state = wen_restart_test_init(&ledger_path);
+        let last_vote = test_state.last_voted_fork_slots[0];
+        let mut slots = test_state.last_voted_fork_slots.clone();
+        slots.reverse();
+        run_and_check_find_bankhash_of_heaviest_fork(&test_state, &slots, last_vote);
+        let new_slot = last_vote + 1;
+        let _ = insert_slots_into_blockstore(
+            test_state.blockstore.clone(),
+            last_vote,
+            &[new_slot],
+            TICKS_PER_SLOT,
+            test_state.last_blockhash,
+        );
+        slots.push(new_slot);
+        run_and_check_find_bankhash_of_heaviest_fork(&test_state, &slots, new_slot);
+        let slot_full_but_not_replayed = last_vote + 2;
+        let _ = insert_slots_into_blockstore(
+            test_state.blockstore.clone(),
+            last_vote,
+            &[slot_full_but_not_replayed],
+            TICKS_PER_SLOT,
+            test_state.last_blockhash,
+        );
+        let new_bank = Bank::new_from_parent(
+            test_state.bank_forks.read().unwrap().get(new_slot).unwrap(),
+            &Pubkey::default(),
+            slot_full_but_not_replayed,
+        );
+        let _ = test_state
+            .bank_forks
+            .write()
+            .unwrap()
+            .insert_from_ledger(new_bank);
+        run_and_check_find_bankhash_of_heaviest_fork(
+            &test_state,
+            &slots,
+            slot_full_but_not_replayed,
+        );
+    }
 }