runtime: Add test for epoch boundary (anza-xyz#8838)

Add test that ensures correct values in `Bank` and
`PartitionedRewardsCalculation` after crossing epoch boundary.

(cherry picked from commit 30d8afd)

# Conflicts:
#	runtime/src/bank/partitioned_epoch_rewards/calculation.rs

Author: vadorovsky

runtime/src/bank/partitioned_epoch_rewards/calculation.rs

@@ -672,8 +672,8 @@ mod tests {
                 partitioned_epoch_rewards::{
                     tests::{
                         build_partitioned_stake_rewards, create_default_reward_bank,
-                        create_reward_bank, create_reward_bank_with_specific_stakes, RewardBank,
-                        SLOTS_PER_EPOCH,
+                        create_reward_bank, create_reward_bank_with_specific_stakes,
+                        populate_vote_accounts_with_votes, RewardBank, SLOTS_PER_EPOCH,
                     },
                     EpochRewardPhase, EpochRewardStatus, PartitionedStakeRewards,
                     StartBlockHeightAndPartitionedRewards,
@@ -682,7 +682,12 @@ mod tests {
                 RewardInfo, VoteReward,
             },
             stake_account::StakeAccount,
+<<<<<<< HEAD
             stakes::Stakes,
+=======
+            stake_utils,
+            stakes::{tests::create_staked_node_accounts, Stakes},
+>>>>>>> 30d8afd32 (runtime: Add test for epoch boundary (#8838))
         },
         agave_feature_set::FeatureSet,
         rayon::ThreadPoolBuilder,
@@ -693,7 +698,10 @@ mod tests {
         solana_stake_interface::state::{Delegation, StakeStateV2},
         solana_vote_interface::state::VoteStateV4,
         solana_vote_program::vote_state,
-        std::sync::{Arc, RwLockReadGuard},
+        std::{
+            collections::HashSet,
+            sync::{Arc, RwLockReadGuard},
+        },
     };
 
     #[test]
@@ -1414,4 +1422,229 @@ mod tests {
         assert_eq!(vote_reward_c.commission, 10);
         assert_eq!(vote_reward_c.vote_rewards, 50);
     }
+<<<<<<< HEAD
+=======
+
+    #[test]
+    fn test_epoch_rewards_cache_multiple_forks() {
+        let (mut genesis_config, _mint_keypair) =
+            create_genesis_config(1_000_000 * LAMPORTS_PER_SOL);
+
+        const NUM_STAKES: usize = 1000;
+
+        for _i in 0..NUM_STAKES {
+            let vote_pubkey = Pubkey::new_unique();
+            let stake_pubkey = Pubkey::new_unique();
+
+            genesis_config.accounts.insert(
+                vote_pubkey,
+                vote_state::create_v4_account_with_authorized(
+                    &vote_pubkey,
+                    &Pubkey::new_unique(),
+                    &Pubkey::new_unique(),
+                    None,
+                    0,
+                    100_000_000_000,
+                )
+                .into(),
+            );
+
+            let stake_lamports = 1_000_000_000_000;
+            let stake_account = stake_utils::create_stake_account(
+                &stake_pubkey,
+                &vote_pubkey,
+                &vote_state::create_v4_account_with_authorized(
+                    &vote_pubkey,
+                    &Pubkey::new_unique(),
+                    &Pubkey::new_unique(),
+                    None,
+                    0,
+                    100_000_000_000,
+                ),
+                &genesis_config.rent,
+                stake_lamports,
+            );
+            genesis_config
+                .accounts
+                .insert(stake_pubkey, stake_account.into());
+        }
+
+        let bank = Arc::new(Bank::new_for_tests(&genesis_config));
+        let slots_per_epoch = bank.epoch_schedule().slots_per_epoch;
+        {
+            let cache = bank.epoch_rewards_calculation_cache.lock().unwrap();
+            assert!(
+                !cache.contains_key(&bank.parent_hash()),
+                "cache should be empty"
+            );
+        }
+
+        let bank_fork1 =
+            Bank::new_from_parent(Arc::clone(&bank), &Pubkey::default(), slots_per_epoch);
+        {
+            let cache = bank_fork1.epoch_rewards_calculation_cache.lock().unwrap();
+            assert!(
+                cache.contains_key(&bank_fork1.parent_hash()),
+                "cache should be populated"
+            );
+        }
+
+        let bank_fork2 = Bank::new_from_parent(bank, &Pubkey::default(), slots_per_epoch);
+        {
+            let cache = bank_fork2.epoch_rewards_calculation_cache.lock().unwrap();
+            assert!(
+                cache.contains_key(&bank_fork2.parent_hash()),
+                "cache should be populated"
+            );
+        }
+    }
+
+    fn add_voters_and_populate(
+        bank: &Arc<Bank>,
+        voters: &mut HashSet<Pubkey>,
+        stakers: &mut HashSet<Pubkey>,
+        count: usize,
+        stake_lamports: u64,
+        commission: u8,
+    ) {
+        for _ in 0..count {
+            let ((vote_pubkey, vote_account), (stake_pubkey, stake_account)) =
+                create_staked_node_accounts(stake_lamports);
+            bank.store_account_and_update_capitalization(&vote_pubkey, &vote_account);
+            bank.store_account_and_update_capitalization(&stake_pubkey, &stake_account);
+            voters.insert(vote_pubkey);
+            stakers.insert(stake_pubkey);
+        }
+        populate_vote_accounts_with_votes(bank, voters.iter().copied(), commission);
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    fn assert_cached_rewards(
+        bank: &Arc<Bank>,
+        expected_cache_len: usize,
+        expected_voters: &HashSet<Pubkey>,
+        expected_stakers: &HashSet<Pubkey>,
+        expected_vote_rewards: u64,
+        expected_stake_rewards: u64,
+        expected_rewards: u64,
+        expected_points: u128,
+        parent_capitalization: Option<u64>,
+    ) {
+        let cache = bank.epoch_rewards_calculation_cache.lock().unwrap();
+        assert_eq!(cache.len(), expected_cache_len);
+        let partitioned = cache.get(&bank.parent_hash()).unwrap().as_ref();
+        let VoteRewardsAccounts {
+            accounts_with_rewards,
+            total_vote_rewards_lamports,
+            ..
+        } = &partitioned.vote_account_rewards;
+        let StakeRewardCalculation {
+            stake_rewards,
+            total_stake_rewards_lamports,
+            ..
+        } = &partitioned.stake_rewards;
+        let point_value = &partitioned.point_value;
+        let voters: HashSet<_> = accounts_with_rewards
+            .iter()
+            .map(|(pubkey, _reward, _acc)| *pubkey)
+            .collect();
+        let stakers: HashSet<_> = stake_rewards
+            .rewards
+            .iter()
+            .filter_map(|reward| reward.as_ref())
+            .map(|reward| reward.stake_pubkey)
+            .collect();
+        assert_eq!(expected_voters, &voters);
+        assert_eq!(expected_stakers, &stakers);
+        assert_eq!(*total_vote_rewards_lamports, expected_vote_rewards);
+        assert_eq!(*total_stake_rewards_lamports, expected_stake_rewards);
+        assert_eq!(point_value.rewards, expected_rewards);
+        assert_eq!(point_value.points, expected_points);
+        if let Some(parent_cap) = parent_capitalization {
+            assert_eq!(bank.capitalization(), parent_cap + expected_vote_rewards);
+        }
+    }
+
+    #[test]
+    fn test_epoch_boundary() {
+        let delegations = 100;
+        let stake_lamports = 2_000_000_000;
+        let stakes: Vec<_> = (0..delegations).map(|_| stake_lamports).collect();
+        let (
+            RewardBank {
+                bank: bank1,
+                voters,
+                stakers,
+                ..
+            },
+            _bank_forks,
+        ) = create_reward_bank_with_specific_stakes(
+            stakes,
+            PartitionedEpochRewardsConfig::default().stake_account_stores_per_block,
+            SLOTS_PER_EPOCH,
+        );
+        let mut voters: HashSet<_> = voters.into_iter().collect();
+        let mut stakers: HashSet<_> = stakers.into_iter().collect();
+
+        // The sysvar account holds the rent-exempt lamport added after
+        // reward calculation, so the bank capitalization exceeds the cached
+        // value by this amount.
+        let epoch_rewards_sysvar_balance = bank1.get_balance(&solana_sysvar::epoch_rewards::id());
+        assert_eq!(epoch_rewards_sysvar_balance, 1);
+
+        assert_cached_rewards(
+            &bank1,
+            1,                     // expected_cache_len
+            &voters,               // expected_voters
+            &stakers,              // expected_stakers
+            0,                     // expected_vote_rewards
+            12300,                 // expected_stake_rewards
+            12392,                 // expected_rewards
+            8_400_000_000_000u128, // expected_points
+            None,                  // parent_capitalization
+        );
+
+        add_voters_and_populate(&bank1, &mut voters, &mut stakers, 5, 5_000_000_000, 10);
+        let parent_capitalization = bank1.capitalization();
+
+        let bank2 = Arc::new(Bank::new_from_parent(
+            Arc::clone(&bank1),
+            &Pubkey::default(),
+            SLOTS_PER_EPOCH * 2,
+        ));
+
+        assert_cached_rewards(
+            &bank2,
+            2,                           // expected_cache_len
+            &voters,                     // expected_voters
+            &stakers,                    // expected_stakers
+            1245,                        // expected_vote_rewards
+            11810,                       // expected_stake_rewards
+            13163,                       // expected_rewards
+            9_450_000_000_000u128,       // expected_points
+            Some(parent_capitalization), // parent_capitalization
+        );
+
+        add_voters_and_populate(&bank2, &mut voters, &mut stakers, 10, 8_000_000_000, 10);
+        let parent_capitalization = bank2.capitalization();
+
+        let bank3 = Arc::new(Bank::new_from_parent(
+            Arc::clone(&bank2),
+            &Pubkey::default(),
+            SLOTS_PER_EPOCH * 3,
+        ));
+
+        assert_cached_rewards(
+            &bank3,
+            3,                           // expected_cache_len
+            &voters,                     // expected_voters
+            &stakers,                    // expected_stakers
+            1525,                        // expected_vote_rewards
+            13930,                       // expected_stake_rewards
+            15629,                       // expected_rewards
+            12_810_000_000_000u128,      // expected_points
+            Some(parent_capitalization), // parent_capitalization
+        );
+    }
+>>>>>>> 30d8afd32 (runtime: Add test for epoch boundary (#8838))
 }

runtime/src/bank/partitioned_epoch_rewards/mod.rs

@@ -391,7 +391,7 @@ mod tests {
         solana_system_transaction as system_transaction,
         solana_vote::vote_transaction,
         solana_vote_interface::state::{VoteStateV4, VoteStateVersions, MAX_LOCKOUT_HISTORY},
-        solana_vote_program::vote_state::{self, TowerSync},
+        solana_vote_program::vote_state::{self, handler::VoteStateHandle, TowerSync},
         std::sync::{Arc, RwLock},
     };
 
@@ -555,27 +555,11 @@ mod tests {
 
         // Fill bank_forks with banks with votes landing in the next slot
         // Create enough banks such that vote account will root
-        for validator_vote_keypairs in &validator_keypairs {
-            let vote_id = validator_vote_keypairs.vote_keypair.pubkey();
-            let mut vote_account = bank.get_account(&vote_id).unwrap();
-            // generate some rewards
-            let mut vote_state =
-                Some(VoteStateV4::deserialize(vote_account.data(), &vote_id).unwrap());
-            for i in 0..MAX_LOCKOUT_HISTORY + 42 {
-                if let Some(v) = vote_state.as_mut() {
-                    vote_state::process_slot_vote_unchecked(v, i as u64)
-                }
-                let versioned = VoteStateVersions::V4(Box::new(vote_state.take().unwrap()));
-                vote_account.set_state(&versioned).unwrap();
-                match versioned {
-                    VoteStateVersions::V4(v) => {
-                        vote_state = Some(*v);
-                    }
-                    _ => panic!("Has to be of type V4"),
-                };
-            }
-            bank.store_account_and_update_capitalization(&vote_id, &vote_account);
-        }
+        populate_vote_accounts_with_votes(
+            &bank,
+            validator_keypairs.iter().map(|k| k.vote_keypair.pubkey()),
+            0,
+        );
 
         // Advance some num slots; usually to the next epoch boundary to update
         // EpochStakes
@@ -603,6 +587,37 @@ mod tests {
         )
     }
 
+    pub(super) fn populate_vote_accounts_with_votes(
+        bank: &Bank,
+        vote_pubkeys: impl IntoIterator<Item = Pubkey>,
+        commission: u8,
+    ) {
+        for vote_pubkey in vote_pubkeys {
+            let mut vote_account = bank
+                .get_account(&vote_pubkey)
+                .unwrap_or_else(|| panic!("missing vote account {vote_pubkey:?}"));
+            let mut vote_state =
+                Some(VoteStateV4::deserialize(vote_account.data(), &vote_pubkey).unwrap());
+            if let Some(state) = vote_state.as_mut() {
+                state.set_commission(commission);
+            }
+            for i in 0..MAX_LOCKOUT_HISTORY + 42 {
+                if let Some(state) = vote_state.as_mut() {
+                    vote_state::process_slot_vote_unchecked(state, i as u64);
+                }
+                let versioned = VoteStateVersions::V4(Box::new(vote_state.take().unwrap()));
+                vote_account.set_state(&versioned).unwrap();
+                match versioned {
+                    VoteStateVersions::V4(v) => {
+                        vote_state = Some(*v);
+                    }
+                    _ => panic!("Has to be of type V4"),
+                };
+            }
+            bank.store_account_and_update_capitalization(&vote_pubkey, &vote_account);
+        }
+    }
+
     #[test]
     fn test_force_reward_interval_end() {
         let (genesis_config, _mint_keypair) = create_genesis_config(1_000_000 * LAMPORTS_PER_SOL);