Revert "Revert "Add more reporting for invalid stake cache members and prune them (backport #21654) (#21740)""

This reverts commit 8028f218a4.

runtime/src/bank.rs

@@ -67,6 +67,8 @@ use {
     dashmap::DashMap,
     itertools::Itertools,
     log::*,
+    num_derive::ToPrimitive,
+    num_traits::ToPrimitive,
     rayon::{
         iter::{IntoParallelIterator, IntoParallelRefIterator, ParallelIterator},
         ThreadPool, ThreadPoolBuilder,
@@ -1110,6 +1112,19 @@ struct VoteWithStakeDelegations {
     delegations: Vec<(Pubkey, (StakeState, AccountSharedData))>,
 }
 
+#[derive(Debug, Clone, PartialEq, ToPrimitive)]
+enum InvalidReason {
+    Missing,
+    BadState,
+    WrongOwner,
+}
+
+struct LoadVoteAndStakeAccountsResult {
+    vote_with_stake_delegations_map: DashMap<Pubkey, VoteWithStakeDelegations>,
+    invalid_stake_keys: DashMap<Pubkey, InvalidReason>,
+    invalid_vote_keys: DashMap<Pubkey, InvalidReason>,
+}
+
 #[derive(Debug, Default)]
 pub struct NewBankOptions {
     pub vote_only_bank: bool,
@@ -2073,9 +2088,11 @@ impl Bank {
         &self,
         thread_pool: &ThreadPool,
         reward_calc_tracer: Option<impl Fn(&RewardCalculationEvent) + Send + Sync>,
-    ) -> DashMap<Pubkey, VoteWithStakeDelegations> {
+    ) -> LoadVoteAndStakeAccountsResult {
         let stakes = self.stakes.read().unwrap();
-        let accounts = DashMap::with_capacity(stakes.vote_accounts().len());
+        let vote_with_stake_delegations_map = DashMap::with_capacity(stakes.vote_accounts().len());
+        let invalid_stake_keys: DashMap<Pubkey, InvalidReason> = DashMap::new();
+        let invalid_vote_keys: DashMap<Pubkey, InvalidReason> = DashMap::new();
 
         thread_pool.install(|| {
             stakes
@@ -2083,87 +2100,134 @@ impl Bank {
                 .par_iter()
                 .for_each(|(stake_pubkey, delegation)| {
                     let vote_pubkey = &delegation.voter_pubkey;
-                    let stake_account = match self.get_account_with_fixed_root(stake_pubkey) {
+                    if invalid_vote_keys.contains_key(vote_pubkey) {
-                        Some(stake_account) => stake_account,
+                        return;
-                        None => return,
+                    }
-                    };
 
-                    // fetch vote account from stakes cache if it hasn't been cached locally
+                    let stake_delegation = match self.get_account_with_fixed_root(stake_pubkey) {
-                    let fetched_vote_account = if !accounts.contains_key(vote_pubkey) {
+                        Some(stake_account) => {
-                        let vote_account = match self.get_account_with_fixed_root(vote_pubkey) {
+                            if stake_account.owner() != &solana_stake_program::id() {
-                            Some(vote_account) => vote_account,
+                                invalid_stake_keys.insert(*stake_pubkey, InvalidReason::WrongOwner);
-                            None => return,
+                                return;
-                        };
+                            }
 
-                        let vote_state: VoteState =
+                            match stake_account.state().ok() {
-                            match StateMut::<VoteStateVersions>::state(&vote_account) {
+                                Some(stake_state) => (*stake_pubkey, (stake_state, stake_account)),
-                                Ok(vote_state) => vote_state.convert_to_current(),
+                                None => {
-                                Err(err) => {
+                                    invalid_stake_keys
-                                    debug!(
+                                        .insert(*stake_pubkey, InvalidReason::BadState);
-                                        "failed to deserialize vote account {}: {}",
-                                        vote_pubkey, err
-                                    );
                                     return;
                                 }
-                            };
+                            }
-
+                        }
-                        Some((vote_state, vote_account))
+                        None => {
-                    } else {
+                            invalid_stake_keys.insert(*stake_pubkey, InvalidReason::Missing);
-                        None
+                            return;
+                        }
                     };
 
-                    let fetched_vote_account_owner = fetched_vote_account
+                    let mut vote_delegations = if let Some(vote_delegations) =
-                        .as_ref()
+                        vote_with_stake_delegations_map.get_mut(vote_pubkey)
-                        .map(|(_vote_state, vote_account)| vote_account.owner());
+                    {
+                        vote_delegations
+                    } else {
+                        let vote_account = match self.get_account_with_fixed_root(vote_pubkey) {
+                            Some(vote_account) => {
+                                if vote_account.owner() != &solana_vote_program::id() {
+                                    invalid_vote_keys
+                                        .insert(*vote_pubkey, InvalidReason::WrongOwner);
+                                    return;
+                                }
+                                vote_account
+                            }
+                            None => {
+                                invalid_vote_keys.insert(*vote_pubkey, InvalidReason::Missing);
+                                return;
+                            }
+                        };
+
+                        let vote_state = if let Ok(vote_state) =
+                            StateMut::<VoteStateVersions>::state(&vote_account)
+                        {
+                            vote_state.convert_to_current()
+                        } else {
+                            invalid_vote_keys.insert(*vote_pubkey, InvalidReason::BadState);
+                            return;
+                        };
+
+                        vote_with_stake_delegations_map
+                            .entry(*vote_pubkey)
+                            .or_insert_with(|| VoteWithStakeDelegations {
+                                vote_state: Arc::new(vote_state),
+                                vote_account,
+                                delegations: vec![],
+                            })
+                    };
 
                     if let Some(reward_calc_tracer) = reward_calc_tracer.as_ref() {
                         reward_calc_tracer(&RewardCalculationEvent::Staking(
                             stake_pubkey,
                             &InflationPointCalculationEvent::Delegation(
                                 *delegation,
-                                fetched_vote_account_owner
+                                solana_vote_program::id(),
-                                    .cloned()
-                                    .unwrap_or_else(solana_vote_program::id),
                             ),
                         ));
                     }
 
-                    // filter invalid delegation accounts
+                    vote_delegations.delegations.push(stake_delegation);
-                    if stake_account.owner() != &solana_stake_program::id()
-                        || (fetched_vote_account_owner.is_some()
-                            && fetched_vote_account_owner != Some(&solana_vote_program::id()))
-                    {
-                        datapoint_warn!(
-                            "bank-stake_delegation_accounts-invalid-account",
-                            ("slot", self.slot() as i64, i64),
-                            ("stake-address", format!("{:?}", stake_pubkey), String),
-                            ("vote-address", format!("{:?}", vote_pubkey), String),
-                        );
-                        return;
-                    }
-
-                    let stake_delegation = match stake_account.state().ok() {
-                        Some(stake_state) => (*stake_pubkey, (stake_state, stake_account)),
-                        None => return,
-                    };
-
-                    if let Some((vote_state, vote_account)) = fetched_vote_account {
-                        accounts
-                            .entry(*vote_pubkey)
-                            .or_insert_with(|| VoteWithStakeDelegations {
-                                vote_state: Arc::new(vote_state),
-                                vote_account,
-                                delegations: vec![],
-                            });
-                    }
-
-                    if let Some(mut stake_delegation_accounts) = accounts.get_mut(vote_pubkey) {
-                        stake_delegation_accounts.delegations.push(stake_delegation);
-                    }
                 });
         });
 
-        accounts
+        LoadVoteAndStakeAccountsResult {
+            vote_with_stake_delegations_map,
+            invalid_vote_keys,
+            invalid_stake_keys,
+        }
+    }
+
+    fn handle_invalid_stakes_cache_keys(
+        &self,
+        invalid_stake_keys: DashMap<Pubkey, InvalidReason>,
+        invalid_vote_keys: DashMap<Pubkey, InvalidReason>,
+    ) {
+        if invalid_stake_keys.is_empty() && invalid_vote_keys.is_empty() {
+            return;
+        }
+
+        // Prune invalid stake delegations and vote accounts that were
+        // not properly evicted in normal operation.
+        let mut maybe_stakes_cache = if self
+            .feature_set
+            .is_active(&feature_set::evict_invalid_stakes_cache_entries::id())
+        {
+            Some(self.stakes.write().unwrap())
+        } else {
+            None
+        };
+
+        for (stake_pubkey, reason) in invalid_stake_keys {
+            if let Some(stakes_cache) = maybe_stakes_cache.as_mut() {
+                stakes_cache.remove_stake_delegation(&stake_pubkey);
+            }
+            datapoint_warn!(
+                "bank-stake_delegation_accounts-invalid-account",
+                ("slot", self.slot() as i64, i64),
+                ("stake-address", format!("{:?}", stake_pubkey), String),
+                ("reason", reason.to_i64().unwrap_or_default(), i64),
+            );
+        }
+
+        for (vote_pubkey, reason) in invalid_vote_keys {
+            if let Some(stakes_cache) = maybe_stakes_cache.as_mut() {
+                stakes_cache.remove_vote_account(&vote_pubkey);
+            }
+            datapoint_warn!(
+                "bank-stake_delegation_accounts-invalid-account",
+                ("slot", self.slot() as i64, i64),
+                ("vote-address", format!("{:?}", vote_pubkey), String),
+                ("reason", reason.to_i64().unwrap_or_default(), i64),
+            );
+        }
     }
 
     /// iterate over all stakes, redeem vote credits for each stake we can
@@ -2177,13 +2241,22 @@ impl Bank {
         thread_pool: &ThreadPool,
     ) -> f64 {
         let stake_history = self.stakes.read().unwrap().history().clone();
-        let vote_and_stake_accounts = self.load_vote_and_stake_accounts_with_thread_pool(
+        let vote_with_stake_delegations_map = {
-            thread_pool,
+            let LoadVoteAndStakeAccountsResult {
-            reward_calc_tracer.as_ref(),
+                vote_with_stake_delegations_map,
-        );
+                invalid_stake_keys,
+                invalid_vote_keys,
+            } = self.load_vote_and_stake_accounts_with_thread_pool(
+                thread_pool,
+                reward_calc_tracer.as_ref(),
+            );
+
+            self.handle_invalid_stakes_cache_keys(invalid_stake_keys, invalid_vote_keys);
+            vote_with_stake_delegations_map
+        };
 
         let points: u128 = thread_pool.install(|| {
-            vote_and_stake_accounts
+            vote_with_stake_delegations_map
                 .par_iter()
                 .map(|entry| {
                     let VoteWithStakeDelegations {
@@ -2214,8 +2287,8 @@ impl Bank {
         // pay according to point value
         let point_value = PointValue { rewards, points };
         let vote_account_rewards: DashMap<Pubkey, (AccountSharedData, u8, u64, bool)> =
-            DashMap::with_capacity(vote_and_stake_accounts.len());
+            DashMap::with_capacity(vote_with_stake_delegations_map.len());
-        let stake_delegation_iterator = vote_and_stake_accounts.into_par_iter().flat_map(
+        let stake_delegation_iterator = vote_with_stake_delegations_map.into_par_iter().flat_map(
             |(
                 vote_pubkey,
                 VoteWithStakeDelegations {
@@ -7780,6 +7853,7 @@ pub(crate) mod tests {
         let thread_pool = ThreadPoolBuilder::new().num_threads(1).build().unwrap();
         let validator_points: u128 = bank0
             .load_vote_and_stake_accounts_with_thread_pool(&thread_pool, null_tracer())
+            .vote_with_stake_delegations_map
             .into_iter()
             .map(
                 |(
@@ -13507,8 +13581,9 @@ pub(crate) mod tests {
         );
         let bank = Arc::new(Bank::new(&genesis_config));
         let thread_pool = ThreadPoolBuilder::new().num_threads(1).build().unwrap();
-        let vote_and_stake_accounts =
+        let vote_and_stake_accounts = bank
-            bank.load_vote_and_stake_accounts_with_thread_pool(&thread_pool, null_tracer());
+            .load_vote_and_stake_accounts_with_thread_pool(&thread_pool, null_tracer())
+            .vote_with_stake_delegations_map;
         assert_eq!(vote_and_stake_accounts.len(), 2);
 
         let mut vote_account = bank
@@ -13548,8 +13623,9 @@ pub(crate) mod tests {
 
         // Accounts must be valid stake and vote accounts
         let thread_pool = ThreadPoolBuilder::new().num_threads(1).build().unwrap();
-        let vote_and_stake_accounts =
+        let vote_and_stake_accounts = bank
-            bank.load_vote_and_stake_accounts_with_thread_pool(&thread_pool, null_tracer());
+            .load_vote_and_stake_accounts_with_thread_pool(&thread_pool, null_tracer())
+            .vote_with_stake_delegations_map;
         assert_eq!(vote_and_stake_accounts.len(), 0);
     }
 

runtime/src/stakes.rs

@@ -148,6 +148,18 @@ impl Stakes {
                 && account.data().len() >= std::mem::size_of::<StakeState>()
     }
 
+    pub fn remove_vote_account(&mut self, vote_pubkey: &Pubkey) {
+        self.vote_accounts.remove(vote_pubkey);
+    }
+
+    pub fn remove_stake_delegation(&mut self, stake_pubkey: &Pubkey) {
+        if let Some(removed_delegation) = self.stake_delegations.remove(stake_pubkey) {
+            let removed_stake = removed_delegation.stake(self.epoch, Some(&self.stake_history));
+            self.vote_accounts
+                .sub_stake(&removed_delegation.voter_pubkey, removed_stake);
+        }
+    }
+
     pub fn store(
         &mut self,
         pubkey: &Pubkey,

sdk/src/feature_set.rs

@@ -281,6 +281,10 @@ pub mod reject_non_rent_exempt_vote_withdraws {
     solana_sdk::declare_id!("7txXZZD6Um59YoLMF7XUNimbMjsqsWhc7g2EniiTrmp1");
 }
 
+pub mod evict_invalid_stakes_cache_entries {
+    solana_sdk::declare_id!("EMX9Q7TVFAmQ9V1CggAkhMzhXSg8ECp7fHrWQX2G1chf");
+}
+
 lazy_static! {
     /// Map of feature identifiers to user-visible description
     pub static ref FEATURE_NAMES: HashMap<Pubkey, &'static str> = [
@@ -351,6 +355,7 @@ lazy_static! {
         (reject_all_elf_rw::id(), "reject all read-write data in program elfs"),
         (spl_token_v3_3_0_release::id(), "spl-token v3.3.0 release"),
         (reject_non_rent_exempt_vote_withdraws::id(), "fail vote withdraw instructions which leave the account non-rent-exempt"),
+        (evict_invalid_stakes_cache_entries::id(), "evict invalid stakes cache entries on epoch boundaries"),
         /*************** ADD NEW FEATURES HERE ***************/
     ]
     .iter()