Optimize VoteRewardsAccounts (#7236)

HaoranYi · web-flow · commit 82c1514b629c · 2025-07-30T17:08:34.000-05:00
* refactor VoteRewardsAccounts

* impl&lt;'a&gt; StorableAccounts&lt;'a&gt; for VoteRewardsAccountsStorable&lt;'a&gt;
diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
@@ -2354,8 +2354,7 @@ impl Bank {
     fn calc_vote_accounts_to_store(vote_account_rewards: VoteRewards) -> VoteRewardsAccounts {
         let len = vote_account_rewards.len();
         let mut result = VoteRewardsAccounts {
-            rewards: Vec::with_capacity(len),
-            accounts_to_store: Vec::with_capacity(len),
+            accounts_with_rewards: Vec::with_capacity(len),
             total_vote_rewards_lamports: 0,
         };
         vote_account_rewards.into_iter().for_each(
@@ -2372,16 +2371,16 @@ impl Bank {
                     return;
                 }
 
-                result.rewards.push((
+                result.accounts_with_rewards.push((
                     vote_pubkey,
                     RewardInfo {
                         reward_type: RewardType::Voting,
                         lamports: vote_rewards as i64,
                         post_balance: vote_account.lamports(),
                         commission: Some(commission),
                     },
+                    vote_account,
                 ));
-                result.accounts_to_store.push((vote_pubkey, vote_account));
                 result.total_vote_rewards_lamports += vote_rewards;
             },
         );
@@ -2391,14 +2390,17 @@ impl Bank {
     fn update_reward_history(
         &self,
         stake_rewards: StakeRewards,
-        vote_rewards: &[(Pubkey, RewardInfo)],
+        vote_rewards: &VoteRewardsAccounts,
     ) {
-        let additional_reserve = stake_rewards.len() + vote_rewards.len();
+        let additional_reserve = stake_rewards.len() + vote_rewards.accounts_with_rewards.len();
         let mut rewards = self.rewards.write().unwrap();
         rewards.reserve(additional_reserve);
-        vote_rewards.iter().for_each(|(vote_pubkey, vote_reward)| {
-            rewards.push((*vote_pubkey, *vote_reward));
-        });
+        vote_rewards
+            .accounts_with_rewards
+            .iter()
+            .for_each(|(vote_pubkey, vote_reward, _)| {
+                rewards.push((*vote_pubkey, *vote_reward));
+            });
         stake_rewards
             .into_iter()
             .filter(|x| x.get_stake_reward() > 0)
diff --git a/runtime/src/bank/partitioned_epoch_rewards/calculation.rs b/runtime/src/bank/partitioned_epoch_rewards/calculation.rs
@@ -3,7 +3,8 @@ use {
         epoch_rewards_hasher::hash_rewards_into_partitions, Bank,
         CalculateRewardsAndDistributeVoteRewardsResult, CalculateValidatorRewardsResult,
         EpochRewardCalculateParamInfo, PartitionedRewardsCalculation, PartitionedStakeReward,
-        StakeRewardCalculation, VoteRewardsAccounts, REWARD_CALCULATION_NUM_BLOCKS,
+        StakeRewardCalculation, VoteRewardsAccounts, VoteRewardsAccountsStorable,
+        REWARD_CALCULATION_NUM_BLOCKS,
     },
     crate::{
         bank::{
@@ -147,7 +148,7 @@ impl Bank {
         self.store_vote_accounts_partitioned(vote_account_rewards, metrics);
 
         // update reward history of JUST vote_rewards, stake_rewards is vec![] here
-        self.update_reward_history(vec![], &vote_account_rewards.rewards[..]);
+        self.update_reward_history(vec![], vote_account_rewards);
 
         let StakeRewardCalculation {
             stake_rewards,
@@ -210,7 +211,11 @@ impl Bank {
         metrics: &RewardsMetrics,
     ) {
         let (_, measure_us) = measure_us!({
-            self.store_accounts((self.slot(), &vote_account_rewards.accounts_to_store[..]));
+            let storable = VoteRewardsAccountsStorable {
+                slot: self.slot(),
+                vote_rewards_accounts: vote_account_rewards,
+            };
+            self.store_accounts(storable);
         });
 
         metrics
@@ -606,10 +611,11 @@ mod tests {
                     post_balance: vote_reward_info.vote_rewards,
                     commission: Some(vote_reward_info.commission),
                 };
-                vote_rewards_account.rewards.push((*vote_key, info));
-                vote_rewards_account
-                    .accounts_to_store
-                    .push((*vote_key, vote_reward_info.vote_account.clone()));
+                vote_rewards_account.accounts_with_rewards.push((
+                    *vote_key,
+                    info,
+                    vote_reward_info.vote_account.clone(),
+                ));
                 vote_rewards_account.total_vote_rewards_lamports += vote_reward_info.vote_rewards;
             });
 
@@ -619,7 +625,7 @@ mod tests {
         bank.store_vote_accounts_partitioned(&vote_rewards_account, &metrics);
         assert_eq!(
             expected_vote_rewards_num,
-            vote_rewards_account.accounts_to_store.len()
+            vote_rewards_account.accounts_with_rewards.len()
         );
         assert_eq!(
             vote_rewards
@@ -654,7 +660,7 @@ mod tests {
         let total_vote_rewards = vote_rewards.total_vote_rewards_lamports;
 
         bank.store_vote_accounts_partitioned(&vote_rewards, &metrics);
-        assert_eq!(expected, vote_rewards.accounts_to_store.len());
+        assert_eq!(expected, vote_rewards.accounts_with_rewards.len());
         assert_eq!(0, total_vote_rewards);
     }
 
@@ -688,9 +694,9 @@ mod tests {
             .stake_reward_calculation;
 
         let total_vote_rewards: u64 = vote_rewards
-            .rewards
+            .accounts_with_rewards
             .iter()
-            .map(|reward| reward.1.lamports)
+            .map(|(_, reward_info, _)| reward_info.lamports)
             .sum::<i64>() as u64;
 
         // assert that total rewards matches the sum of vote rewards and stake rewards
@@ -800,26 +806,26 @@ mod tests {
         let vote_state = VoteState::deserialize(vote_account.data()).unwrap();
 
         assert_eq!(
-            vote_rewards_accounts.rewards.len(),
-            vote_rewards_accounts.accounts_to_store.len()
+            vote_rewards_accounts.accounts_with_rewards.len(),
+            vote_rewards_accounts.accounts_with_rewards.len()
         );
-        assert_eq!(vote_rewards_accounts.rewards.len(), 1);
-        let rewards = &vote_rewards_accounts.rewards[0];
-        let account = &vote_rewards_accounts.accounts_to_store[0].1;
+        assert_eq!(vote_rewards_accounts.accounts_with_rewards.len(), 1);
+        let (vote_pubkey_from_result, rewards, account) =
+            &vote_rewards_accounts.accounts_with_rewards[0];
         let vote_rewards = 0;
         let commission = vote_state.commission;
         assert_eq!(account.lamports(), vote_account.lamports());
         assert!(accounts_equal(account, &vote_account));
         assert_eq!(
-            rewards.1,
+            *rewards,
             RewardInfo {
                 reward_type: RewardType::Voting,
                 lamports: vote_rewards as i64,
                 post_balance: vote_account.lamports(),
                 commission: Some(commission),
             }
         );
-        assert_eq!(&rewards.0, vote_pubkey);
+        assert_eq!(vote_pubkey_from_result, vote_pubkey);
 
         assert_eq!(stake_reward_calculation.stake_rewards.len(), 1);
         let expected_reward = {
diff --git a/runtime/src/bank/partitioned_epoch_rewards/mod.rs b/runtime/src/bank/partitioned_epoch_rewards/mod.rs
@@ -9,10 +9,13 @@ use {
         inflation_rewards::points::PointValue, stake_account::StakeAccount,
         stake_history::StakeHistory,
     },
-    solana_account::AccountSharedData,
+    solana_account::{AccountSharedData, ReadableAccount},
     solana_accounts_db::{
-        partitioned_rewards::PartitionedEpochRewardsConfig, stake_rewards::StakeReward,
+        partitioned_rewards::PartitionedEpochRewardsConfig,
+        stake_rewards::StakeReward,
+        storable_accounts::{AccountForStorage, StorableAccounts},
     },
+    solana_clock::Slot,
     solana_pubkey::Pubkey,
     solana_reward_info::RewardInfo,
     solana_stake_interface::state::{Delegation, Stake},
@@ -81,15 +84,59 @@ pub(crate) enum EpochRewardPhase {
 
 #[derive(Debug, Default)]
 pub(super) struct VoteRewardsAccounts {
-    /// reward info for each vote account pubkey.
-    /// This type is used by `update_reward_history()`
-    pub(super) rewards: Vec<(Pubkey, RewardInfo)>,
-    /// account to be stored, corresponds to pubkey in `rewards`
-    pub(super) accounts_to_store: Vec<(Pubkey, AccountSharedData)>,
+    /// accounts with rewards to be stored
+    pub(super) accounts_with_rewards: Vec<(Pubkey, RewardInfo, AccountSharedData)>,
     /// total lamports across all `vote_rewards`
     pub(super) total_vote_rewards_lamports: u64,
 }
 
+/// Wrapper struct to implement StorableAccounts for VoteRewardsAccounts
+pub(super) struct VoteRewardsAccountsStorable<'a> {
+    pub slot: Slot,
+    pub vote_rewards_accounts: &'a VoteRewardsAccounts,
+}
+
+impl<'a> StorableAccounts<'a> for VoteRewardsAccountsStorable<'a> {
+    fn account<Ret>(
+        &self,
+        index: usize,
+        mut callback: impl for<'local> FnMut(AccountForStorage<'local>) -> Ret,
+    ) -> Ret {
+        let (pubkey, _, account) = &self.vote_rewards_accounts.accounts_with_rewards[index];
+        callback((pubkey, account).into())
+    }
+
+    fn is_zero_lamport(&self, index: usize) -> bool {
+        self.vote_rewards_accounts.accounts_with_rewards[index]
+            .2
+            .lamports()
+            == 0
+    }
+
+    fn data_len(&self, index: usize) -> usize {
+        self.vote_rewards_accounts.accounts_with_rewards[index]
+            .2
+            .data()
+            .len()
+    }
+
+    fn pubkey(&self, index: usize) -> &Pubkey {
+        &self.vote_rewards_accounts.accounts_with_rewards[index].0
+    }
+
+    fn slot(&self, _index: usize) -> Slot {
+        self.target_slot()
+    }
+
+    fn target_slot(&self) -> Slot {
+        self.slot
+    }
+
+    fn len(&self) -> usize {
+        self.vote_rewards_accounts.accounts_with_rewards.len()
+    }
+}
+
 #[derive(Debug, Default)]
 /// result of calculating the stake rewards at end of epoch
 pub(super) struct StakeRewardCalculation {
diff --git a/runtime/src/bank/tests.rs b/runtime/src/bank/tests.rs
@@ -11205,8 +11205,11 @@ fn test_system_instruction_unsigned_transaction() {
 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());
+    assert_eq!(
+        result.accounts_with_rewards.len(),
+        result.accounts_with_rewards.len()
+    );
+    assert!(result.accounts_with_rewards.is_empty());
 }
 
 #[test]
@@ -11224,8 +11227,11 @@ fn test_calc_vote_accounts_to_store_overflow() {
         },
     );
     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());
+    assert_eq!(
+        result.accounts_with_rewards.len(),
+        result.accounts_with_rewards.len()
+    );
+    assert!(result.accounts_with_rewards.is_empty());
 }
 
 #[test]
@@ -11245,22 +11251,24 @@ fn test_calc_vote_accounts_to_store_normal() {
                 },
             );
             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].1;
+            assert_eq!(
+                result.accounts_with_rewards.len(),
+                result.accounts_with_rewards.len()
+            );
+            assert_eq!(result.accounts_with_rewards.len(), 1);
+            let (pubkey_result, rewards, account) = &result.accounts_with_rewards[0];
             _ = vote_account.checked_add_lamports(vote_rewards);
             assert!(accounts_equal(account, &vote_account));
             assert_eq!(
-                rewards.1,
+                *rewards,
                 RewardInfo {
                     reward_type: RewardType::Voting,
                     lamports: vote_rewards as i64,
                     post_balance: vote_account.lamports(),
                     commission: Some(commission),
                 }
             );
-            assert_eq!(rewards.0, pubkey);
+            assert_eq!(*pubkey_result, pubkey);
         }
     }
 }
