From 8fd839041838da0944f2e556752b79d7b6edfc1e Mon Sep 17 00:00:00 2001
From: Ashwin Sekar <ashwin@solana.com>
Date: Tue, 9 Jul 2024 18:56:29 +0000
Subject: [PATCH] add tests

---
 core/src/consensus.rs              |  8 ++++
 core/src/replay_stage.rs           | 77 ++++++++++++++++++++++++++++++
 ledger/src/blockstore_processor.rs | 13 +++++
 3 files changed, 98 insertions(+)

diff --git a/core/src/consensus.rs b/core/src/consensus.rs
index cd1201429fec11..d10c30d6c6ce6a 100644
--- a/core/src/consensus.rs
+++ b/core/src/consensus.rs
@@ -3243,6 +3243,10 @@ pub mod test {
         assert!(!blockstore.is_root(1));
         assert!(!blockstore.is_root(3));
         assert!(!blockstore.is_root(4));
+        for i in 0..4 {
+            blockstore.store_duplicate_slot(i, vec![], vec![]).unwrap();
+            assert!(blockstore.has_duplicate_shreds_in_slot(i));
+        }
 
         let mut tower = Tower::default();
         tower.vote_state.root_slot = Some(4);
@@ -3254,9 +3258,13 @@ pub mod test {
         .unwrap();
 
         assert!(!blockstore.is_root(0));
+        assert!(blockstore.has_duplicate_shreds_in_slot(0));
         assert!(blockstore.is_root(1));
+        assert!(!blockstore.has_duplicate_shreds_in_slot(1));
         assert!(!blockstore.is_root(3));
+        assert!(blockstore.has_duplicate_shreds_in_slot(3));
         assert!(blockstore.is_root(4));
+        assert!(!blockstore.has_duplicate_shreds_in_slot(4));
     }
 
     #[test]
diff --git a/core/src/replay_stage.rs b/core/src/replay_stage.rs
index 627e0175c89e71..2cef81bf233c59 100644
--- a/core/src/replay_stage.rs
+++ b/core/src/replay_stage.rs
@@ -4521,6 +4521,7 @@ pub(crate) mod tests {
             replay_stage::ReplayStage,
             vote_simulator::{self, VoteSimulator},
         },
+        blockstore_processor::{fill_blockstore_slot_with_ticks, ProcessOptions},
         crossbeam_channel::unbounded,
         itertools::Itertools,
         solana_entry::entry::{self, Entry},
@@ -9032,4 +9033,80 @@ pub(crate) mod tests {
         assert_eq!(tower.vote_state, expected_tower.vote_state);
         assert_eq!(tower.node_pubkey, expected_tower.node_pubkey);
     }
+
+    #[test]
+    fn test_initialize_progress_and_fork_choice_with_duplicates() {
+        let GenesisConfigInfo {
+            mut genesis_config, ..
+        } = create_genesis_config(123);
+
+        let ticks_per_slot = 1;
+        genesis_config.ticks_per_slot = ticks_per_slot;
+        let (ledger_path, blockhash) =
+            solana_ledger::create_new_tmp_ledger_auto_delete!(&genesis_config);
+        let blockstore = Blockstore::open(ledger_path.path()).unwrap();
+
+        /*
+          Bank forks with:
+               slot 0
+                 |
+               slot 1 -> Duplicate before restart
+        */
+
+        let mut last_hash = blockhash;
+        for i in 0..1 {
+            last_hash =
+                fill_blockstore_slot_with_ticks(&blockstore, ticks_per_slot, i + 1, i, last_hash);
+        }
+        blockstore.set_roots([0, 1].iter()).unwrap();
+
+        // Set up bank0
+        let bank_forks = BankForks::new_rw_arc(Bank::new_for_tests(&genesis_config));
+        let bank0 = bank_forks.read().unwrap().get_with_scheduler(0).unwrap();
+
+        // Mark block 1 as duplicate
+        blockstore.store_duplicate_slot(1, vec![], vec![]).unwrap();
+
+        let bank1 = bank_forks.write().unwrap().insert(Bank::new_from_parent(
+            bank0.clone_without_scheduler(),
+            &Pubkey::default(),
+            1,
+        ));
+        bank_forks
+            .write()
+            .unwrap()
+            .set_root(
+                1,
+                &solana_runtime::accounts_background_service::AbsRequestSender::default(),
+                None,
+            )
+            .unwrap();
+
+        let leader_schedule_cache = LeaderScheduleCache::new_from_bank(&bank1);
+
+        // process_blockstore_from_root() from slot 1 onwards, should remove the duplicate proof from before the restart
+        blockstore_processor::process_blockstore_from_root(
+            &blockstore,
+            &bank_forks,
+            &leader_schedule_cache,
+            &ProcessOptions::default(),
+            None,
+            None,
+            None,
+            &AbsRequestSender::default(),
+        )
+        .unwrap();
+
+        // Verify that the duplicate information has been cleared and that fork choice can be initialized
+        let (_progress, fork_choice) =
+            ReplayStage::initialize_progress_and_fork_choice_with_locked_bank_forks(
+                &bank_forks,
+                &Pubkey::new_unique(),
+                &Pubkey::new_unique(),
+                &blockstore,
+            );
+
+        assert_eq!(fork_choice.tree_root().0, 1);
+        assert_eq!(fork_choice.best_overall_slot().0, 1);
+    }
 }
diff --git a/ledger/src/blockstore_processor.rs b/ledger/src/blockstore_processor.rs
index c21ebda1cd55c1..84d0c33485493c 100644
--- a/ledger/src/blockstore_processor.rs
+++ b/ledger/src/blockstore_processor.rs
@@ -3992,6 +3992,14 @@ pub mod tests {
             accounts_db_test_hash_calculation: true,
             ..ProcessOptions::default()
         };
+
+        // Mark all blocks 1 - 6 as duplicate
+        for slot in 1..6 {
+            blockstore
+                .store_duplicate_slot(slot, vec![], vec![])
+                .unwrap();
+        }
+
         let recyclers = VerifyRecyclers::default();
         let replay_tx_thread_pool = create_thread_pool(1);
         process_bank_0(
@@ -4065,6 +4073,11 @@ pub mod tests {
 
         // Check that bank forks has the correct banks
         verify_fork_infos(&bank_forks);
+
+        // Verify that the duplicate information has been cleared and that fork choice can be initialized
+        for slot in 1..6 {
+            assert!(!blockstore.has_duplicate_shreds_in_slot(slot));
+        }
     }
 
     #[test]