v3.1: Optimize LockoutIntervals in collect_vote_lockouts (backport of #8865) (#8965)

Optimize LockoutIntervals in collect_vote_lockouts (#8865)

* Optimize LockoutIntervals in collect_vote_lockouts

* Filter lockout_intervals by slot > root

* Revert "Filter lockout_intervals by slot > root"

This reverts commit 5a802aab9290da24524ffcff09cc15bc592b7a06.

* Update comments

(cherry picked from commit d9ad69e28a2788872b2b9171f8f5b319c8daa6f0)

Co-authored-by: Zach Brown <zach@zb.dev>

core/src/consensus.rs

@@ -19,7 +19,7 @@ use {
         tower_storage::{SavedTower, SavedTowerVersions, TowerStorage},
         tower_vote_state::TowerVoteState,
     },
-    crate::replay_stage::DUPLICATE_THRESHOLD,
+    crate::{consensus::progress_map::LockoutInterval, replay_stage::DUPLICATE_THRESHOLD},
     chrono::prelude::*,
     solana_clock::{Slot, UnixTimestamp},
     solana_hash::Hash,
@@ -41,10 +41,7 @@ use {
     std::{
         cmp::Ordering,
         collections::{HashMap, HashSet},
-        ops::{
-            Bound::{Included, Unbounded},
-            Deref,
-        },
+        ops::Deref,
     },
     thiserror::Error,
 };
@@ -168,8 +165,8 @@ pub(crate) struct ComputedBankState {
     pub voted_stakes: VotedStakes,
     pub total_stake: Stake,
     pub fork_stake: Stake,
-    // Tree of intervals of lockouts of the form [slot, slot + slot.lockout],
-    // keyed by end of the range
+    /// Flat list of intervals of lockouts of the form {voter, start, end}
+    /// ([`crate::consensus::progress_map::LockoutInterval`]).
     pub lockout_intervals: LockoutIntervals,
     pub my_latest_landed_vote: Option<Slot>,
 }
@@ -416,9 +413,13 @@ impl Tower {
             HashMap::with_capacity_and_hasher(total_slots, ahash::RandomState::default());
         let mut total_stake = 0;
 
-        // Tree of intervals of lockouts of the form [slot, slot + slot.lockout],
-        // keyed by end of the range
-        let mut lockout_intervals = LockoutIntervals::new();
+        let total_votes = vote_accounts
+            .values()
+            .filter(|(voted_stake, _)| *voted_stake != 0)
+            .map(|(_, account)| account.vote_state_view().votes_len())
+            .sum();
+        // Flat list of intervals of lockouts of the form {voter, start, end}.
+        let mut lockout_intervals = LockoutIntervals::with_capacity(total_votes);
         let mut my_latest_landed_vote = None;
         for (&key, (voted_stake, account)) in vote_accounts.iter() {
             let voted_stake = *voted_stake;
@@ -427,12 +428,11 @@ impl Tower {
             }
             trace!("{vote_account_pubkey} {key} with stake {voted_stake}");
             let mut vote_state = TowerVoteState::from(account.vote_state_view());
-            for vote in &vote_state.votes {
-                lockout_intervals
-                    .entry(vote.last_locked_out_slot())
-                    .or_default()
-                    .push((vote.slot(), key));
-            }
+            lockout_intervals.extend(vote_state.votes.iter().map(|v| LockoutInterval {
+                start: v.slot(),
+                end: v.last_locked_out_slot(),
+                voter: key,
+            }));
 
             if key == *vote_account_pubkey {
                 my_latest_landed_vote = vote_state.nth_recent_lockout(0).map(|l| l.slot());
@@ -1176,37 +1176,41 @@ impl Tower {
                 .fork_stats
                 .lockout_intervals;
             // Find any locked out intervals for vote accounts in this bank with
-            // `lockout_interval_end` >= `last_vote`, which implies they are locked out at
-            // `last_vote` on another fork.
-            for (_lockout_interval_end, intervals_keyed_by_end) in
-                lockout_intervals.range((Included(last_voted_slot), Unbounded))
+            // `lockout_interval_end` >= `last_vote`, which implies that the most recent tower slot is locked out
+            // at `last_vote` on another fork. We also consider the remaining tower slots, however these older slots
+            // could be from before the fork so we must filter by their ancestry.
+            for LockoutInterval {
+                start: lockout_interval_start,
+                voter: vote_account_pubkey,
+                ..
+            } in lockout_intervals
+                .iter()
+                .filter(|interval| interval.end >= last_voted_slot)
             {
-                for (lockout_interval_start, vote_account_pubkey) in intervals_keyed_by_end {
-                    if locked_out_vote_accounts.contains(vote_account_pubkey) {
-                        continue;
-                    }
+                if locked_out_vote_accounts.contains(vote_account_pubkey) {
+                    continue;
+                }
 
-                    // Only count lockouts on slots that are:
-                    // 1) Not ancestors of `last_vote`, meaning being on different fork
-                    // 2) Not from before the current root as we can't determine if
-                    // anything before the root was an ancestor of `last_vote` or not
-                    if !last_vote_ancestors.contains(lockout_interval_start) && {
-                        // Given a `lockout_interval_start` < root that appears in a
-                        // bank for a `candidate_slot`, it must be that `lockout_interval_start`
-                        // is an ancestor of the current root, because `candidate_slot` is a
-                        // descendant of the current root
-                        *lockout_interval_start > root
-                    } {
-                        let stake = epoch_vote_accounts
-                            .get(vote_account_pubkey)
-                            .map(|(stake, _)| *stake)
-                            .unwrap_or(0);
-                        locked_out_stake += stake;
-                        if (locked_out_stake as f64 / total_stake as f64) > SWITCH_FORK_THRESHOLD {
-                            return SwitchForkDecision::SwitchProof(switch_proof);
-                        }
-                        locked_out_vote_accounts.insert(vote_account_pubkey);
+                // Only count lockouts on slots that are:
+                // 1) Not ancestors of `last_vote`, meaning being on different fork
+                // 2) Not from before the current root as we can't determine if
+                // anything before the root was an ancestor of `last_vote` or not
+                if !last_vote_ancestors.contains(lockout_interval_start) && {
+                    // Given a `lockout_interval_start` < root that appears in a
+                    // bank for a `candidate_slot`, it must be that `lockout_interval_start`
+                    // is an ancestor of the current root, because `candidate_slot` is a
+                    // descendant of the current root
+                    *lockout_interval_start > root
+                } {
+                    let stake = epoch_vote_accounts
+                        .get(vote_account_pubkey)
+                        .map(|(stake, _)| *stake)
+                        .unwrap_or(0);
+                    locked_out_stake += stake;
+                    if (locked_out_stake as f64 / total_stake as f64) > SWITCH_FORK_THRESHOLD {
+                        return SwitchForkDecision::SwitchProof(switch_proof);
                     }
+                    locked_out_vote_accounts.insert(vote_account_pubkey);
                 }
             }
         }

core/src/consensus/progress_map.rs

@@ -12,7 +12,7 @@ use {
     solana_runtime::{bank::Bank, bank_forks::BankForks},
     solana_vote::vote_account::VoteAccountsHashMap,
     std::{
-        collections::{BTreeMap, HashMap, HashSet},
+        collections::{HashMap, HashSet},
         sync::{Arc, RwLock},
         time::Instant,
     },
@@ -20,7 +20,15 @@ use {
 
 type VotedSlot = Slot;
 type ExpirationSlot = Slot;
-pub type LockoutIntervals = BTreeMap<ExpirationSlot, Vec<(VotedSlot, Pubkey)>>;
+
+#[derive(Clone, Copy, Debug)]
+pub struct LockoutInterval {
+    pub voter: Pubkey,
+    pub start: VotedSlot,
+    pub end: ExpirationSlot,
+}
+
+pub type LockoutIntervals = Vec<LockoutInterval>;
 
 #[derive(Debug)]
 pub struct ValidatorStakeInfo {

core/src/vote_simulator.rs

@@ -7,7 +7,7 @@ use {
             fork_choice::{select_vote_and_reset_forks, SelectVoteAndResetForkResult},
             heaviest_subtree_fork_choice::HeaviestSubtreeForkChoice,
             latest_validator_votes_for_frozen_banks::LatestValidatorVotesForFrozenBanks,
-            progress_map::{ForkProgress, ProgressMap},
+            progress_map::{ForkProgress, LockoutInterval, ProgressMap},
             tower_vote_state::TowerVoteState,
             Tower,
         },
@@ -286,9 +286,11 @@ impl VoteSimulator {
             .or_insert_with(|| ForkProgress::new(Hash::default(), None, None, 0, 0))
             .fork_stats
             .lockout_intervals
-            .entry(lockout_interval.1)
-            .or_default()
-            .push((lockout_interval.0, *vote_account_pubkey));
+            .push(LockoutInterval {
+                start: lockout_interval.0,
+                end: lockout_interval.1,
+                voter: *vote_account_pubkey,
+            });
     }
 
     pub fn clear_lockout_intervals(&mut self, slot: Slot) {

vote/src/vote_state_view.rs

@@ -133,6 +133,11 @@ impl VoteStateView {
         })
     }
 
+    #[inline]
+    pub fn votes_len(&self) -> usize {
+        self.votes_view().len()
+    }
+
     pub fn last_lockout(&self) -> Option<Lockout> {
         self.votes_view().last().map(|item| {
             Lockout::new_with_confirmation_count(item.slot(), item.confirmation_count())