Refactor: Add separate methods for retrieving stake info from VoteAccounts (#26224)

runtime/src/bank.rs

@@ -2533,8 +2533,8 @@ impl Bank {
                     .stakes_cache
                     .stakes()
                     .vote_accounts()
-                    .iter()
-                    .map(|(pubkey, (stake, _))| (*pubkey, *stake))
+                    .delegated_stakes_iter()
+                    .map(|(pubkey, stake)| (*pubkey, stake))
                     .collect();
                 info!(
                     "new epoch stakes, epoch: {}, stakes: {:#?}, total_stake: {}",
@@ -2854,7 +2854,7 @@ impl Bank {
                         let vote_account = match self.get_account_with_fixed_root(vote_pubkey) {
                             Some(vote_account) => {
                                 match cached_vote_account {
-                                    Some((_stake, cached_vote_account))
+                                    Some(cached_vote_account)
                                         if cached_vote_account == &vote_account => {}
                                     _ => {
                                         invalid_cached_vote_accounts.fetch_add(1, Relaxed);
@@ -2958,7 +2958,7 @@ impl Bank {
         let solana_vote_program: Pubkey = solana_vote_program::id();
         let vote_accounts_cache_miss_count = AtomicUsize::default();
         let get_vote_account = |vote_pubkey: &Pubkey| -> Option<VoteAccount> {
-            if let Some((_stake, vote_account)) = cached_vote_accounts.get(vote_pubkey) {
+            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
@@ -7070,7 +7070,7 @@ impl Bank {
     /// Vote account for the given vote account pubkey.
     pub fn get_vote_account(&self, vote_account: &Pubkey) -> Option<VoteAccount> {
         let stakes = self.stakes_cache.stakes();
-        let (_stake, ref vote_account) = stakes.vote_accounts().get(vote_account)?;
+        let vote_account = stakes.vote_accounts().get(vote_account)?;
         Some(vote_account.clone())
     }
 

runtime/src/stakes.rs

@@ -292,7 +292,7 @@ impl Stakes<StakeAccount> {
         self.vote_accounts = self
             .vote_accounts
             .iter()
-            .map(|(&vote_pubkey, (_ /*stake*/, vote_account))| {
+            .map(|(&vote_pubkey, vote_account)| {
                 let delegated_stake = delegated_stakes
                     .get(&vote_pubkey)
                     .copied()
@@ -320,7 +320,7 @@ impl Stakes<StakeAccount> {
     /// Sum the lamports of the vote accounts and the delegated stake
     pub fn vote_balance_and_staked(&self) -> u64 {
         let get_stake = |stake_account: &StakeAccount| stake_account.delegation().stake;
-        let get_lamports = |(_, (_, vote_account)): (_, &(_, VoteAccount))| vote_account.lamports();
+        let get_lamports = |(_, vote_account): (_, &VoteAccount)| vote_account.lamports();
 
         self.stake_delegations.values().map(get_stake).sum::<u64>()
             + self.vote_accounts.iter().map(get_lamports).sum::<u64>()
@@ -377,8 +377,7 @@ impl Stakes<StakeAccount> {
     }
 
     pub(crate) fn highest_staked_node(&self) -> Option<Pubkey> {
-        let key = |(_pubkey, (stake, _vote_account)): &(_, &(u64, _))| *stake;
-        let (_pubkey, (_stake, vote_account)) = self.vote_accounts.iter().max_by_key(key)?;
+        let vote_account = self.vote_accounts.find_max_by_delegated_stake()?;
         Some(vote_account.vote_state().as_ref().ok()?.node_pubkey)
     }
 }
@@ -571,7 +570,7 @@ pub mod tests {
                 let vote_accounts = stakes.vote_accounts();
                 assert!(vote_accounts.get(&vote_pubkey).is_some());
                 assert_eq!(
-                    vote_accounts.get(&vote_pubkey).unwrap().0,
+                    vote_accounts.get_delegated_stake(&vote_pubkey),
                     stake.stake(i, None)
                 );
             }
@@ -583,7 +582,7 @@ pub mod tests {
                 let vote_accounts = stakes.vote_accounts();
                 assert!(vote_accounts.get(&vote_pubkey).is_some());
                 assert_eq!(
-                    vote_accounts.get(&vote_pubkey).unwrap().0,
+                    vote_accounts.get_delegated_stake(&vote_pubkey),
                     stake.stake(i, None)
                 ); // stays old stake, because only 10 is activated
             }
@@ -597,7 +596,7 @@ pub mod tests {
                 let vote_accounts = stakes.vote_accounts();
                 assert!(vote_accounts.get(&vote_pubkey).is_some());
                 assert_eq!(
-                    vote_accounts.get(&vote_pubkey).unwrap().0,
+                    vote_accounts.get_delegated_stake(&vote_pubkey),
                     stake.stake(i, None)
                 ); // now stake of 42 is activated
             }
@@ -608,7 +607,7 @@ pub mod tests {
                 let stakes = stakes_cache.stakes();
                 let vote_accounts = stakes.vote_accounts();
                 assert!(vote_accounts.get(&vote_pubkey).is_some());
-                assert_eq!(vote_accounts.get(&vote_pubkey).unwrap().0, 0);
+                assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 0);
             }
         }
     }
@@ -654,7 +653,7 @@ pub mod tests {
             let stakes = stakes_cache.stakes();
             let vote_accounts = stakes.vote_accounts();
             assert!(vote_accounts.get(&vote_pubkey).is_some());
-            assert_eq!(vote_accounts.get(&vote_pubkey).unwrap().0, 10);
+            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 10);
         }
 
         vote_account.set_lamports(0);
@@ -664,6 +663,7 @@ pub mod tests {
             let stakes = stakes_cache.stakes();
             let vote_accounts = stakes.vote_accounts();
             assert!(vote_accounts.get(&vote_pubkey).is_none());
+            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 0);
         }
 
         vote_account.set_lamports(1);
@@ -673,7 +673,7 @@ pub mod tests {
             let stakes = stakes_cache.stakes();
             let vote_accounts = stakes.vote_accounts();
             assert!(vote_accounts.get(&vote_pubkey).is_some());
-            assert_eq!(vote_accounts.get(&vote_pubkey).unwrap().0, 10);
+            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 10);
         }
 
         // Vote account too big
@@ -687,6 +687,7 @@ pub mod tests {
             let stakes = stakes_cache.stakes();
             let vote_accounts = stakes.vote_accounts();
             assert!(vote_accounts.get(&vote_pubkey).is_none());
+            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 0);
         }
 
         // Vote account uninitialized
@@ -699,6 +700,7 @@ pub mod tests {
             let stakes = stakes_cache.stakes();
             let vote_accounts = stakes.vote_accounts();
             assert!(vote_accounts.get(&vote_pubkey).is_none());
+            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 0);
         }
 
         vote_account.set_data(cache_data);
@@ -708,7 +710,7 @@ pub mod tests {
             let stakes = stakes_cache.stakes();
             let vote_accounts = stakes.vote_accounts();
             assert!(vote_accounts.get(&vote_pubkey).is_some());
-            assert_eq!(vote_accounts.get(&vote_pubkey).unwrap().0, 10);
+            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 10);
         }
     }
 
@@ -738,11 +740,11 @@ pub mod tests {
             let vote_accounts = stakes.vote_accounts();
             assert!(vote_accounts.get(&vote_pubkey).is_some());
             assert_eq!(
-                vote_accounts.get(&vote_pubkey).unwrap().0,
+                vote_accounts.get_delegated_stake(&vote_pubkey),
                 stake.stake(stakes.epoch, Some(&stakes.stake_history))
             );
             assert!(vote_accounts.get(&vote_pubkey2).is_some());
-            assert_eq!(vote_accounts.get(&vote_pubkey2).unwrap().0, 0);
+            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey2), 0);
         }
 
         // delegates to vote_pubkey2
@@ -752,10 +754,10 @@ pub mod tests {
             let stakes = stakes_cache.stakes();
             let vote_accounts = stakes.vote_accounts();
             assert!(vote_accounts.get(&vote_pubkey).is_some());
-            assert_eq!(vote_accounts.get(&vote_pubkey).unwrap().0, 0);
+            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 0);
             assert!(vote_accounts.get(&vote_pubkey2).is_some());
             assert_eq!(
-                vote_accounts.get(&vote_pubkey2).unwrap().0,
+                vote_accounts.get_delegated_stake(&vote_pubkey2),
                 stake.stake(stakes.epoch, Some(&stakes.stake_history))
             );
         }
@@ -782,7 +784,7 @@ pub mod tests {
             let stakes = stakes_cache.stakes();
             let vote_accounts = stakes.vote_accounts();
             assert!(vote_accounts.get(&vote_pubkey).is_some());
-            assert_eq!(vote_accounts.get(&vote_pubkey).unwrap().0, 20);
+            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 20);
         }
     }
 
@@ -801,7 +803,7 @@ pub mod tests {
             let stakes = stakes_cache.stakes();
             let vote_accounts = stakes.vote_accounts();
             assert_eq!(
-                vote_accounts.get(&vote_pubkey).unwrap().0,
+                vote_accounts.get_delegated_stake(&vote_pubkey),
                 stake.stake(stakes.epoch, Some(&stakes.stake_history))
             );
         }
@@ -811,7 +813,7 @@ pub mod tests {
             let stakes = stakes_cache.stakes();
             let vote_accounts = stakes.vote_accounts();
             assert_eq!(
-                vote_accounts.get(&vote_pubkey).unwrap().0,
+                vote_accounts.get_delegated_stake(&vote_pubkey),
                 stake.stake(stakes.epoch, Some(&stakes.stake_history))
             );
         }
@@ -834,7 +836,7 @@ pub mod tests {
             let stakes = stakes_cache.stakes();
             let vote_accounts = stakes.vote_accounts();
             assert!(vote_accounts.get(&vote_pubkey).is_some());
-            assert_eq!(vote_accounts.get(&vote_pubkey).unwrap().0, 10);
+            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 10);
         }
 
         // not a stake account, and whacks above entry
@@ -846,7 +848,7 @@ pub mod tests {
             let stakes = stakes_cache.stakes();
             let vote_accounts = stakes.vote_accounts();
             assert!(vote_accounts.get(&vote_pubkey).is_some());
-            assert_eq!(vote_accounts.get(&vote_pubkey).unwrap().0, 0);
+            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 0);
         }
     }
 
