improves performance in replay-stage (#14217)

bank::vote_accounts returns a hash-map which is slow to iterate, but all uses
only require an iterator:
https://github.com/solana-labs/solana/blob/b3dc98856/runtime/src/bank.rs#L4300-L4306
Similarly, calculate_stake_weighted_timestamp takes a hash-map whereas it only
requires an iterator:
https://github.com/solana-labs/solana/blob/b3dc98856/sdk/src/stake_weighted_timestamp.rs#L21-L28

core/src/consensus.rs

@@ -216,6 +216,7 @@ impl Tower {
     where
         F: IntoIterator<Item = (Pubkey, (u64, ArcVoteAccount))>,
     {
+        let mut vote_slots = HashSet::new();
         let mut voted_stakes = HashMap::new();
         let mut total_stake = 0;
         let mut bank_weight = 0;
@@ -278,7 +279,7 @@ impl Tower {
 
             for vote in &vote_state.votes {
                 bank_weight += vote.lockout() as u128 * voted_stake as u128;
-                Self::populate_ancestor_voted_stakes(&mut voted_stakes, &vote, ancestors);
+                vote_slots.insert(vote.slot);
             }
 
             if start_root != vote_state.root_slot {
@@ -289,7 +290,7 @@ impl Tower {
                     };
                     trace!("ROOT: {}", vote.slot);
                     bank_weight += vote.lockout() as u128 * voted_stake as u128;
-                    Self::populate_ancestor_voted_stakes(&mut voted_stakes, &vote, ancestors);
+                    vote_slots.insert(vote.slot);
                 }
             }
             if let Some(root) = vote_state.root_slot {
@@ -298,7 +299,7 @@ impl Tower {
                     slot: root,
                 };
                 bank_weight += vote.lockout() as u128 * voted_stake as u128;
-                Self::populate_ancestor_voted_stakes(&mut voted_stakes, &vote, ancestors);
+                vote_slots.insert(vote.slot);
             }
 
             // The last vote in the vote stack is a simulated vote on bank_slot, which
@@ -326,6 +327,9 @@ impl Tower {
             total_stake += voted_stake;
         }
 
+        // TODO: populate_ancestor_voted_stakes only adds zeros. Comment why
+        // that is necessary (if so).
+        Self::populate_ancestor_voted_stakes(&mut voted_stakes, vote_slots, ancestors);
         ComputedBankState {
             voted_stakes,
             total_stake,
@@ -766,20 +770,19 @@ impl Tower {
     /// Update lockouts for all the ancestors
     pub(crate) fn populate_ancestor_voted_stakes(
         voted_stakes: &mut VotedStakes,
-        vote: &Lockout,
+        vote_slots: impl IntoIterator<Item = Slot>,
         ancestors: &HashMap<Slot, HashSet<Slot>>,
     ) {
         // If there's no ancestors, that means this slot must be from before the current root,
         // in which case the lockouts won't be calculated in bank_weight anyways, so ignore
         // this slot
-        let vote_slot_ancestors = ancestors.get(&vote.slot);
-        if vote_slot_ancestors.is_none() {
-            return;
-        }
-        let mut slot_with_ancestors = vec![vote.slot];
-        slot_with_ancestors.extend(vote_slot_ancestors.unwrap());
-        for slot in slot_with_ancestors {
-            voted_stakes.entry(slot).or_default();
+        for vote_slot in vote_slots {
+            if let Some(slot_ancestors) = ancestors.get(&vote_slot) {
+                voted_stakes.entry(vote_slot).or_default();
+                for slot in slot_ancestors {
+                    voted_stakes.entry(*slot).or_default();
+                }
+            }
         }
     }
 
@@ -793,15 +796,11 @@ impl Tower {
     ) {
         // If there's no ancestors, that means this slot must be from
         // before the current root, so ignore this slot
-        let vote_slot_ancestors = ancestors.get(&voted_slot);
-        if vote_slot_ancestors.is_none() {
-            return;
-        }
-        let mut slot_with_ancestors = vec![voted_slot];
-        slot_with_ancestors.extend(vote_slot_ancestors.unwrap());
-        for slot in slot_with_ancestors {
-            let current = voted_stakes.entry(slot).or_default();
-            *current += voted_stake;
+        if let Some(vote_slot_ancestors) = ancestors.get(&voted_slot) {
+            *voted_stakes.entry(voted_slot).or_default() += voted_stake;
+            for slot in vote_slot_ancestors {
+                *voted_stakes.entry(*slot).or_default() += voted_stake;
+            }
         }
     }
 

core/src/validator.rs

@@ -1154,10 +1154,10 @@ fn get_stake_percent_in_gossip(bank: &Bank, cluster_info: &ClusterInfo, log: boo
     let my_shred_version = cluster_info.my_shred_version();
     let my_id = cluster_info.id();
 
-    for (activated_stake, vote_account) in bank.vote_accounts().values() {
+    for (_, (activated_stake, vote_account)) in bank.vote_accounts() {
         total_activated_stake += activated_stake;
 
-        if *activated_stake == 0 {
+        if activated_stake == 0 {
             continue;
         }
         let vote_state_node_pubkey = vote_account
@@ -1179,13 +1179,13 @@ fn get_stake_percent_in_gossip(bank: &Bank, cluster_info: &ClusterInfo, log: boo
                 online_stake += activated_stake;
             } else {
                 wrong_shred_stake += activated_stake;
-                wrong_shred_nodes.push((*activated_stake, vote_state_node_pubkey));
+                wrong_shred_nodes.push((activated_stake, vote_state_node_pubkey));
             }
         } else if vote_state_node_pubkey == my_id {
             online_stake += activated_stake; // This node is online
         } else {
             offline_stake += activated_stake;
-            offline_nodes.push((*activated_stake, vote_state_node_pubkey));
+            offline_nodes.push((activated_stake, vote_state_node_pubkey));
         }
     }
 

runtime/src/bank.rs

@@ -1877,35 +1877,36 @@ impl Bank {
         epoch_start_timestamp: Option<(Slot, UnixTimestamp)>,
     ) -> Option<UnixTimestamp> {
         let mut get_timestamp_estimate_time = Measure::start("get_timestamp_estimate");
-        let recent_timestamps: HashMap<Pubkey, (Slot, UnixTimestamp)> = self
-            .vote_accounts()
-            .into_iter()
-            .filter_map(|(pubkey, (_, account))| {
-                account.vote_state().as_ref().ok().and_then(|state| {
-                    let timestamp_slot = state.last_timestamp.slot;
-                    if (self
-                        .feature_set
-                        .is_active(&feature_set::timestamp_bounding::id())
-                        && self.slot().checked_sub(timestamp_slot)?
-                            <= self.epoch_schedule().slots_per_epoch)
-                        || self.slot().checked_sub(timestamp_slot)?
-                            <= DEPRECATED_TIMESTAMP_SLOT_RANGE as u64
+        let timestamp_bounding_enabled = self
+            .feature_set
+            .is_active(&feature_set::timestamp_bounding::id());
+        let slots_per_epoch = self.epoch_schedule().slots_per_epoch;
+        let recent_timestamps =
+            self.vote_accounts()
+                .into_iter()
+                .filter_map(|(pubkey, (_, account))| {
+                    let vote_state = account.vote_state();
+                    let vote_state = vote_state.as_ref().ok()?;
+                    let slot_delta = self.slot().checked_sub(vote_state.last_timestamp.slot)?;
+                    if (timestamp_bounding_enabled && slot_delta <= slots_per_epoch)
+                        || slot_delta <= DEPRECATED_TIMESTAMP_SLOT_RANGE as u64
                     {
                         Some((
                             pubkey,
-                            (state.last_timestamp.slot, state.last_timestamp.timestamp),
+                            (
+                                vote_state.last_timestamp.slot,
+                                vote_state.last_timestamp.timestamp,
+                            ),
                         ))
                     } else {
                         None
                     }
-                })
-            })
-            .collect();
+                });
         let slot_duration = Duration::from_nanos(self.ns_per_slot as u64);
         let epoch = self.epoch_schedule().get_epoch(self.slot());
         let stakes = self.epoch_vote_accounts(epoch)?;
         let stake_weighted_timestamp = calculate_stake_weighted_timestamp(
-            &recent_timestamps,
+            recent_timestamps,
             stakes,
             self.slot(),
             slot_duration,
@@ -3169,25 +3170,24 @@ impl Bank {
     //
     // Ref: collect_fees
     #[allow(clippy::needless_collect)]
-    fn distribute_rent_to_validators(
-        &self,
-        vote_account_hashmap: &HashMap<Pubkey, (u64, ArcVoteAccount)>,
-        rent_to_be_distributed: u64,
-    ) {
+    fn distribute_rent_to_validators<I>(&self, vote_accounts: I, rent_to_be_distributed: u64)
+    where
+        I: IntoIterator<Item = (Pubkey, (u64, ArcVoteAccount))>,
+    {
         let mut total_staked = 0;
 
         // Collect the stake associated with each validator.
         // Note that a validator may be present in this vector multiple times if it happens to have
         // more than one staked vote account somehow
-        let mut validator_stakes = vote_account_hashmap
-            .iter()
+        let mut validator_stakes = vote_accounts
+            .into_iter()
             .filter_map(|(_vote_pubkey, (staked, account))| {
-                if *staked == 0 {
+                if staked == 0 {
                     None
                 } else {
-                    total_staked += *staked;
+                    total_staked += staked;
                     let node_pubkey = account.vote_state().as_ref().ok()?.node_pubkey;
-                    Some((node_pubkey, *staked))
+                    Some((node_pubkey, staked))
                 }
             })
             .collect::<Vec<(Pubkey, u64)>>();
@@ -3287,7 +3287,7 @@ impl Bank {
             return;
         }
 
-        self.distribute_rent_to_validators(&self.vote_accounts(), rent_to_be_distributed);
+        self.distribute_rent_to_validators(self.vote_accounts(), rent_to_be_distributed);
     }
 
     fn collect_rent(
@@ -4326,8 +4326,14 @@ impl Bank {
     ///   attributed to each account
     /// Note: This clones the entire vote-accounts hashmap. For a single
     /// account lookup use get_vote_account instead.
-    pub fn vote_accounts(&self) -> HashMap<Pubkey, (u64 /*stake*/, ArcVoteAccount)> {
-        self.stakes.read().unwrap().vote_accounts().clone()
+    pub fn vote_accounts(&self) -> Vec<(Pubkey, (u64 /*stake*/, ArcVoteAccount))> {
+        self.stakes
+            .read()
+            .unwrap()
+            .vote_accounts()
+            .iter()
+            .map(|(k, v)| (*k, v.clone()))
+            .collect()
     }
 
     /// Vote account for the given vote account pubkey along with the stake.
@@ -5697,7 +5703,7 @@ pub(crate) mod tests {
 
         let bank = Bank::new(&genesis_config);
         let old_validator_lamports = bank.get_balance(&validator_pubkey);
-        bank.distribute_rent_to_validators(&bank.vote_accounts(), RENT_TO_BE_DISTRIBUTED);
+        bank.distribute_rent_to_validators(bank.vote_accounts(), RENT_TO_BE_DISTRIBUTED);
         let new_validator_lamports = bank.get_balance(&validator_pubkey);
         assert_eq!(
             new_validator_lamports,
@@ -5711,7 +5717,7 @@ pub(crate) mod tests {
         let bank = std::panic::AssertUnwindSafe(Bank::new(&genesis_config));
         let old_validator_lamports = bank.get_balance(&validator_pubkey);
         let new_validator_lamports = std::panic::catch_unwind(|| {
-            bank.distribute_rent_to_validators(&bank.vote_accounts(), RENT_TO_BE_DISTRIBUTED);
+            bank.distribute_rent_to_validators(bank.vote_accounts(), RENT_TO_BE_DISTRIBUTED);
             bank.get_balance(&validator_pubkey)
         });
 
@@ -8484,7 +8490,7 @@ pub(crate) mod tests {
 
         bank.process_transaction(&transaction).unwrap();
 
-        let vote_accounts = bank.vote_accounts();
+        let vote_accounts = bank.vote_accounts().into_iter().collect::<HashMap<_, _>>();
 
         assert_eq!(vote_accounts.len(), 2);
 

sdk/src/stake_weighted_timestamp.rs

@@ -5,6 +5,7 @@ use solana_sdk::{
     pubkey::Pubkey,
 };
 use std::{
+    borrow::Borrow,
     collections::{BTreeMap, HashMap},
     time::Duration,
 };
@@ -18,14 +19,19 @@ pub enum EstimateType {
     Unbounded, // Deprecated.  Remove in the Solana v1.6.0 timeframe
 }
 
-pub fn calculate_stake_weighted_timestamp<T>(
-    unique_timestamps: &HashMap<Pubkey, (Slot, UnixTimestamp)>,
+pub fn calculate_stake_weighted_timestamp<I, K, V, T>(
+    unique_timestamps: I,
     stakes: &HashMap<Pubkey, (u64, T /*Account|ArcVoteAccount*/)>,
     slot: Slot,
     slot_duration: Duration,
     estimate_type: EstimateType,
     epoch_start_timestamp: Option<(Slot, UnixTimestamp)>,
-) -> Option<UnixTimestamp> {
+) -> Option<UnixTimestamp>
+where
+    I: IntoIterator<Item = (K, V)>,
+    K: Borrow<Pubkey>,
+    V: Borrow<(Slot, UnixTimestamp)>,
+{
     match estimate_type {
         EstimateType::Bounded => calculate_bounded_stake_weighted_timestamp(
             unique_timestamps,
@@ -43,17 +49,23 @@ pub fn calculate_stake_weighted_timestamp<T>(
     }
 }
 
-fn calculate_unbounded_stake_weighted_timestamp<T>(
-    unique_timestamps: &HashMap<Pubkey, (Slot, UnixTimestamp)>,
+fn calculate_unbounded_stake_weighted_timestamp<I, K, V, T>(
+    unique_timestamps: I,
     stakes: &HashMap<Pubkey, (u64, T /*Account|ArcVoteAccount*/)>,
     slot: Slot,
     slot_duration: Duration,
-) -> Option<UnixTimestamp> {
+) -> Option<UnixTimestamp>
+where
+    I: IntoIterator<Item = (K, V)>,
+    K: Borrow<Pubkey>,
+    V: Borrow<(Slot, UnixTimestamp)>,
+{
     let (stake_weighted_timestamps_sum, total_stake) = unique_timestamps
-        .iter()
-        .filter_map(|(vote_pubkey, (timestamp_slot, timestamp))| {
+        .into_iter()
+        .filter_map(|(vote_pubkey, slot_timestamp)| {
+            let (timestamp_slot, timestamp) = slot_timestamp.borrow();
             let offset = (slot - timestamp_slot) as u32 * slot_duration;
-            stakes.get(&vote_pubkey).map(|(stake, _account)| {
+            stakes.get(vote_pubkey.borrow()).map(|(stake, _account)| {
                 (
                     (*timestamp as u128 + offset.as_secs() as u128) * *stake as u128,
                     stake,
@@ -70,20 +82,26 @@ fn calculate_unbounded_stake_weighted_timestamp<T>(
     }
 }
 
-fn calculate_bounded_stake_weighted_timestamp<T>(
-    unique_timestamps: &HashMap<Pubkey, (Slot, UnixTimestamp)>,
+fn calculate_bounded_stake_weighted_timestamp<I, K, V, T>(
+    unique_timestamps: I,
     stakes: &HashMap<Pubkey, (u64, T /*Account|ArcVoteAccount*/)>,
     slot: Slot,
     slot_duration: Duration,
     epoch_start_timestamp: Option<(Slot, UnixTimestamp)>,
-) -> Option<UnixTimestamp> {
+) -> Option<UnixTimestamp>
+where
+    I: IntoIterator<Item = (K, V)>,
+    K: Borrow<Pubkey>,
+    V: Borrow<(Slot, UnixTimestamp)>,
+{
     let mut stake_per_timestamp: BTreeMap<UnixTimestamp, u128> = BTreeMap::new();
     let mut total_stake = 0;
-    for (vote_pubkey, (timestamp_slot, timestamp)) in unique_timestamps.iter() {
+    for (vote_pubkey, slot_timestamp) in unique_timestamps {
+        let (timestamp_slot, timestamp) = slot_timestamp.borrow();
         let offset = slot.saturating_sub(*timestamp_slot) as u32 * slot_duration;
         let estimate = timestamp + offset.as_secs() as i64;
         let stake = stakes
-            .get(&vote_pubkey)
+            .get(vote_pubkey.borrow())
             .map(|(stake, _account)| stake)
             .unwrap_or(&0);
         stake_per_timestamp