From b89ce94d33c4019bd332dbab0ca76ca5dbc5ceea Mon Sep 17 00:00:00 2001
From: "Jeff Washington (jwash)" <jeff.washington@solana.com>
Date: Mon, 12 Jun 2023 16:38:30 -0500
Subject: [PATCH] add store_vote_accounts_partitioned (#32061)

* add store_vote_accounts_partitioned

* pr feedback
---
 runtime/src/bank.rs       | 29 +++++++++++++-
 runtime/src/bank/tests.rs | 80 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 108 insertions(+), 1 deletion(-)

diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
index 232912c81..17e5aadd6 100644
--- a/runtime/src/bank.rs
+++ b/runtime/src/bank.rs
@@ -1105,7 +1105,7 @@ struct VoteReward {
 
 type VoteRewards = DashMap<Pubkey, VoteReward>;
 #[allow(dead_code)]
-#[derive(Debug)]
+#[derive(Debug, Default)]
 struct VoteRewardsAccounts {
     /// reward info for each vote account pubkey.
     /// This type is used by `update_reward_history()`
@@ -3164,6 +3164,33 @@ impl Bank {
             .sum::<i64>() as u64
     }
 
+    #[allow(dead_code)]
+    fn store_vote_accounts_partitioned(
+        &self,
+        vote_account_rewards: VoteRewardsAccounts,
+        metrics: &mut RewardsMetrics,
+    ) -> Vec<(Pubkey, RewardInfo)> {
+        let (_, measure_us) = measure_us!({
+            // reformat data to make it not sparse.
+            // `StorableAccounts` does not efficiently handle sparse data.
+            // Not all entries in `vote_account_rewards.accounts_to_store` have a Some(account) to store.
+            let to_store = vote_account_rewards
+                .accounts_to_store
+                .iter()
+                .filter_map(|account| account.as_ref())
+                .enumerate()
+                .map(|(i, account)| (&vote_account_rewards.rewards[i].0, account))
+                .collect::<Vec<_>>();
+            self.store_accounts((self.slot(), &to_store[..], self.include_slot_in_hash()));
+        });
+
+        metrics
+            .store_vote_accounts_us
+            .fetch_add(measure_us, Relaxed);
+
+        vote_account_rewards.rewards
+    }
+
     fn store_vote_accounts(
         &self,
         vote_account_rewards: VoteRewards,
diff --git a/runtime/src/bank/tests.rs b/runtime/src/bank/tests.rs
index 664b1639b..4bae7b299 100644
--- a/runtime/src/bank/tests.rs
+++ b/runtime/src/bank/tests.rs
@@ -159,6 +159,30 @@ impl StakeReward {
     }
 }
 
+impl VoteReward {
+    pub fn new_random() -> Self {
+        let mut rng = rand::thread_rng();
+
+        let validator_pubkey = solana_sdk::pubkey::new_rand();
+        let validator_stake_lamports = rng.gen_range(1, 200);
+        let validator_voting_keypair = Keypair::new();
+
+        let validator_vote_account = vote_state::create_account(
+            &validator_voting_keypair.pubkey(),
+            &validator_pubkey,
+            rng.gen_range(1, 20),
+            validator_stake_lamports,
+        );
+
+        Self {
+            vote_account: validator_vote_account,
+            commission: rng.gen_range(1, 20),
+            vote_rewards: rng.gen_range(1, 200),
+            vote_needs_store: rng.gen_range(1, 20) > 10,
+        }
+    }
+}
+
 #[test]
 fn test_race_register_tick_freeze() {
     solana_logger::setup();
@@ -12661,6 +12685,62 @@ fn test_update_reward_history_in_partition_empty() {
     assert_eq!(num_in_history, 0);
 }
 
+#[test]
+fn test_store_vote_accounts_partitioned() {
+    let (genesis_config, _mint_keypair) = create_genesis_config(1_000_000 * LAMPORTS_PER_SOL);
+    let bank = Bank::new_for_tests(&genesis_config);
+
+    let expected_vote_rewards_num = 100;
+
+    let vote_rewards = (0..expected_vote_rewards_num)
+        .map(|_| Some((Pubkey::new_unique(), VoteReward::new_random())))
+        .collect::<Vec<_>>();
+
+    let mut vote_rewards_account = VoteRewardsAccounts::default();
+    vote_rewards.iter().for_each(|e| {
+        if let Some(p) = &e {
+            let info = RewardInfo {
+                reward_type: RewardType::Voting,
+                lamports: p.1.vote_rewards as i64,
+                post_balance: p.1.vote_rewards,
+                commission: Some(p.1.commission),
+            };
+            vote_rewards_account.rewards.push((p.0, info));
+            vote_rewards_account
+                .accounts_to_store
+                .push(e.as_ref().map(|p| p.1.vote_account.clone()));
+        }
+    });
+
+    let mut metrics = RewardsMetrics::default();
+
+    let stored_vote_accounts =
+        bank.store_vote_accounts_partitioned(vote_rewards_account, &mut metrics);
+    assert_eq!(expected_vote_rewards_num, stored_vote_accounts.len());
+
+    // load accounts to make sure they were stored correctly
+    vote_rewards.iter().for_each(|e| {
+        if let Some(p) = &e {
+            let (k, account) = (p.0, p.1.vote_account.clone());
+            let loaded_account = bank.load_slow_with_fixed_root(&bank.ancestors, &k).unwrap();
+            assert!(accounts_equal(&loaded_account.0, &account));
+        }
+    });
+}
+
+#[test]
+fn test_store_vote_accounts_partitioned_empty() {
+    let (genesis_config, _mint_keypair) = create_genesis_config(1_000_000 * LAMPORTS_PER_SOL);
+    let bank = Bank::new_for_tests(&genesis_config);
+
+    let expected = 0;
+    let vote_rewards = VoteRewardsAccounts::default();
+    let mut metrics = RewardsMetrics::default();
+
+    let stored_vote_accounts = bank.store_vote_accounts_partitioned(vote_rewards, &mut metrics);
+    assert_eq!(expected, stored_vote_accounts.len());
+}
+
 #[test]
 fn test_system_instruction_allocate() {
     let (genesis_config, mint_keypair) = create_genesis_config(sol_to_lamports(1.0));