@@ -861,10 +863,17 @@ pub mod tests {
         let stakes_cache = StakesCache::default();
         impl Stakes<StakeAccount> {
             pub fn vote_balance_and_warmed_staked(&self) -> u64 {
-                self.vote_accounts
+                let vote_balance: u64 = self
+                    .vote_accounts
                     .iter()
-                    .map(|(_pubkey, (staked, account))| staked + account.lamports())
-                    .sum()
+                    .map(|(_pubkey, account)| account.lamports())
+                    .sum();
+                let warmed_stake: u64 = self
+                    .vote_accounts
+                    .delegated_stakes_iter()
+                    .map(|(_pubkey, stake)| stake)
+                    .sum();
+                vote_balance + warmed_stake
             }
         }
 

runtime/src/vote_account.rs

@@ -120,8 +120,10 @@ impl VoteAccounts {
         self.staked_nodes.read().unwrap().clone()
     }
 
-    pub fn get(&self, pubkey: &Pubkey) -> Option<&(/*stake:*/ u64, VoteAccount)> {
-        self.vote_accounts.get(pubkey)
+    pub fn get(&self, pubkey: &Pubkey) -> Option<&VoteAccount> {
+        self.vote_accounts
+            .get(pubkey)
+            .map(|(_stake, vote_account)| vote_account)
     }
 
     pub fn get_delegated_stake(&self, pubkey: &Pubkey) -> u64 {
@@ -131,8 +133,22 @@ impl VoteAccounts {
             .unwrap_or_default()
     }
 
-    pub(crate) fn iter(&self) -> impl Iterator<Item = (&Pubkey, &(u64, VoteAccount))> {
-        self.vote_accounts.iter()
+    pub(crate) fn iter(&self) -> impl Iterator<Item = (&Pubkey, &VoteAccount)> {
+        self.vote_accounts
+            .iter()
+            .map(|(vote_pubkey, (_stake, vote_account))| (vote_pubkey, vote_account))
+    }
+
+    pub(crate) fn delegated_stakes_iter(&self) -> impl Iterator<Item = (&Pubkey, u64)> {
+        self.vote_accounts
+            .iter()
+            .map(|(vote_pubkey, (stake, ..))| (vote_pubkey, *stake))
+    }
+
+    pub(crate) fn find_max_by_delegated_stake(&self) -> Option<&VoteAccount> {
+        let key = |(_pubkey, (stake, _vote_account)): &(_, &(u64, _))| *stake;
+        let (_pubkey, (_stake, vote_account)) = self.vote_accounts.iter().max_by_key(key)?;
+        Some(vote_account)
     }
 
     pub(crate) fn insert(&mut self, pubkey: Pubkey, (stake, vote_account): (u64, VoteAccount)) {
@@ -611,7 +627,7 @@ mod tests {
         ));
         assert_ne!(vote_accounts_hashmap, vote_accounts.vote_accounts);
         let other = (more_stake, vote_account);
-        for (pk, value) in vote_accounts.iter() {
+        for (pk, value) in vote_accounts.vote_accounts.iter() {
             if *pk != pubkey {
                 assert_eq!(value, &vote_accounts_hashmap[pk]);
             } else {