Use BankForks on tests - Part 4 (solana-labs#34271)

* Use BankForks on tests - Part 4

* Ensure the correct slot is set

bench-tps/src/bench.rs

@@ -1160,17 +1160,17 @@ mod tests {
         },
     };
 
-    fn bank_with_all_features(genesis_config: &GenesisConfig) -> Bank {
+    fn bank_with_all_features(genesis_config: &GenesisConfig) -> Arc<Bank> {
         let mut bank = Bank::new_for_tests(genesis_config);
         bank.feature_set = Arc::new(FeatureSet::all_enabled());
-        bank
+        bank.wrap_with_bank_forks_for_tests().0
     }
 
     #[test]
     fn test_bench_tps_bank_client() {
         let (genesis_config, id) = create_genesis_config(sol_to_lamports(10_000.0));
         let bank = bank_with_all_features(&genesis_config);
-        let client = Arc::new(BankClient::new(bank));
+        let client = Arc::new(BankClient::new_shared(bank));
 
         let config = Config {
             id,
@@ -1191,7 +1191,7 @@ mod tests {
     fn test_bench_tps_fund_keys() {
         let (genesis_config, id) = create_genesis_config(sol_to_lamports(10_000.0));
         let bank = bank_with_all_features(&genesis_config);
-        let client = Arc::new(BankClient::new(bank));
+        let client = Arc::new(BankClient::new_shared(bank));
         let keypair_count = 20;
         let lamports = 20;
         let rent = client.get_minimum_balance_for_rent_exemption(0).unwrap();
@@ -1216,7 +1216,7 @@ mod tests {
         let fee_rate_governor = FeeRateGovernor::new(11, 0);
         genesis_config.fee_rate_governor = fee_rate_governor;
         let bank = bank_with_all_features(&genesis_config);
-        let client = Arc::new(BankClient::new(bank));
+        let client = Arc::new(BankClient::new_shared(bank));
         let keypair_count = 20;
         let lamports = 20;
         let rent = client.get_minimum_balance_for_rent_exemption(0).unwrap();
@@ -1234,7 +1234,7 @@ mod tests {
     fn test_bench_tps_create_durable_nonce() {
         let (genesis_config, id) = create_genesis_config(sol_to_lamports(10_000.0));
         let bank = bank_with_all_features(&genesis_config);
-        let client = Arc::new(BankClient::new(bank));
+        let client = Arc::new(BankClient::new_shared(bank));
         let keypair_count = 10;
         let lamports = 10_000_000;
 

core/src/banking_stage/consumer.rs

@@ -1724,14 +1724,19 @@ mod tests {
             mint_keypair,
             ..
         } = create_slow_genesis_config(10_000);
-        let bank = Bank::new_no_wallclock_throttle_for_tests(&genesis_config).0;
+        let (bank, bank_forks) = Bank::new_no_wallclock_throttle_for_tests(&genesis_config);
         let keypair = Keypair::new();
 
         let address_table_key = Pubkey::new_unique();
         let address_table_state = generate_new_address_lookup_table(None, 2);
         store_address_lookup_table(&bank, address_table_key, address_table_state);
 
-        let bank = Arc::new(Bank::new_from_parent(bank, &Pubkey::new_unique(), 1));
+        let new_bank = Bank::new_from_parent(bank, &Pubkey::new_unique(), 2);
+        let bank = bank_forks
+            .write()
+            .unwrap()
+            .insert(new_bank)
+            .clone_without_scheduler();
         let message = VersionedMessage::V0(v0::Message {
             header: MessageHeader {
                 num_required_signatures: 1,

core/src/commitment_service.rs

@@ -263,6 +263,20 @@ mod tests {
         },
     };
 
+    fn new_bank_from_parent_with_bank_forks(
+        bank_forks: &RwLock<BankForks>,
+        parent: Arc<Bank>,
+        collector_id: &Pubkey,
+        slot: Slot,
+    ) -> Arc<Bank> {
+        let bank = Bank::new_from_parent(parent, collector_id, slot);
+        bank_forks
+            .write()
+            .unwrap()
+            .insert(bank)
+            .clone_without_scheduler()
+    }
+
     #[test]
     fn test_get_highest_super_majority_root() {
         assert_eq!(get_highest_super_majority_root(vec![], 10), 0);
@@ -508,14 +522,18 @@ mod tests {
             vec![100; 1],
         );
 
-        let bank0 = Bank::new_for_tests(&genesis_config);
-        let bank_forks = BankForks::new_rw_arc(bank0);
+        let (_bank0, bank_forks) = Bank::new_with_bank_forks_for_tests(&genesis_config);
 
         // Fill bank_forks with banks with votes landing in the next slot
         // Create enough banks such that vote account will root slots 0 and 1
         for x in 0..33 {
             let previous_bank = bank_forks.read().unwrap().get(x).unwrap();
-            let bank = Bank::new_from_parent(previous_bank.clone(), &Pubkey::default(), x + 1);
+            let bank = new_bank_from_parent_with_bank_forks(
+                bank_forks.as_ref(),
+                previous_bank.clone(),
+                &Pubkey::default(),
+                x + 1,
+            );
             let vote = vote_transaction::new_vote_transaction(
                 vec![x],
                 previous_bank.hash(),
@@ -526,7 +544,6 @@ mod tests {
                 None,
             );
             bank.process_transaction(&vote).unwrap();
-            bank_forks.write().unwrap().insert(bank);
         }
 
         let working_bank = bank_forks.read().unwrap().working_bank();
@@ -543,7 +560,12 @@ mod tests {
 
         // Add an additional bank/vote that will root slot 2
         let bank33 = bank_forks.read().unwrap().get(33).unwrap();
-        let bank34 = Bank::new_from_parent(bank33.clone(), &Pubkey::default(), 34);
+        let bank34 = new_bank_from_parent_with_bank_forks(
+            bank_forks.as_ref(),
+            bank33.clone(),
+            &Pubkey::default(),
+            34,
+        );
         let vote33 = vote_transaction::new_vote_transaction(
             vec![33],
             bank33.hash(),
@@ -554,7 +576,6 @@ mod tests {
             None,
         );
         bank34.process_transaction(&vote33).unwrap();
-        bank_forks.write().unwrap().insert(bank34);
 
         let working_bank = bank_forks.read().unwrap().working_bank();
         let root = get_vote_account_root_slot(
@@ -587,8 +608,12 @@ mod tests {
         // Add a forked bank. Because the vote for bank 33 landed in the non-ancestor, the vote
         // account's root (and thus the highest_super_majority_root) rolls back to slot 1
         let bank33 = bank_forks.read().unwrap().get(33).unwrap();
-        let bank35 = Bank::new_from_parent(bank33, &Pubkey::default(), 35);
-        bank_forks.write().unwrap().insert(bank35);
+        let _bank35 = new_bank_from_parent_with_bank_forks(
+            bank_forks.as_ref(),
+            bank33,
+            &Pubkey::default(),
+            35,
+        );
 
         let working_bank = bank_forks.read().unwrap().working_bank();
         let ancestors = working_bank.status_cache_ancestors();
@@ -613,7 +638,12 @@ mod tests {
         // continues normally
         for x in 35..=37 {
             let previous_bank = bank_forks.read().unwrap().get(x).unwrap();
-            let bank = Bank::new_from_parent(previous_bank.clone(), &Pubkey::default(), x + 1);
+            let bank = new_bank_from_parent_with_bank_forks(
+                bank_forks.as_ref(),
+                previous_bank.clone(),
+                &Pubkey::default(),
+                x + 1,
+            );
             let vote = vote_transaction::new_vote_transaction(
                 vec![x],
                 previous_bank.hash(),
@@ -624,7 +654,6 @@ mod tests {
                 None,
             );
             bank.process_transaction(&vote).unwrap();
-            bank_forks.write().unwrap().insert(bank);
         }
 
         let working_bank = bank_forks.read().unwrap().working_bank();

core/src/replay_stage.rs

@@ -4160,6 +4160,20 @@ pub(crate) mod tests {
         trees::{tr, Tree},
     };
 
+    fn new_bank_from_parent_with_bank_forks(
+        bank_forks: &RwLock<BankForks>,
+        parent: Arc<Bank>,
+        collector_id: &Pubkey,
+        slot: Slot,
+    ) -> Arc<Bank> {
+        let bank = Bank::new_from_parent(parent, collector_id, slot);
+        bank_forks
+            .write()
+            .unwrap()
+            .insert(bank)
+            .clone_without_scheduler()
+    }
+
     #[test]
     fn test_is_partition_detected() {
         let (VoteSimulator { bank_forks, .. }, _) = setup_default_forks(1, None::<GenerateVotes>);
@@ -4933,7 +4947,12 @@ pub(crate) mod tests {
         for i in 1..=3 {
             let prev_bank = bank_forks.read().unwrap().get(i - 1).unwrap();
             let slot = prev_bank.slot() + 1;
-            let bank = Bank::new_from_parent(prev_bank, &Pubkey::default(), slot);
+            let bank = new_bank_from_parent_with_bank_forks(
+                bank_forks.as_ref(),
+                prev_bank,
+                &Pubkey::default(),
+                slot,
+            );
             let _res = bank.transfer(
                 10,
                 &genesis_config_info.mint_keypair,
@@ -4942,7 +4961,7 @@ pub(crate) mod tests {
             for _ in 0..genesis_config.ticks_per_slot {
                 bank.register_default_tick_for_test();
             }
-            bank_forks.write().unwrap().insert(bank);
+
             let arc_bank = bank_forks.read().unwrap().get(i).unwrap();
             leader_vote(i - 1, &arc_bank, &leader_voting_pubkey);
             ReplayStage::update_commitment_cache(
@@ -5094,10 +5113,6 @@ pub(crate) mod tests {
             None,
         );
 
-        let bank1 = Bank::new_from_parent(bank0.clone(), &my_node_pubkey, 1);
-        bank1.process_transaction(&vote_tx).unwrap();
-        bank1.freeze();
-
         // Test confirmations
         let ancestors = bank_forks.read().unwrap().ancestors();
         let mut frozen_banks: Vec<_> = bank_forks
@@ -5138,8 +5153,16 @@ pub(crate) mod tests {
             assert!(confirmed_forks.is_empty());
         }
 
+        let bank1 = new_bank_from_parent_with_bank_forks(
+            bank_forks.as_ref(),
+            bank0.clone(),
+            &my_node_pubkey,
+            1,
+        );
+        bank1.process_transaction(&vote_tx).unwrap();
+        bank1.freeze();
+
         // Insert the bank that contains a vote for slot 0, which confirms slot 0
-        bank_forks.write().unwrap().insert(bank1);
         progress.insert(
             1,
             ForkProgress::new(bank0.last_blockhash(), None, None, 0, 0),
@@ -7325,7 +7348,12 @@ pub(crate) mod tests {
         let (voting_sender, voting_receiver) = unbounded();
 
         // Simulate landing a vote for slot 0 landing in slot 1
-        let bank1 = Arc::new(Bank::new_from_parent(bank0.clone(), &Pubkey::default(), 1));
+        let bank1 = new_bank_from_parent_with_bank_forks(
+            bank_forks.as_ref(),
+            bank0.clone(),
+            &Pubkey::default(),
+            1,
+        );
         bank1.fill_bank_with_ticks_for_tests();
         tower.record_bank_vote(&bank0);
         ReplayStage::push_vote(
@@ -7366,7 +7394,12 @@ pub(crate) mod tests {
 
         // Trying to refresh the vote for bank 0 in bank 1 or bank 2 won't succeed because
         // the last vote has landed already
-        let bank2 = Arc::new(Bank::new_from_parent(bank1.clone(), &Pubkey::default(), 2));
+        let bank2 = new_bank_from_parent_with_bank_forks(
+            bank_forks.as_ref(),
+            bank1.clone(),
+            &Pubkey::default(),
+            2,
+        );
         bank2.fill_bank_with_ticks_for_tests();
         bank2.freeze();
         for refresh_bank in &[&bank1, &bank2] {
@@ -7460,8 +7493,12 @@ pub(crate) mod tests {
             let mut parent_bank = bank2.clone();
             for _ in 0..MAX_PROCESSING_AGE {
                 let slot = parent_bank.slot() + 1;
-                parent_bank =
-                    Arc::new(Bank::new_from_parent(parent_bank, &Pubkey::default(), slot));
+                parent_bank = new_bank_from_parent_with_bank_forks(
+                    bank_forks.as_ref(),
+                    parent_bank,
+                    &Pubkey::default(),
+                    slot,
+                );
                 parent_bank.fill_bank_with_ticks_for_tests();
                 parent_bank.freeze();
             }
@@ -7516,11 +7553,12 @@ pub(crate) mod tests {
 
         // Processing the vote transaction should be valid
         let expired_bank_child_slot = expired_bank.slot() + 1;
-        let expired_bank_child = Arc::new(Bank::new_from_parent(
+        let expired_bank_child = new_bank_from_parent_with_bank_forks(
+            bank_forks.as_ref(),
             expired_bank.clone(),
             &Pubkey::default(),
             expired_bank_child_slot,
-        ));
+        );
         expired_bank_child.process_transaction(vote_tx).unwrap();
         let vote_account = expired_bank_child
             .get_vote_account(&my_vote_pubkey)
@@ -7536,11 +7574,12 @@ pub(crate) mod tests {
         // 1) The vote for slot 1 hasn't landed
         // 2) The latest refresh vote transaction's recent blockhash (the sibling's hash) doesn't exist
         // This will still not refresh because `MAX_VOTE_REFRESH_INTERVAL_MILLIS` has not expired yet
-        let expired_bank_sibling = Arc::new(Bank::new_from_parent(
+        let expired_bank_sibling = new_bank_from_parent_with_bank_forks(
+            bank_forks.as_ref(),
             bank2,
             &Pubkey::default(),
             expired_bank_child.slot() + 1,
-        ));
+        );
         expired_bank_sibling.fill_bank_with_ticks_for_tests();
         expired_bank_sibling.freeze();
         // Set the last refresh to now, shouldn't refresh because the last refresh just happened.
@@ -7628,7 +7667,12 @@ pub(crate) mod tests {
             parent_bank.last_blockhash()
         );
         assert_eq!(tower.last_voted_slot().unwrap(), parent_bank.slot());
-        let bank = Bank::new_from_parent(parent_bank, &Pubkey::default(), my_slot);
+        let bank = new_bank_from_parent_with_bank_forks(
+            bank_forks,
+            parent_bank,
+            &Pubkey::default(),
+            my_slot,
+        );
         bank.fill_bank_with_ticks_for_tests();
         if make_it_landing {
             bank.process_transaction(vote_tx).unwrap();
@@ -7644,7 +7688,6 @@ pub(crate) mod tests {
                 0,
             )
         });
-        bank_forks.write().unwrap().insert(bank);
         bank_forks.read().unwrap().get(my_slot).unwrap()
     }
 
@@ -7685,8 +7728,12 @@ pub(crate) mod tests {
         // Add a new fork starting from 0 with bigger slot number, we assume it has a bigger
         // weight, but we cannot switch because of lockout.
         let other_fork_slot = 1;
-        let other_fork_bank =
-            Bank::new_from_parent(bank0.clone(), &Pubkey::default(), other_fork_slot);
+        let other_fork_bank = new_bank_from_parent_with_bank_forks(
+            bank_forks.as_ref(),
+            bank0.clone(),
+            &Pubkey::default(),
+            other_fork_slot,
+        );
         other_fork_bank.fill_bank_with_ticks_for_tests();
         other_fork_bank.freeze();
         progress.entry(other_fork_slot).or_insert_with(|| {
@@ -7699,7 +7746,6 @@ pub(crate) mod tests {
                 0,
             )
         });
-        bank_forks.write().unwrap().insert(other_fork_bank);
 
         let (voting_sender, voting_receiver) = unbounded();
         let mut cursor = Cursor::default();
@@ -7745,7 +7791,12 @@ pub(crate) mod tests {
         let last_voted_slot = tower.last_voted_slot().unwrap();
         while new_bank.is_in_slot_hashes_history(&last_voted_slot) {
             let new_slot = new_bank.slot() + 1;
-            let bank = Bank::new_from_parent(new_bank, &Pubkey::default(), new_slot);
+            let bank = new_bank_from_parent_with_bank_forks(
+                bank_forks.as_ref(),
+                new_bank,
+                &Pubkey::default(),
+                new_slot,
+            );
             bank.fill_bank_with_ticks_for_tests();
             bank.freeze();
             progress.entry(new_slot).or_insert_with(|| {
@@ -7758,7 +7809,6 @@ pub(crate) mod tests {
                     0,
                 )
             });
-            bank_forks.write().unwrap().insert(bank);
             new_bank = bank_forks.read().unwrap().get(new_slot).unwrap();
         }
         let tip_of_voted_fork = new_bank.slot();