add calc_vote_accounts_to_store (#32063)

* add calc_vote_accounts_to_store

* update comment

runtime/src/bank.rs

@@ -1104,6 +1104,17 @@ struct VoteReward {
 }
 
 type VoteRewards = DashMap<Pubkey, VoteReward>;
+#[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<Option<AccountSharedData>>,
+}
 pub(crate) type StakeRewards = Vec<StakeReward>;
 
 #[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<Pubkey, VoteReward>,
+    ) -> 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,

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,
+            accounts_equal, create_account_shared_data_with_fields as create_account, from_account,
-            AccountSharedData, ReadableAccount, WritableAccount,
+            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);
+            }
+        }
+    }
+}