diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs index a5476de4e0..404af2ba65 100644 --- a/runtime/src/bank.rs +++ b/runtime/src/bank.rs @@ -1104,6 +1104,17 @@ struct VoteReward { } type VoteRewards = DashMap; +#[allow(dead_code)] +#[derive(Debug)] +struct VoteRewardsAccounts { + /// reward info for each vote account pubkey. + /// This type is used by `update_reward_history()` + rewards: Vec<(Pubkey, RewardInfo)>, + /// corresponds to pubkey in `rewards` + /// Some if account is to be stored. + /// None if to be skipped. + accounts_to_store: Vec>, +} pub(crate) type StakeRewards = Vec; #[derive(Debug, Default)] @@ -3169,6 +3180,55 @@ impl Bank { vote_rewards } + #[allow(dead_code)] + /// return reward info for each vote account + /// return account data for each vote account that needs to be stored + /// This return value is a little awkward at the moment so that downstream existing code in the non-partitioned rewards code path can be re-used without duplication or modification. + /// This function is copied from the existing code path's `store_vote_accounts`. + /// The primary differences: + /// - we want this fn to have no side effects (such as actually storing vote accounts) so that we + /// can compare the expected results with the current code path + /// - we want to be able to batch store the vote accounts later for improved performance/cache updating + fn calc_vote_accounts_to_store( + vote_account_rewards: DashMap, + ) -> VoteRewardsAccounts { + let len = vote_account_rewards.len(); + let mut result = VoteRewardsAccounts { + rewards: Vec::with_capacity(len), + accounts_to_store: Vec::with_capacity(len), + }; + vote_account_rewards.into_iter().for_each( + |( + vote_pubkey, + VoteReward { + mut vote_account, + commission, + vote_rewards, + vote_needs_store, + }, + )| { + if let Err(err) = vote_account.checked_add_lamports(vote_rewards) { + debug!("reward redemption failed for {}: {:?}", vote_pubkey, err); + return; + } + + result.rewards.push(( + vote_pubkey, + RewardInfo { + reward_type: RewardType::Voting, + lamports: vote_rewards as i64, + post_balance: vote_account.lamports(), + commission: Some(commission), + }, + )); + result + .accounts_to_store + .push(vote_needs_store.then_some(vote_account)); + }, + ); + result + } + fn update_reward_history( &self, stake_rewards: StakeRewards, diff --git a/runtime/src/bank/tests.rs b/runtime/src/bank/tests.rs index 8946b67fc7..3abefa9c7c 100644 --- a/runtime/src/bank/tests.rs +++ b/runtime/src/bank/tests.rs @@ -42,8 +42,8 @@ use { }, solana_sdk::{ account::{ - create_account_shared_data_with_fields as create_account, from_account, Account, - AccountSharedData, ReadableAccount, WritableAccount, + accounts_equal, create_account_shared_data_with_fields as create_account, from_account, + Account, AccountSharedData, ReadableAccount, WritableAccount, }, account_utils::StateMut, bpf_loader, @@ -12860,3 +12860,126 @@ fn test_system_instruction_unsigned_transaction() { ); assert_eq!(bank_client.get_balance(&mallory_pubkey).unwrap(), amount); } + +#[test] +fn test_calc_vote_accounts_to_store_empty() { + let vote_account_rewards = DashMap::default(); + let result = Bank::calc_vote_accounts_to_store(vote_account_rewards); + assert_eq!(result.rewards.len(), result.accounts_to_store.len()); + assert!(result.rewards.is_empty()); +} + +#[test] +fn test_calc_vote_accounts_to_store_overflow() { + let vote_account_rewards = DashMap::default(); + let pubkey = solana_sdk::pubkey::new_rand(); + let mut vote_account = AccountSharedData::default(); + vote_account.set_lamports(u64::MAX); + vote_account_rewards.insert( + pubkey, + VoteReward { + vote_account, + commission: 0, + vote_rewards: 1, // enough to overflow + vote_needs_store: false, + }, + ); + let result = Bank::calc_vote_accounts_to_store(vote_account_rewards); + assert_eq!(result.rewards.len(), result.accounts_to_store.len()); + assert!(result.rewards.is_empty()); +} + +#[test] +fn test_calc_vote_accounts_to_store_three() { + let vote_account_rewards = DashMap::default(); + let pubkey = solana_sdk::pubkey::new_rand(); + let pubkey2 = solana_sdk::pubkey::new_rand(); + let pubkey3 = solana_sdk::pubkey::new_rand(); + let mut vote_account = AccountSharedData::default(); + vote_account.set_lamports(u64::MAX); + vote_account_rewards.insert( + pubkey, + VoteReward { + vote_account: vote_account.clone(), + commission: 0, + vote_rewards: 0, + vote_needs_store: false, // don't store + }, + ); + vote_account_rewards.insert( + pubkey2, + VoteReward { + vote_account: vote_account.clone(), + commission: 0, + vote_rewards: 0, + vote_needs_store: true, // this one needs storing + }, + ); + vote_account_rewards.insert( + pubkey3, + VoteReward { + vote_account: vote_account.clone(), + commission: 0, + vote_rewards: 0, + vote_needs_store: false, // don't store + }, + ); + let result = Bank::calc_vote_accounts_to_store(vote_account_rewards); + assert_eq!(result.rewards.len(), result.accounts_to_store.len()); + assert_eq!(result.rewards.len(), 3); + result.rewards.iter().enumerate().for_each(|(i, (k, _))| { + // pubkey2 is some(account), others should be none + if k == &pubkey2 { + assert!(accounts_equal( + result.accounts_to_store[i].as_ref().unwrap(), + &vote_account + )); + } else { + assert!(result.accounts_to_store[i].is_none()); + } + }); +} + +#[test] +fn test_calc_vote_accounts_to_store_normal() { + let pubkey = solana_sdk::pubkey::new_rand(); + for commission in 0..2 { + for vote_rewards in 0..2 { + for vote_needs_store in [false, true] { + let vote_account_rewards = DashMap::default(); + let mut vote_account = AccountSharedData::default(); + vote_account.set_lamports(1); + vote_account_rewards.insert( + pubkey, + VoteReward { + vote_account: vote_account.clone(), + commission, + vote_rewards, + vote_needs_store, + }, + ); + let result = Bank::calc_vote_accounts_to_store(vote_account_rewards); + assert_eq!(result.rewards.len(), result.accounts_to_store.len()); + assert_eq!(result.rewards.len(), 1); + let rewards = &result.rewards[0]; + let account = &result.accounts_to_store[0]; + _ = vote_account.checked_add_lamports(vote_rewards); + if vote_needs_store { + assert!(accounts_equal(account.as_ref().unwrap(), &vote_account)); + } else { + assert!(account.is_none()); + } + assert_eq!( + rewards.1, + RewardInfo { + reward_type: RewardType::Voting, + lamports: vote_rewards as i64, + post_balance: vote_account.lamports(), + commission: Some(commission), + } + ); + assert_eq!(rewards.0, pubkey); + } + } + } +}