diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs index c2af04c249..a5476de4e0 100644 --- a/runtime/src/bank.rs +++ b/runtime/src/bank.rs @@ -3184,6 +3184,20 @@ impl Bank { .for_each(|x| rewards.push((x.stake_pubkey, x.stake_reward_info))); } + #[allow(dead_code)] + /// insert non-zero stake rewards to self.rewards + /// Return the number of rewards inserted + fn update_reward_history_in_partition(&self, stake_rewards: &[StakeReward]) -> usize { + let mut rewards = self.rewards.write().unwrap(); + rewards.reserve(stake_rewards.len()); + let initial_len = rewards.len(); + stake_rewards + .iter() + .filter(|x| x.get_stake_reward() > 0) + .for_each(|x| rewards.push((x.stake_pubkey, x.stake_reward_info))); + rewards.len().saturating_sub(initial_len) + } + fn update_recent_blockhashes_locked(&self, locked_blockhash_queue: &BlockhashQueue) { #[allow(deprecated)] self.update_sysvar_account(&sysvar::recent_blockhashes::id(), |account| { diff --git a/runtime/src/bank/tests.rs b/runtime/src/bank/tests.rs index 8ad6418735..8946b67fc7 100644 --- a/runtime/src/bank/tests.rs +++ b/runtime/src/bank/tests.rs @@ -12606,6 +12606,63 @@ fn test_store_stake_accounts_in_partition_empty() { assert_eq!(expected_total, total_rewards_in_lamports); } +#[test] +fn test_update_reward_history_in_partition() { + for zero_reward in [false, true] { + let (genesis_config, _mint_keypair) = create_genesis_config(1_000_000 * LAMPORTS_PER_SOL); + let bank = Bank::new_for_tests(&genesis_config); + + let mut expected_num = 100; + + let mut stake_rewards = (0..expected_num) + .map(|_| StakeReward::new_random()) + .collect::>(); + + let mut rng = rand::thread_rng(); + let i_zero = rng.gen_range(0, expected_num); + if zero_reward { + // pick one entry to have zero rewards so it gets ignored + stake_rewards[i_zero].stake_reward_info.lamports = 0; + } + + let num_in_history = bank.update_reward_history_in_partition(&stake_rewards); + + if zero_reward { + stake_rewards.remove(i_zero); + // -1 because one of them had zero rewards and was ignored + expected_num -= 1; + } + + bank.rewards + .read() + .unwrap() + .iter() + .zip(stake_rewards.iter()) + .for_each(|((k, reward_info), expected_stake_reward)| { + assert_eq!( + ( + &expected_stake_reward.stake_pubkey, + &expected_stake_reward.stake_reward_info + ), + (k, reward_info) + ); + }); + + assert_eq!(num_in_history, expected_num); + } +} + +#[test] +fn test_update_reward_history_in_partition_empty() { + let (genesis_config, _mint_keypair) = create_genesis_config(1_000_000 * LAMPORTS_PER_SOL); + let bank = Bank::new_for_tests(&genesis_config); + + let stake_rewards = vec![]; + + let num_in_history = bank.update_reward_history_in_partition(&stake_rewards); + assert_eq!(num_in_history, 0); +} + #[test] fn test_system_instruction_allocate() { let (genesis_config, mint_keypair) = create_genesis_config(sol_to_lamports(1.0));