add calculate_stake_vote_rewards for partitioned rewards (#32066)

* add calculate_stake_vote_rewards for partitioned rewards

* check stake reward in test

* pr feedback

* add comment

---------

Co-authored-by: HaoranYi <haoran.yi@solana.com>

runtime/src/bank.rs

@@ -1217,6 +1217,16 @@ impl WorkingSlot for Bank {
     }
 }
 
+#[allow(dead_code)]
+#[derive(Debug, Default)]
+/// result of calculating the stake rewards at end of epoch
+struct StakeRewardCalculation {
+    /// each individual stake account to reward
+    stake_rewards: StakeRewards,
+    /// total lamports across all `stake_rewards`
+    total_stake_rewards_lamports: u64,
+}
+
 impl Bank {
     pub fn default_for_tests() -> Self {
         Self::default_with_accounts(Accounts::default_for_tests())
@@ -3125,6 +3135,148 @@ impl Bank {
         (points > 0).then_some(PointValue { rewards, points })
     }
 
+    #[allow(dead_code)]
+    /// Calculates epoch rewards for stake/vote accounts
+    /// Returns vote rewards, stake rewards, and the sum of all stake rewards in lamports
+    fn calculate_stake_vote_rewards(
+        &self,
+        reward_calculate_params: &EpochRewardCalculateParamInfo,
+        rewarded_epoch: Epoch,
+        point_value: PointValue,
+        thread_pool: &ThreadPool,
+        reward_calc_tracer: Option<impl RewardCalcTracer>,
+        metrics: &mut RewardsMetrics,
+    ) -> (VoteRewardsAccounts, StakeRewardCalculation) {
+        let credits_auto_rewind = self.credits_auto_rewind();
+        let EpochRewardCalculateParamInfo {
+            stake_history,
+            stake_delegations,
+            cached_vote_accounts,
+        } = reward_calculate_params;
+
+        let solana_vote_program: Pubkey = solana_vote_program::id();
+
+        let get_vote_account = |vote_pubkey: &Pubkey| -> Option<VoteAccount> {
+            if let Some(vote_account) = cached_vote_accounts.get(vote_pubkey) {
+                return Some(vote_account.clone());
+            }
+            // If accounts-db contains a valid vote account, then it should
+            // already have been cached in cached_vote_accounts; so the code
+            // below is only for sanity checking, and can be removed once
+            // the cache is deemed to be reliable.
+            let account = self.get_account_with_fixed_root(vote_pubkey)?;
+            VoteAccount::try_from(account).ok()
+        };
+
+        let vote_account_rewards: VoteRewards = DashMap::new();
+        let total_stake_rewards = AtomicU64::default();
+        let (stake_rewards, measure_stake_rewards_us) = measure_us!(thread_pool.install(|| {
+            stake_delegations
+                .par_iter()
+                .filter_map(|(stake_pubkey, stake_account)| {
+                    // curry closure to add the contextual stake_pubkey
+                    let reward_calc_tracer = reward_calc_tracer.as_ref().map(|outer| {
+                        // inner
+                        move |inner_event: &_| {
+                            outer(&RewardCalculationEvent::Staking(stake_pubkey, inner_event))
+                        }
+                    });
+
+                    let stake_pubkey = **stake_pubkey;
+                    let stake_account = (*stake_account).to_owned();
+
+                    let delegation = stake_account.delegation();
+                    let (mut stake_account, stake_state) =
+                        <(AccountSharedData, StakeState)>::from(stake_account);
+                    let vote_pubkey = delegation.voter_pubkey;
+                    let vote_account = match get_vote_account(&vote_pubkey) {
+                        Some(vote_account) => vote_account,
+                        None => {
+                            return None;
+                        }
+                    };
+                    if vote_account.owner() != &solana_vote_program {
+                        return None;
+                    }
+                    let vote_state = match vote_account.vote_state().cloned() {
+                        Ok(vote_state) => vote_state,
+                        Err(_) => {
+                            return None;
+                        }
+                    };
+
+                    let pre_lamport = stake_account.lamports();
+
+                    let redeemed = stake_state::redeem_rewards(
+                        rewarded_epoch,
+                        stake_state,
+                        &mut stake_account,
+                        &vote_state,
+                        &point_value,
+                        Some(stake_history),
+                        reward_calc_tracer.as_ref(),
+                        credits_auto_rewind,
+                    );
+
+                    let post_lamport = stake_account.lamports();
+
+                    if let Ok((stakers_reward, voters_reward)) = redeemed {
+                        debug!(
+                            "calculated reward: {} {} {} {}",
+                            stake_pubkey, pre_lamport, post_lamport, stakers_reward
+                        );
+
+                        // track voter rewards
+                        let mut voters_reward_entry = vote_account_rewards
+                            .entry(vote_pubkey)
+                            .or_insert(VoteReward {
+                                vote_account: vote_account.into(),
+                                commission: vote_state.commission,
+                                vote_rewards: 0,
+                                vote_needs_store: false,
+                            });
+
+                        voters_reward_entry.vote_needs_store = true;
+                        voters_reward_entry.vote_rewards = voters_reward_entry
+                            .vote_rewards
+                            .saturating_add(voters_reward);
+
+                        let post_balance = stake_account.lamports();
+                        total_stake_rewards.fetch_add(stakers_reward, Relaxed);
+                        return Some(StakeReward {
+                            stake_pubkey,
+                            stake_reward_info: RewardInfo {
+                                reward_type: RewardType::Staking,
+                                lamports: i64::try_from(stakers_reward).unwrap(),
+                                post_balance,
+                                commission: Some(vote_state.commission),
+                            },
+                            stake_account,
+                        });
+                    } else {
+                        debug!(
+                            "stake_state::redeem_rewards() failed for {}: {:?}",
+                            stake_pubkey, redeemed
+                        );
+                    }
+                    None
+                })
+                .collect()
+        }));
+        let (vote_rewards, measure_vote_rewards_us) =
+            measure_us!(Self::calc_vote_accounts_to_store(vote_account_rewards));
+
+        metrics.redeem_rewards_us += measure_stake_rewards_us + measure_vote_rewards_us;
+
+        (
+            vote_rewards,
+            StakeRewardCalculation {
+                stake_rewards,
+                total_stake_rewards_lamports: total_stake_rewards.load(Relaxed),
+            },
+        )
+    }
+
     fn redeem_rewards(
         &self,
         vote_with_stake_delegations_map: DashMap<Pubkey, VoteWithStakeDelegations>,

runtime/src/bank/tests.rs

@@ -12506,7 +12506,7 @@ fn test_squash_timing_add_assign() {
 }
 
 /// Helper function to create a bank that pays some rewards
-fn create_reward_bank(expected_num_delegations: usize) -> Bank {
+fn create_reward_bank(expected_num_delegations: usize) -> (Bank, Vec<Pubkey>, Vec<Pubkey>) {
     let validator_keypairs = (0..expected_num_delegations)
         .map(|_| ValidatorVoteKeypairs::new_rand())
         .collect::<Vec<_>>();
@@ -12521,7 +12521,7 @@ fn create_reward_bank(expected_num_delegations: usize) -> Bank {
 
     // 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 {
+    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
@@ -12541,7 +12541,17 @@ fn create_reward_bank(expected_num_delegations: usize) -> Bank {
         }
         bank.store_account_and_update_capitalization(&vote_id, &vote_account);
     }
-    bank
+    (
+        bank,
+        validator_keypairs
+            .iter()
+            .map(|k| k.vote_keypair.pubkey())
+            .collect(),
+        validator_keypairs
+            .iter()
+            .map(|k| k.stake_keypair.pubkey())
+            .collect(),
+    )
 }
 
 #[test]
@@ -12549,7 +12559,7 @@ fn test_rewards_point_calculation() {
     solana_logger::setup();
 
     let expected_num_delegations = 100;
-    let bank = create_reward_bank(expected_num_delegations);
+    let (bank, _, _) = create_reward_bank(expected_num_delegations);
 
     let thread_pool = ThreadPoolBuilder::new().num_threads(1).build().unwrap();
     let mut rewards_metrics = RewardsMetrics::default();
@@ -13061,3 +13071,91 @@ fn test_calc_vote_accounts_to_store_normal() {
         }
     }
 }
+
+#[test]
+fn test_calculate_stake_vote_rewards() {
+    solana_logger::setup();
+
+    let expected_num_delegations = 1;
+    let (bank, voters, stakers) = create_reward_bank(expected_num_delegations);
+
+    let vote_pubkey = voters.first().unwrap();
+    let mut vote_account = bank
+        .load_slow_with_fixed_root(&bank.ancestors, vote_pubkey)
+        .unwrap()
+        .0;
+
+    let stake_pubkey = stakers.first().unwrap();
+    let stake_account = bank
+        .load_slow_with_fixed_root(&bank.ancestors, stake_pubkey)
+        .unwrap()
+        .0;
+
+    let thread_pool = ThreadPoolBuilder::new().num_threads(1).build().unwrap();
+    let mut rewards_metrics = RewardsMetrics::default();
+
+    let point_value = PointValue {
+        rewards: 100000, // lamports to split
+        points: 1000,    // over these points
+    };
+    let tracer = |_event: &RewardCalculationEvent| {};
+    let reward_calc_tracer = Some(tracer);
+    let rewarded_epoch = bank.epoch();
+    let stakes: RwLockReadGuard<Stakes<StakeAccount<Delegation>>> = bank.stakes_cache.stakes();
+    let reward_calculate_param = bank.get_epoch_reward_calculate_param_info(&stakes);
+    let (vote_rewards_accounts, stake_reward_calculation) = bank.calculate_stake_vote_rewards(
+        &reward_calculate_param,
+        rewarded_epoch,
+        point_value,
+        &thread_pool,
+        reward_calc_tracer,
+        &mut rewards_metrics,
+    );
+
+    assert_eq!(
+        vote_rewards_accounts.rewards.len(),
+        vote_rewards_accounts.accounts_to_store.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];
+    let vote_rewards = 0;
+    let commision = 0;
+    _ = vote_account.checked_add_lamports(vote_rewards);
+    assert_eq!(
+        account.as_ref().unwrap().lamports(),
+        vote_account.lamports()
+    );
+    assert!(accounts_equal(account.as_ref().unwrap(), &vote_account));
+    assert_eq!(
+        rewards.1,
+        RewardInfo {
+            reward_type: RewardType::Voting,
+            lamports: vote_rewards as i64,
+            post_balance: vote_account.lamports(),
+            commission: Some(commision),
+        }
+    );
+    assert_eq!(&rewards.0, vote_pubkey);
+
+    assert_eq!(stake_reward_calculation.stake_rewards.len(), 1);
+    assert_eq!(
+        &stake_reward_calculation.stake_rewards[0].stake_pubkey,
+        stake_pubkey
+    );
+
+    let original_stake_lamport = stake_account.lamports();
+    let rewards = stake_reward_calculation.stake_rewards[0]
+        .stake_reward_info
+        .lamports;
+    let expected_reward_info = RewardInfo {
+        reward_type: RewardType::Staking,
+        lamports: rewards,
+        post_balance: original_stake_lamport + rewards as u64,
+        commission: Some(commision),
+    };
+    assert_eq!(
+        stake_reward_calculation.stake_rewards[0].stake_reward_info,
+        expected_reward_info,
+    );
+}