Skip to main content

solana_runtime/
stakes.rs

1//! Stakes serve as a cache of stake and vote accounts to derive
2//! node stakes
3#[cfg(feature = "frozen-abi")]
4use solana_frozen_abi::stable_abi::{context::SequenceLenMax, sample_collection_sized};
5use {
6    crate::{
7        alpenglow_epoch_type::RewardEpochDelegatedStakes, stake_account,
8        stake_history::StakeHistory,
9    },
10    imbl::HashMap as ImblHashMap,
11    log::error,
12    num_derive::ToPrimitive,
13    rayon::{ThreadPool, prelude::*},
14    solana_account::{AccountSharedData, ReadableAccount},
15    solana_accounts_db::utils::create_account_shared_data,
16    solana_clock::Epoch,
17    solana_leader_schedule::SlotLeader,
18    solana_pubkey::Pubkey,
19    solana_stake_history::StakeHistoryGetEntry,
20    solana_stake_interface::{
21        program as stake_program,
22        state::{Delegation, StakeActivationStatus},
23    },
24    solana_vote::vote_account::{VoteAccount, VoteAccounts, VoteAccountsHashMap},
25    solana_vote_interface::state::VoteStateVersions,
26    std::{
27        collections::HashMap,
28        sync::{Arc, RwLock, RwLockReadGuard},
29    },
30    thiserror::Error,
31    wincode::{SchemaWrite, containers::FromIntoIterator, len::BincodeLen},
32};
33#[cfg(feature = "dev-context-only-utils")]
34use {
35    qualifier_attr::{field_qualifiers, qualifiers},
36    solana_stake_interface::state::Stake,
37};
38
39mod serde_stakes;
40#[cfg_attr(feature = "dev-context-only-utils", qualifiers(pub))]
41pub(crate) use serde_stakes::DeserializableDelegationStakes;
42pub use serde_stakes::SerdeStakesToStakeFormat;
43#[derive(Debug, Error)]
44pub enum Error {
45    #[error("Invalid delegation: {0}")]
46    InvalidDelegation(Pubkey),
47    #[error(transparent)]
48    InvalidStakeAccount(#[from] stake_account::Error),
49    #[error("Stake account not found: {0}")]
50    StakeAccountNotFound(Pubkey),
51    #[error("Vote account mismatch: {0}")]
52    VoteAccountMismatch(Pubkey),
53    #[error("Vote account not cached: {0}")]
54    VoteAccountNotCached(Pubkey),
55    #[error("Vote account not found: {0}")]
56    VoteAccountNotFound(Pubkey),
57}
58
59#[derive(Debug, Clone, PartialEq, Eq, ToPrimitive)]
60pub enum InvalidCacheEntryReason {
61    Missing,
62    BadState,
63    WrongOwner,
64}
65
66type StakeAccount = stake_account::StakeAccount<Delegation>;
67pub(crate) type DelegatedStakes = ImblHashMap<Pubkey, u64>;
68
69/// Maximum number of inert delegations evicted from StakesCache at epoch
70/// boundaries. This number is chosen to cover normal deactivation pressure
71/// while not imposing undue cost in event of a backlog.
72const MAX_INERT_STAKES_REMOVED_PER_EPOCH: usize = 10_000;
73
74/// StakeStateV2::Stake, a stake with a delegation.
75const STAKE_STATE_STAKE_TAG: [u8; 4] = 2u32.to_le_bytes();
76
77fn stakes_cache_account_shared_data(account: &AccountSharedData) -> AccountSharedData {
78    if account.capacity() > account.data().len() {
79        create_account_shared_data(account)
80    } else {
81        account.clone()
82    }
83}
84
85/// Results of [`Stakes::calculate_activated_stake`] at the epoch boundary.
86pub(crate) struct EpochBoundaryStakes {
87    pub(crate) stake_history: StakeHistory,
88    pub(crate) vote_accounts: VoteAccounts,
89    pub(crate) delegated_stakes: DelegatedStakes,
90    pub(crate) reward_epoch_delegated_stakes: RewardEpochDelegatedStakes,
91    pub(crate) inert_stake_delegations: Vec<Pubkey>,
92}
93
94#[derive(Default, Debug)]
95pub(crate) struct StakesCache(RwLock<Stakes<StakeAccount>>);
96
97impl StakesCache {
98    pub(crate) fn new(stakes: Stakes<StakeAccount>) -> Self {
99        Self(RwLock::new(stakes))
100    }
101
102    pub(crate) fn stakes(&self) -> RwLockReadGuard<'_, Stakes<StakeAccount>> {
103        self.0.read().unwrap()
104    }
105
106    pub(crate) fn check_and_store(
107        &self,
108        pubkey: &Pubkey,
109        account: &AccountSharedData,
110        new_rate_activation_epoch: Option<Epoch>,
111        in_epoch_rewards_period: bool,
112        remove_inactive_stakes: bool,
113    ) {
114        // TODO: If the account is already cached as a vote or stake account
115        // but the owner changes, then this needs to evict the account from
116        // the cache. see:
117        // https://github.com/solana-labs/solana/pull/24200#discussion_r849935444
118        let owner = account.owner();
119        // Zero lamport accounts are not stored in accounts-db
120        // and so should be removed from cache as well.
121        if account.lamports() == 0 {
122            if solana_vote_program::check_id(owner) {
123                let _old_vote_account = {
124                    let mut stakes = self.0.write().unwrap();
125                    stakes.remove_vote_account(pubkey)
126                };
127            } else if stake_program::check_id(owner) {
128                let mut stakes = self.0.write().unwrap();
129                stakes.remove_stake_delegation(pubkey, new_rate_activation_epoch);
130            }
131            return;
132        }
133        debug_assert_ne!(account.lamports(), 0u64);
134        if solana_vote_program::check_id(owner) {
135            if VoteStateVersions::is_correct_size_and_initialized(account.data()) {
136                match VoteAccount::try_from(stakes_cache_account_shared_data(account)) {
137                    Ok(vote_account) => {
138                        // drop the old account after releasing the lock
139                        let _old_vote_account = {
140                            let mut stakes = self.0.write().unwrap();
141                            stakes.upsert_vote_account(pubkey, vote_account)
142                        };
143                    }
144                    Err(_) => {
145                        // drop the old account after releasing the lock
146                        let _old_vote_account = {
147                            let mut stakes = self.0.write().unwrap();
148                            stakes.remove_vote_account(pubkey)
149                        };
150                    }
151                }
152            } else {
153                // drop the old account after releasing the lock
154                let _old_vote_account = {
155                    let mut stakes = self.0.write().unwrap();
156                    stakes.remove_vote_account(pubkey)
157                };
158            };
159        } else if stake_program::check_id(owner) {
160            if account.data().starts_with(&STAKE_STATE_STAKE_TAG)
161                && let Ok(stake_account) =
162                    StakeAccount::try_from(stakes_cache_account_shared_data(account))
163            {
164                let mut stakes = self.0.write().unwrap();
165                stakes.upsert_stake_delegation(
166                    *pubkey,
167                    stake_account,
168                    new_rate_activation_epoch,
169                    in_epoch_rewards_period,
170                    remove_inactive_stakes,
171                );
172            } else {
173                let mut stakes = self.0.write().unwrap();
174                stakes.remove_stake_delegation(pubkey, new_rate_activation_epoch);
175            }
176        }
177    }
178
179    pub(crate) fn activate_epoch(
180        &self,
181        next_epoch: Epoch,
182        stake_history: StakeHistory,
183        vote_accounts: VoteAccounts,
184        delegated_stakes: DelegatedStakes,
185        inert_stake_delegations: &[Pubkey],
186    ) {
187        let mut stakes = self.0.write().unwrap();
188        stakes.activate_epoch(
189            next_epoch,
190            stake_history,
191            vote_accounts,
192            delegated_stakes,
193            inert_stake_delegations,
194        )
195    }
196
197    pub(crate) fn refresh_delegated_stakes(&self, new_rate_activation_epoch: Option<Epoch>) {
198        let mut stakes = self.0.write().unwrap();
199        stakes.refresh_delegated_stakes(new_rate_activation_epoch);
200    }
201
202    pub(crate) fn remove_inert_stake_delegations(
203        &self,
204        new_rate_activation_epoch: Option<Epoch>,
205        in_epoch_rewards_period: bool,
206    ) {
207        let mut stakes = self.0.write().unwrap();
208        stakes.remove_inert_stake_delegations(new_rate_activation_epoch, in_epoch_rewards_period);
209    }
210}
211
212/// The generic type T is either Delegation or StakeAccount.
213/// [`Stakes<Delegation>`] is equivalent to the old code and is used for backward
214/// compatibility in [`crate::bank::BankFieldsToDeserialize`].
215/// But banks cache [`Stakes<StakeAccount>`] which includes the entire stake
216/// account and StakeStateV2 deserialized from the account. Doing so, will remove
217/// the need to load the stake account from accounts-db when working with
218/// stake-delegations.
219#[cfg_attr(feature = "frozen-abi", derive(StableAbi, StableAbiSample))]
220#[derive(Default, Clone, PartialEq, Debug, SchemaWrite)]
221#[cfg_attr(
222    feature = "dev-context-only-utils",
223    field_qualifiers(
224        vote_accounts(pub),
225        stake_delegations(pub),
226        delegated_stakes(pub),
227        unused(pub),
228        epoch(pub),
229        stake_history(pub),
230    )
231)]
232pub struct Stakes<T: Clone> {
233    /// vote accounts
234    vote_accounts: VoteAccounts,
235
236    /// stake_delegations
237    #[cfg_attr(
238        feature = "frozen-abi",
239        stable_abi_sample(with = "sample_collection_sized(rng, SequenceLenMax(1))")
240    )]
241    #[wincode(with = "FromIntoIterator<ImblHashMap<Pubkey, T>, BincodeLen>")]
242    stake_delegations: ImblHashMap<Pubkey, T>,
243
244    /// current effective stake delegated to each vote account pubkey
245    #[cfg_attr(feature = "frozen-abi", stable_abi_sample(with = "Default::default()"))]
246    #[wincode(skip)]
247    delegated_stakes: DelegatedStakes,
248
249    /// unused
250    unused: u64,
251
252    /// current epoch, used to calculate current stake
253    epoch: Epoch,
254
255    /// history of staking levels
256    stake_history: StakeHistory,
257}
258
259impl<T: Clone> Stakes<T> {
260    pub fn new(vote_accounts: VoteAccounts, epoch: Epoch) -> Stakes<T> {
261        Stakes {
262            vote_accounts,
263            epoch,
264            stake_delegations: ImblHashMap::new(),
265            delegated_stakes: DelegatedStakes::default(),
266            unused: 0,
267            stake_history: StakeHistory::default(),
268        }
269    }
270
271    pub fn clone_and_filter_for_vat(
272        &self,
273        max_vote_accounts: usize,
274        minimum_vote_account_balance: u64,
275        block_revenue_sharing: bool,
276    ) -> Stakes<T> {
277        Self::new(
278            self.vote_accounts.clone_and_filter_for_vat(
279                max_vote_accounts,
280                minimum_vote_account_balance,
281                block_revenue_sharing,
282            ),
283            self.epoch,
284        )
285    }
286
287    pub fn vote_accounts(&self) -> &VoteAccounts {
288        &self.vote_accounts
289    }
290
291    pub(crate) fn staked_nodes(&self) -> Arc<HashMap<Pubkey, u64>> {
292        self.vote_accounts.staked_nodes()
293    }
294
295    /// Destructure self and return the fields needed by EpochStakes
296    pub(crate) fn into_epoch_stakes_fields(self) -> (Epoch, VoteAccounts, StakeHistory) {
297        let Self {
298            vote_accounts,
299            stake_delegations: _,
300            delegated_stakes: _,
301            unused: _,
302            epoch,
303            stake_history,
304        } = self;
305        (epoch, vote_accounts, stake_history)
306    }
307}
308
309impl Stakes<StakeAccount> {
310    pub(crate) fn new_from_accounts_for_genesis<'a, T: ReadableAccount + 'a>(
311        new_rate_activation_epoch: Option<Epoch>,
312        accounts: impl IntoIterator<Item = (&'a Pubkey, &'a T)>,
313    ) -> Self {
314        let stake_history = StakeHistory::default();
315        let mut vote_accounts = VoteAccountsHashMap::default();
316        let mut delegated_stakes = DelegatedStakes::default();
317        let mut stake_delegations = ImblHashMap::new();
318        let epoch = 0;
319
320        for (pubkey, account) in accounts {
321            if account.lamports() == 0 {
322                continue;
323            }
324
325            if solana_vote_program::check_id(account.owner()) {
326                if VoteStateVersions::is_correct_size_and_initialized(account.data())
327                    && let Ok(vote_account) =
328                        VoteAccount::try_from(create_account_shared_data(account))
329                {
330                    vote_accounts.insert(*pubkey, (0, vote_account));
331                }
332            } else if stake_program::check_id(account.owner())
333                && let Ok(stake_account) =
334                    StakeAccount::try_from(create_account_shared_data(account))
335            {
336                let delegation = stake_account.delegation();
337                let stake = delegation.stake_v2(epoch, &stake_history, new_rate_activation_epoch);
338                if stake != 0 {
339                    *delegated_stakes.entry(delegation.voter_pubkey).or_default() += stake;
340                }
341                stake_delegations.insert(*pubkey, stake_account);
342            }
343        }
344
345        let mut vote_accounts = VoteAccounts::from(Arc::new(vote_accounts));
346        for (vote_pubkey, stake) in &delegated_stakes {
347            vote_accounts.add_stake(vote_pubkey, *stake);
348        }
349
350        Self {
351            vote_accounts,
352            stake_delegations,
353            delegated_stakes,
354            unused: 0,
355            epoch,
356            stake_history,
357        }
358    }
359
360    /// Creates a Stake<StakeAccount> from DeserializableDelegationStakes by loading the
361    /// full account state for respective stake pubkeys. get_account function
362    /// should return the account at the respective slot where stakes where
363    /// cached.
364    #[cfg_attr(feature = "dev-context-only-utils", qualifiers(pub))]
365    pub(crate) fn load_from_deserialized_delegations<F>(
366        stakes: DeserializableDelegationStakes,
367        get_account: F,
368    ) -> Result<Self, Error>
369    where
370        F: Fn(&Pubkey) -> Option<AccountSharedData> + Sync,
371    {
372        let stake_delegations = stakes
373            .stake_delegations
374            .into_par_iter()
375            // We use fold/reduce to aggregate the results, which does a bit more work than calling
376            // collect()/collect_vec_list() and then imbl::HashMap::from_iter(collected.into_iter()),
377            // but it does it in background threads, so effectively it's faster.
378            .try_fold(ImblHashMap::new, |mut map, (pubkey, delegation)| {
379                let Some(stake_account) = get_account(&pubkey) else {
380                    return Err(Error::StakeAccountNotFound(pubkey));
381                };
382
383                // Assert that all valid vote-accounts referenced in stake delegations are already
384                // contained in `stakes.vote_account`.
385                let voter_pubkey = &delegation.voter_pubkey;
386                if stakes.vote_accounts.get(voter_pubkey).is_none()
387                    && let Some(account) = get_account(voter_pubkey)
388                    && VoteStateVersions::is_correct_size_and_initialized(account.data())
389                    && VoteAccount::try_from(account.clone()).is_ok()
390                {
391                    error!("vote account not cached: {voter_pubkey}, {account:?}");
392                    return Err(Error::VoteAccountNotCached(*voter_pubkey));
393                }
394
395                let stake_account = StakeAccount::try_from(stake_account)?;
396                // Sanity check that the delegation is consistent with what is
397                // stored in the account.
398                if stake_account.delegation() == &delegation {
399                    map.insert(pubkey, stake_account);
400                    Ok(map)
401                } else {
402                    Err(Error::InvalidDelegation(pubkey))
403                }
404            })
405            .try_reduce(ImblHashMap::new, |a, b| Ok(a.union(b)))?;
406
407        // Assert that cached vote accounts are consistent with accounts-db.
408        //
409        // This currently includes ~5500 accounts, parallelizing brings minor
410        // (sub 2s) improvements.
411        for (pubkey, vote_account) in stakes.vote_accounts.iter() {
412            let Some(account) = get_account(pubkey) else {
413                return Err(Error::VoteAccountNotFound(*pubkey));
414            };
415            let vote_account = vote_account.account();
416            if vote_account != &account {
417                error!("vote account mismatch: {pubkey}, {vote_account:?}, {account:?}");
418                return Err(Error::VoteAccountMismatch(*pubkey));
419            }
420        }
421
422        Ok(Self {
423            vote_accounts: stakes.vote_accounts.clone(),
424            stake_delegations,
425            delegated_stakes: DelegatedStakes::default(),
426            unused: stakes.unused,
427            epoch: stakes.epoch,
428            stake_history: stakes.stake_history,
429        })
430    }
431
432    #[cfg(feature = "dev-context-only-utils")]
433    pub fn new_for_tests(
434        epoch: Epoch,
435        vote_accounts: VoteAccounts,
436        stake_delegations: ImblHashMap<Pubkey, StakeAccount>,
437    ) -> Self {
438        let stake_history = StakeHistory::default();
439        let delegated_stakes =
440            Self::calculate_delegated_stakes(&stake_delegations, epoch, &stake_history, None);
441        Self {
442            vote_accounts,
443            stake_delegations,
444            delegated_stakes,
445            unused: 0,
446            epoch,
447            stake_history,
448        }
449    }
450
451    pub(crate) fn history(&self) -> &StakeHistory {
452        &self.stake_history
453    }
454
455    pub(crate) fn calculate_activated_stake(
456        &self,
457        next_epoch: Epoch,
458        thread_pool: &ThreadPool,
459        new_rate_activation_epoch: Option<Epoch>,
460        stake_delegations: &[(&Pubkey, &StakeAccount)],
461        remove_inactive_stakes: bool,
462    ) -> EpochBoundaryStakes {
463        // Wrap up the prev epoch by adding new stake history entry for the
464        // prev epoch.
465        let identity = || {
466            (
467                StakeActivationStatus::default(),
468                HashMap::default(),
469                Vec::default(),
470            )
471        };
472        let (stake_history_entry, effective_delegated_stakes, inert_stake_delegations) =
473            thread_pool.install(|| {
474                stake_delegations
475                    .par_iter()
476                    .fold(
477                        identity,
478                        |(acc, mut delegated_stakes, mut inert_stake_delegations),
479                         (stake_pubkey, stake_account)| {
480                            let delegation = stake_account.delegation();
481                            let activation_status = delegation.stake_activating_and_deactivating_v2(
482                                self.epoch,
483                                &self.stake_history,
484                                new_rate_activation_epoch,
485                            );
486                            // A delegation with neither effective nor activating
487                            // stake can never contribute stake again, so it can
488                            // be dropped from the cache.
489                            if remove_inactive_stakes
490                                && activation_status.effective == 0
491                                && activation_status.activating == 0
492                            {
493                                inert_stake_delegations.push(**stake_pubkey);
494                            }
495                            *delegated_stakes.entry(delegation.voter_pubkey).or_default() +=
496                                activation_status.effective;
497                            (acc + activation_status, delegated_stakes, inert_stake_delegations)
498                        },
499                    )
500                    .reduce(
501                        identity,
502                        |(activation_status_a, delegated_stakes_a, inert_a),
503                         (activation_status_b, delegated_stakes_b, inert_b)| {
504                            (
505                                activation_status_a + activation_status_b,
506                                merge_delegated_stakes(delegated_stakes_a, delegated_stakes_b),
507                                merge_inert_stake_delegations(inert_a, inert_b),
508                            )
509                        },
510                    )
511            });
512        let mut stake_history = self.stake_history.clone();
513        stake_history.add(self.epoch, stake_history_entry);
514        // Refresh the stake distribution of vote accounts for the next epoch,
515        // using new stake history.
516        let (vote_accounts, delegated_stakes) = refresh_vote_accounts(
517            thread_pool,
518            next_epoch,
519            &self.vote_accounts,
520            stake_delegations,
521            &stake_history,
522            new_rate_activation_epoch,
523        );
524        let reward_epoch_delegated_stakes = RewardEpochDelegatedStakes {
525            epoch: self.epoch,
526            delegated_stakes: effective_delegated_stakes,
527        };
528        EpochBoundaryStakes {
529            stake_history,
530            vote_accounts,
531            delegated_stakes,
532            reward_epoch_delegated_stakes,
533            inert_stake_delegations,
534        }
535    }
536
537    pub(crate) fn activate_epoch(
538        &mut self,
539        next_epoch: Epoch,
540        stake_history: StakeHistory,
541        vote_accounts: VoteAccounts,
542        delegated_stakes: DelegatedStakes,
543        inert_stake_delegations: &[Pubkey],
544    ) {
545        self.epoch = next_epoch;
546        self.stake_history = stake_history;
547        self.vote_accounts = vote_accounts;
548        self.delegated_stakes = delegated_stakes;
549        // With SIMD-0599, inert delegations contribute no stake or points and
550        // earn no rewards, so we can drop them here with no consensus effect.
551        // We bound the number of removals for performance reasons. If we have
552        // more inactives than MAX_INERT_STAKES_REMOVED_PER_EPOCH, this will
553        // remove a random subset on each validator due to how imbl is seeded,
554        // again with no consensus effect.
555        //
556        // Before SIMD-0599, inert_stake_delegations is empty, so this is a no-op.
557        for stake_pubkey in inert_stake_delegations
558            .iter()
559            .take(MAX_INERT_STAKES_REMOVED_PER_EPOCH)
560        {
561            self.stake_delegations.remove(stake_pubkey);
562        }
563    }
564
565    fn calculate_delegated_stakes(
566        stake_delegations: &ImblHashMap<Pubkey, StakeAccount>,
567        epoch: Epoch,
568        stake_history: &StakeHistory,
569        new_rate_activation_epoch: Option<Epoch>,
570    ) -> DelegatedStakes {
571        let mut delegated_stakes = DelegatedStakes::new();
572        for stake_account in stake_delegations.values() {
573            let delegation = stake_account.delegation();
574            let stake = delegation.stake_v2(epoch, stake_history, new_rate_activation_epoch);
575            if stake != 0 {
576                *delegated_stakes.entry(delegation.voter_pubkey).or_default() += stake;
577            }
578        }
579        delegated_stakes
580    }
581
582    fn refresh_delegated_stakes(&mut self, new_rate_activation_epoch: Option<Epoch>) {
583        self.delegated_stakes = Self::calculate_delegated_stakes(
584            &self.stake_delegations,
585            self.epoch,
586            &self.stake_history,
587            new_rate_activation_epoch,
588        );
589    }
590
591    /// Removes inactive delegations, used when restoring from snapshot. As in
592    /// `upsert_stake_delegation`, during the rewards period we keep stakes that
593    /// were active last epoch in case they are owed rewards.
594    fn remove_inert_stake_delegations(
595        &mut self,
596        new_rate_activation_epoch: Option<Epoch>,
597        in_epoch_rewards_period: bool,
598    ) {
599        let inert_stake_delegations: Vec<Pubkey> = self
600            .stake_delegations
601            .iter()
602            .filter(|(_stake_pubkey, stake_account)| {
603                let delegation = stake_account.delegation();
604
605                // Inactive this epoch.
606                let is_inactive_now = is_delegation_inert(
607                    delegation,
608                    self.epoch,
609                    &self.stake_history,
610                    new_rate_activation_epoch,
611                );
612
613                // Either we are out of rewards (and prior epoch doesn't
614                // matter), or it was also inactive last epoch.
615                is_inactive_now
616                    && (!in_epoch_rewards_period
617                        || is_delegation_inert(
618                            delegation,
619                            self.epoch.saturating_sub(1),
620                            &self.stake_history,
621                            new_rate_activation_epoch,
622                        ))
623            })
624            .map(|(stake_pubkey, _stake_account)| *stake_pubkey)
625            .collect();
626        // Inert delegations contribute no stake to `delegated_stakes` or
627        // `vote_accounts`, so no further bookkeeping is needed.
628        for stake_pubkey in &inert_stake_delegations {
629            self.stake_delegations.remove(stake_pubkey);
630        }
631    }
632
633    fn add_delegated_stake(&mut self, voter_pubkey: Pubkey, stake: u64) {
634        if stake == 0 {
635            return;
636        }
637        *self.delegated_stakes.entry(voter_pubkey).or_default() += stake;
638    }
639
640    fn sub_delegated_stake(&mut self, voter_pubkey: &Pubkey, stake: u64) {
641        if stake == 0 {
642            return;
643        }
644        let current_stake = self
645            .delegated_stakes
646            .get_mut(voter_pubkey)
647            .expect("subtraction from missing delegated stake");
648        *current_stake = current_stake
649            .checked_sub(stake)
650            .expect("subtraction value exceeds delegated stake");
651        if *current_stake == 0 {
652            self.delegated_stakes.remove(voter_pubkey);
653        }
654    }
655
656    fn remove_vote_account(&mut self, vote_pubkey: &Pubkey) -> Option<VoteAccount> {
657        self.vote_accounts.remove(vote_pubkey).map(|(_, a)| a)
658    }
659
660    fn remove_stake_delegation(
661        &mut self,
662        stake_pubkey: &Pubkey,
663        new_rate_activation_epoch: Option<Epoch>,
664    ) {
665        if let Some(stake_account) = self.stake_delegations.remove(stake_pubkey) {
666            let removed_delegation = stake_account.delegation();
667            let removed_stake = removed_delegation.stake_v2(
668                self.epoch,
669                &self.stake_history,
670                new_rate_activation_epoch,
671            );
672            self.sub_delegated_stake(&removed_delegation.voter_pubkey, removed_stake);
673            self.vote_accounts
674                .sub_stake(&removed_delegation.voter_pubkey, removed_stake);
675        }
676    }
677
678    fn upsert_vote_account(
679        &mut self,
680        vote_pubkey: &Pubkey,
681        vote_account: VoteAccount,
682    ) -> Option<VoteAccount> {
683        debug_assert_ne!(vote_account.lamports(), 0u64);
684
685        let calculate_delegated_stake = || {
686            self.delegated_stakes
687                .get(vote_pubkey)
688                .copied()
689                .unwrap_or_default()
690        };
691        self.vote_accounts
692            .insert(*vote_pubkey, vote_account, calculate_delegated_stake)
693    }
694
695    fn upsert_stake_delegation(
696        &mut self,
697        stake_pubkey: Pubkey,
698        stake_account: StakeAccount,
699        new_rate_activation_epoch: Option<Epoch>,
700        in_epoch_rewards_period: bool,
701        remove_inactive_stakes: bool,
702    ) {
703        debug_assert_ne!(stake_account.lamports(), 0u64);
704        let delegation = stake_account.delegation();
705        let voter_pubkey = delegation.voter_pubkey;
706        let activation_status = delegation.stake_activating_and_deactivating_v2(
707            self.epoch,
708            &self.stake_history,
709            new_rate_activation_epoch,
710        );
711
712        // An inactive delegation contributes no stake, so it can generally be
713        // excluded from cache. However, if we are in the rewards period, we
714        // must also check that it was not active in the previous epoch.
715        if remove_inactive_stakes
716            && activation_status.effective == 0
717            && activation_status.activating == 0
718        {
719            let may_be_awaiting_rewards = in_epoch_rewards_period
720                && !is_delegation_inert(
721                    delegation,
722                    self.epoch.saturating_sub(1),
723                    &self.stake_history,
724                    new_rate_activation_epoch,
725                );
726            if !may_be_awaiting_rewards {
727                self.remove_stake_delegation(&stake_pubkey, new_rate_activation_epoch);
728                return;
729            }
730        }
731        let stake = activation_status.effective;
732        match self.stake_delegations.insert(stake_pubkey, stake_account) {
733            None => {
734                self.add_delegated_stake(voter_pubkey, stake);
735                self.vote_accounts.add_stake(&voter_pubkey, stake);
736            }
737            Some(old_stake_account) => {
738                let old_delegation = old_stake_account.delegation();
739                let old_voter_pubkey = old_delegation.voter_pubkey;
740                let old_stake = old_delegation.stake_v2(
741                    self.epoch,
742                    &self.stake_history,
743                    new_rate_activation_epoch,
744                );
745                if voter_pubkey != old_voter_pubkey || stake != old_stake {
746                    self.sub_delegated_stake(&old_voter_pubkey, old_stake);
747                    self.add_delegated_stake(voter_pubkey, stake);
748                    self.vote_accounts.sub_stake(&old_voter_pubkey, old_stake);
749                    self.vote_accounts.add_stake(&voter_pubkey, stake);
750                }
751            }
752        }
753    }
754
755    /// Returns a reference to the map of stake delegations.
756    ///
757    /// # Performance
758    ///
759    /// `[imbl::HashMap]` is a [hash array mapped trie (HAMT)][hamt], which means
760    /// that inserts, deletions and lookups are average-case O(1) and
761    /// worst-case O(log n). However, the performance of iterations is poor due
762    /// to depth-first traversal and jumps. Currently it's also impossible to
763    /// iterate over it with [`rayon`].
764    ///
765    /// [hamt]: https://en.wikipedia.org/wiki/Hash_array_mapped_trie
766    pub(crate) fn stake_delegations(&self) -> &ImblHashMap<Pubkey, StakeAccount> {
767        &self.stake_delegations
768    }
769
770    /// Collects stake delegations into a vector, which then can be used for
771    /// parallel iteration with [`rayon`].
772    ///
773    /// # Performance
774    ///
775    /// The execution of this method takes ~200ms and it collects elements of
776    /// the [`imbl::HashMap`], which is a [hash array mapped trie (HAMT)][hamt],
777    /// so that operation involves a depth-first traversal with jumps. However,
778    /// it's still a reasonable tradeoff if the caller iterates over these
779    /// elements.
780    ///
781    /// [hamt]: https://en.wikipedia.org/wiki/Hash_array_mapped_trie
782    pub(crate) fn stake_delegations_vec(&self) -> Vec<(&Pubkey, &StakeAccount)> {
783        self.stake_delegations.iter().collect()
784    }
785
786    pub(crate) fn highest_staked_node(&self) -> Option<SlotLeader> {
787        let (vote_address, vote_account) = self.vote_accounts.find_max_by_delegated_stake()?;
788        Some(SlotLeader {
789            id: *vote_account.node_pubkey(),
790            vote_address: *vote_address,
791        })
792    }
793}
794
795/// Macro to generate `From<Stakes<From>> for Stakes<To>` impls.
796#[cfg(feature = "dev-context-only-utils")]
797macro_rules! impl_stake_format_conversion {
798    ($from:ty, $to:ty, |$binding:ident| $expr:expr) => {
799        /// This conversion is memory intensive so should only be used in development contexts.
800        impl From<Stakes<$from>> for Stakes<$to> {
801            fn from(stakes: Stakes<$from>) -> Self {
802                let Stakes {
803                    vote_accounts,
804                    stake_delegations,
805                    delegated_stakes: _,
806                    unused,
807                    epoch,
808                    stake_history,
809                } = stakes;
810                let stake_delegations = stake_delegations
811                    .into_iter()
812                    .map(|(pubkey, $binding)| (pubkey, $expr))
813                    .collect();
814                Self {
815                    vote_accounts,
816                    stake_delegations,
817                    delegated_stakes: DelegatedStakes::default(),
818                    unused,
819                    epoch,
820                    stake_history,
821                }
822            }
823        }
824    };
825}
826
827#[cfg(feature = "dev-context-only-utils")]
828impl_stake_format_conversion!(StakeAccount, Delegation, |sa| *sa.delegation());
829
830#[cfg(feature = "dev-context-only-utils")]
831impl_stake_format_conversion!(StakeAccount, Stake, |sa| *sa.stake());
832
833#[cfg(feature = "dev-context-only-utils")]
834impl_stake_format_conversion!(Stake, Delegation, |stake| stake.delegation);
835
836pub(crate) fn is_delegation_inert<T: StakeHistoryGetEntry>(
837    delegation: &Delegation,
838    epoch: Epoch,
839    stake_history: &T,
840    new_rate_activation_epoch: Option<Epoch>,
841) -> bool {
842    let activation_status = delegation.stake_activating_and_deactivating_v2(
843        epoch,
844        stake_history,
845        new_rate_activation_epoch,
846    );
847    activation_status.effective == 0 && activation_status.activating == 0
848}
849
850fn merge_delegated_stakes(
851    mut stakes: HashMap</*voter:*/ Pubkey, /*stake:*/ u64>,
852    other: HashMap</*voter:*/ Pubkey, /*stake:*/ u64>,
853) -> HashMap</*voter:*/ Pubkey, /*stake:*/ u64> {
854    if stakes.len() < other.len() {
855        return merge_delegated_stakes(other, stakes);
856    }
857    for (pubkey, stake) in other {
858        *stakes.entry(pubkey).or_default() += stake;
859    }
860    stakes
861}
862
863fn merge_inert_stake_delegations(mut inert: Vec<Pubkey>, other: Vec<Pubkey>) -> Vec<Pubkey> {
864    if inert.len() < other.len() {
865        return merge_inert_stake_delegations(other, inert);
866    }
867    inert.extend(other);
868    inert
869}
870
871fn refresh_vote_accounts(
872    thread_pool: &ThreadPool,
873    epoch: Epoch,
874    vote_accounts: &VoteAccounts,
875    stake_delegations: &[(&Pubkey, &StakeAccount)],
876    stake_history: &StakeHistory,
877    new_rate_activation_epoch: Option<Epoch>,
878) -> (VoteAccounts, DelegatedStakes) {
879    fn merge(mut stakes: DelegatedStakes, other: DelegatedStakes) -> DelegatedStakes {
880        if stakes.len() < other.len() {
881            return merge(other, stakes);
882        }
883        for (pubkey, stake) in other {
884            *stakes.entry(pubkey).or_default() += stake;
885        }
886        stakes
887    }
888    let delegated_stakes = thread_pool.install(|| {
889        stake_delegations
890            .par_iter()
891            .fold(
892                DelegatedStakes::default,
893                |mut delegated_stakes, (_stake_pubkey, stake_account)| {
894                    let delegation = stake_account.delegation();
895                    let stake =
896                        delegation.stake_v2(epoch, stake_history, new_rate_activation_epoch);
897                    if stake != 0 {
898                        *delegated_stakes.entry(delegation.voter_pubkey).or_default() += stake;
899                    }
900                    delegated_stakes
901                },
902            )
903            .reduce(DelegatedStakes::default, merge)
904    });
905    let vote_accounts = vote_accounts
906        .iter()
907        .map(|(&vote_pubkey, vote_account)| {
908            let delegated_stake = delegated_stakes
909                .get(&vote_pubkey)
910                .copied()
911                .unwrap_or_default();
912            (vote_pubkey, (delegated_stake, vote_account.clone()))
913        })
914        .collect();
915    (vote_accounts, delegated_stakes)
916}
917
918#[cfg(test)]
919pub(crate) mod tests {
920    use {
921        super::*,
922        crate::stake_utils,
923        rayon::ThreadPoolBuilder,
924        solana_account::{WritableAccount, state_traits::StateMutWincode as _},
925        solana_pubkey::Pubkey,
926        solana_rent::Rent,
927        solana_stake_interface::{self as stake, state::StakeStateV2},
928        solana_vote_interface::state::{BLS_PUBLIC_KEY_COMPRESSED_SIZE, VoteStateV4},
929        solana_vote_program::vote_state,
930        test_case::test_matrix,
931    };
932
933    impl Stakes<Delegation> {
934        /// Convert deserialized stakes into runtime stakes representation
935        pub(crate) fn from_deserialized(stakes: DeserializableDelegationStakes) -> Self {
936            Self {
937                vote_accounts: stakes.vote_accounts,
938                stake_delegations: ImblHashMap::from_iter(stakes.stake_delegations),
939                delegated_stakes: DelegatedStakes::default(),
940                unused: stakes.unused,
941                epoch: stakes.epoch,
942                stake_history: stakes.stake_history,
943            }
944        }
945    }
946
947    //  set up some dummies for a staked node
948    pub(crate) fn create_staked_node_accounts_with_activation_and_deactivation_epochs(
949        stake: u64,
950        rent: &Rent,
951        activation_epoch: Epoch,
952        deactivation_epoch: Epoch,
953    ) -> ((Pubkey, AccountSharedData), (Pubkey, AccountSharedData)) {
954        let vote_pubkey = solana_pubkey::new_rand();
955        let node_pubkey = solana_pubkey::new_rand();
956        let vote_account = vote_state::create_v4_account_with_authorized(
957            &node_pubkey,
958            &vote_pubkey,
959            [0u8; BLS_PUBLIC_KEY_COMPRESSED_SIZE],
960            &vote_pubkey,
961            0,
962            &vote_pubkey,
963            0,
964            &node_pubkey,
965            rent.minimum_balance(VoteStateV4::size_of()),
966        );
967        let stake_pubkey = solana_pubkey::new_rand();
968        (
969            (vote_pubkey, vote_account),
970            (
971                stake_pubkey,
972                create_stake_account_with_activation_and_deactivation_epochs(
973                    stake,
974                    &vote_pubkey,
975                    &stake_pubkey,
976                    rent,
977                    activation_epoch,
978                    deactivation_epoch,
979                ),
980            ),
981        )
982    }
983
984    //  set up some dummies for a staked node     ((     vote      )  (     stake     ))
985    pub(crate) fn create_staked_node_accounts(
986        stake: u64,
987        rent: &Rent,
988    ) -> ((Pubkey, AccountSharedData), (Pubkey, AccountSharedData)) {
989        create_staked_node_accounts_with_activation_and_deactivation_epochs(
990            stake,
991            rent,
992            Epoch::MAX,
993            Epoch::MAX,
994        )
995    }
996
997    //   add stake to a vote_pubkey
998    pub(crate) fn create_stake_account_with_activation_and_deactivation_epochs(
999        stake: u64,
1000        vote_pubkey: &Pubkey,
1001        stake_pubkey: &Pubkey,
1002        rent: &Rent,
1003        activation_epoch: Epoch,
1004        deactivation_epoch: Epoch,
1005    ) -> AccountSharedData {
1006        let node_pubkey = solana_pubkey::new_rand();
1007        let lamports = rent.minimum_balance(StakeStateV2::size_of()) + stake;
1008        stake_utils::create_stake_account_with_activation_and_deactivation_epochs(
1009            stake_pubkey,
1010            vote_pubkey,
1011            &vote_state::create_v4_account_with_authorized(
1012                &node_pubkey,
1013                vote_pubkey,
1014                [0u8; BLS_PUBLIC_KEY_COMPRESSED_SIZE],
1015                vote_pubkey,
1016                0,
1017                vote_pubkey,
1018                0,
1019                &node_pubkey,
1020                1,
1021            ),
1022            rent,
1023            lamports,
1024            activation_epoch,
1025            deactivation_epoch,
1026        )
1027    }
1028
1029    //   add stake to a vote_pubkey                               (   stake    )
1030    pub(crate) fn create_stake_account(
1031        stake: u64,
1032        vote_pubkey: &Pubkey,
1033        stake_pubkey: &Pubkey,
1034        rent: &Rent,
1035    ) -> AccountSharedData {
1036        let node_pubkey = solana_pubkey::new_rand();
1037        let lamports = rent.minimum_balance(StakeStateV2::size_of()) + stake;
1038        stake_utils::create_stake_account(
1039            stake_pubkey,
1040            vote_pubkey,
1041            &vote_state::create_v4_account_with_authorized(
1042                &node_pubkey,
1043                vote_pubkey,
1044                [0u8; BLS_PUBLIC_KEY_COMPRESSED_SIZE],
1045                vote_pubkey,
1046                0,
1047                vote_pubkey,
1048                0,
1049                &node_pubkey,
1050                1,
1051            ),
1052            rent,
1053            lamports,
1054        )
1055    }
1056
1057    #[test]
1058    fn test_stakes_basic() {
1059        for i in 0..4 {
1060            let stakes_cache = StakesCache::new(Stakes {
1061                epoch: i,
1062                ..Stakes::default()
1063            });
1064            let rent = Rent::default();
1065
1066            let ((vote_pubkey, vote_account), (stake_pubkey, mut stake_account)) =
1067                create_staked_node_accounts(10, &rent);
1068
1069            stakes_cache.check_and_store(&vote_pubkey, &vote_account, None, false, false);
1070            stakes_cache.check_and_store(&stake_pubkey, &stake_account, None, false, false);
1071            let stake_state: StakeStateV2 = stake_account.state().unwrap();
1072            let stake = stake_state.stake().unwrap();
1073            {
1074                let stakes = stakes_cache.stakes();
1075                let vote_accounts = stakes.vote_accounts();
1076                assert!(vote_accounts.get(&vote_pubkey).is_some());
1077                let expected_stake = stake.stake_v2(i, &StakeHistory::default(), None);
1078                assert_eq!(
1079                    vote_accounts.get_delegated_stake(&vote_pubkey),
1080                    expected_stake
1081                );
1082            }
1083
1084            stake_account.set_lamports(42);
1085            stakes_cache.check_and_store(&stake_pubkey, &stake_account, None, false, false);
1086            {
1087                let stakes = stakes_cache.stakes();
1088                let vote_accounts = stakes.vote_accounts();
1089                assert!(vote_accounts.get(&vote_pubkey).is_some());
1090                let expected_stake = stake.stake_v2(i, &StakeHistory::default(), None);
1091                assert_eq!(
1092                    vote_accounts.get_delegated_stake(&vote_pubkey),
1093                    expected_stake
1094                ); // stays old stake, because only 10 is activated
1095            }
1096
1097            // activate more
1098            let mut stake_account =
1099                create_stake_account(42, &vote_pubkey, &solana_pubkey::new_rand(), &rent);
1100            stakes_cache.check_and_store(&stake_pubkey, &stake_account, None, false, false);
1101            let stake_state: StakeStateV2 = stake_account.state().unwrap();
1102            let stake = stake_state.stake().unwrap();
1103            {
1104                let stakes = stakes_cache.stakes();
1105                let vote_accounts = stakes.vote_accounts();
1106                assert!(vote_accounts.get(&vote_pubkey).is_some());
1107                let expected_stake = stake.stake_v2(i, &StakeHistory::default(), None);
1108                assert_eq!(
1109                    vote_accounts.get_delegated_stake(&vote_pubkey),
1110                    expected_stake
1111                ); // now stake of 42 is activated
1112            }
1113
1114            stake_account.set_lamports(0);
1115            stakes_cache.check_and_store(&stake_pubkey, &stake_account, None, false, false);
1116            {
1117                let stakes = stakes_cache.stakes();
1118                let vote_accounts = stakes.vote_accounts();
1119                assert!(vote_accounts.get(&vote_pubkey).is_some());
1120                assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 0);
1121            }
1122        }
1123    }
1124
1125    #[test]
1126    fn test_stakes_highest() {
1127        let stakes_cache = StakesCache::default();
1128        let rent = Rent::default();
1129
1130        assert_eq!(stakes_cache.stakes().highest_staked_node(), None);
1131
1132        let ((vote_pubkey, vote_account), (stake_pubkey, stake_account)) =
1133            create_staked_node_accounts(10, &rent);
1134
1135        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None, false, false);
1136        stakes_cache.check_and_store(&stake_pubkey, &stake_account, None, false, false);
1137
1138        let ((vote11_pubkey, vote11_account), (stake11_pubkey, stake11_account)) =
1139            create_staked_node_accounts(20, &rent);
1140
1141        stakes_cache.check_and_store(&vote11_pubkey, &vote11_account, None, false, false);
1142        stakes_cache.check_and_store(&stake11_pubkey, &stake11_account, None, false, false);
1143
1144        let vote11_node_pubkey = VoteStateV4::deserialize(vote11_account.data(), &vote11_pubkey)
1145            .unwrap()
1146            .node_pubkey;
1147
1148        let highest_staked_node = stakes_cache.stakes().highest_staked_node();
1149        assert_eq!(
1150            highest_staked_node,
1151            Some(SlotLeader {
1152                id: vote11_node_pubkey,
1153                vote_address: vote11_pubkey,
1154            })
1155        );
1156    }
1157
1158    #[test]
1159    fn test_stakes_vote_account_disappear_reappear() {
1160        let stakes_cache = StakesCache::new(Stakes {
1161            epoch: 4,
1162            ..Stakes::default()
1163        });
1164        let rent = Rent::default();
1165
1166        let ((vote_pubkey, mut vote_account), (stake_pubkey, stake_account)) =
1167            create_staked_node_accounts(10, &rent);
1168
1169        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None, false, false);
1170        stakes_cache.check_and_store(&stake_pubkey, &stake_account, None, false, false);
1171
1172        {
1173            let stakes = stakes_cache.stakes();
1174            let vote_accounts = stakes.vote_accounts();
1175            assert!(vote_accounts.get(&vote_pubkey).is_some());
1176            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 10);
1177        }
1178
1179        vote_account.set_lamports(0);
1180        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None, false, false);
1181
1182        {
1183            let stakes = stakes_cache.stakes();
1184            let vote_accounts = stakes.vote_accounts();
1185            assert!(vote_accounts.get(&vote_pubkey).is_none());
1186            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 0);
1187        }
1188
1189        vote_account.set_lamports(1);
1190        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None, false, false);
1191
1192        {
1193            let stakes = stakes_cache.stakes();
1194            let vote_accounts = stakes.vote_accounts();
1195            assert!(vote_accounts.get(&vote_pubkey).is_some());
1196            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 10);
1197        }
1198
1199        // Vote account too big
1200        let cache_data = vote_account.data().to_vec();
1201        let mut pushed = vote_account.data().to_vec();
1202        pushed.push(0);
1203        vote_account.set_data_from_slice(&pushed);
1204        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None, false, false);
1205
1206        {
1207            let stakes = stakes_cache.stakes();
1208            let vote_accounts = stakes.vote_accounts();
1209            assert!(vote_accounts.get(&vote_pubkey).is_none());
1210            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 0);
1211        }
1212
1213        // Vote account uninitialized
1214        vote_account.set_data_from_slice(&vec![0; VoteStateV4::size_of()]);
1215        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None, false, false);
1216
1217        {
1218            let stakes = stakes_cache.stakes();
1219            let vote_accounts = stakes.vote_accounts();
1220            assert!(vote_accounts.get(&vote_pubkey).is_none());
1221            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 0);
1222        }
1223
1224        vote_account.set_data_from_slice(&cache_data);
1225        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None, false, false);
1226
1227        {
1228            let stakes = stakes_cache.stakes();
1229            let vote_accounts = stakes.vote_accounts();
1230            assert!(vote_accounts.get(&vote_pubkey).is_some());
1231            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 10);
1232        }
1233    }
1234
1235    #[test]
1236    fn test_stakes_change_delegate() {
1237        let stakes_cache = StakesCache::new(Stakes {
1238            epoch: 4,
1239            ..Stakes::default()
1240        });
1241        let rent = Rent::default();
1242
1243        let ((vote_pubkey, vote_account), (stake_pubkey, stake_account)) =
1244            create_staked_node_accounts(10, &rent);
1245
1246        let ((vote_pubkey2, vote_account2), (_stake_pubkey2, stake_account2)) =
1247            create_staked_node_accounts(10, &rent);
1248
1249        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None, false, false);
1250        stakes_cache.check_and_store(&vote_pubkey2, &vote_account2, None, false, false);
1251
1252        // delegates to vote_pubkey
1253        stakes_cache.check_and_store(&stake_pubkey, &stake_account, None, false, false);
1254
1255        let stake_state: StakeStateV2 = stake_account.state().unwrap();
1256        let stake = stake_state.stake().unwrap();
1257
1258        {
1259            let stakes = stakes_cache.stakes();
1260            let vote_accounts = stakes.vote_accounts();
1261            assert!(vote_accounts.get(&vote_pubkey).is_some());
1262            let expected_stake = stake.stake_v2(stakes.epoch, &stakes.stake_history, None);
1263            assert_eq!(
1264                vote_accounts.get_delegated_stake(&vote_pubkey),
1265                expected_stake
1266            );
1267            assert!(vote_accounts.get(&vote_pubkey2).is_some());
1268            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey2), 0);
1269        }
1270
1271        // delegates to vote_pubkey2
1272        stakes_cache.check_and_store(&stake_pubkey, &stake_account2, None, false, false);
1273
1274        {
1275            let stakes = stakes_cache.stakes();
1276            let vote_accounts = stakes.vote_accounts();
1277            assert!(vote_accounts.get(&vote_pubkey).is_some());
1278            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 0);
1279            assert!(vote_accounts.get(&vote_pubkey2).is_some());
1280            let expected_stake = stake.stake_v2(stakes.epoch, &stakes.stake_history, None);
1281            assert_eq!(
1282                vote_accounts.get_delegated_stake(&vote_pubkey2),
1283                expected_stake
1284            );
1285        }
1286    }
1287    #[test]
1288    fn test_stakes_multiple_stakers() {
1289        let stakes_cache = StakesCache::new(Stakes {
1290            epoch: 4,
1291            ..Stakes::default()
1292        });
1293        let rent = Rent::default();
1294
1295        let ((vote_pubkey, vote_account), (stake_pubkey, stake_account)) =
1296            create_staked_node_accounts(10, &rent);
1297
1298        let stake_pubkey2 = solana_pubkey::new_rand();
1299        let stake_account2 = create_stake_account(10, &vote_pubkey, &stake_pubkey2, &rent);
1300
1301        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None, false, false);
1302
1303        // delegates to vote_pubkey
1304        stakes_cache.check_and_store(&stake_pubkey, &stake_account, None, false, false);
1305        stakes_cache.check_and_store(&stake_pubkey2, &stake_account2, None, false, false);
1306
1307        {
1308            let stakes = stakes_cache.stakes();
1309            let vote_accounts = stakes.vote_accounts();
1310            assert!(vote_accounts.get(&vote_pubkey).is_some());
1311            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 20);
1312        }
1313    }
1314
1315    #[test]
1316    fn test_activate_epoch() {
1317        let stakes_cache = StakesCache::default();
1318        let rent = Rent::default();
1319
1320        let ((vote_pubkey, vote_account), (stake_pubkey, stake_account)) =
1321            create_staked_node_accounts(10, &rent);
1322
1323        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None, false, false);
1324        stakes_cache.check_and_store(&stake_pubkey, &stake_account, None, false, false);
1325        let stake_state: StakeStateV2 = stake_account.state().unwrap();
1326        let stake = stake_state.stake().unwrap();
1327
1328        let initial_expected_stake = {
1329            let stakes = stakes_cache.stakes();
1330            stake.stake_v2(stakes.epoch, &stakes.stake_history, None)
1331        };
1332        {
1333            let stakes = stakes_cache.stakes();
1334            let vote_accounts = stakes.vote_accounts();
1335            assert_eq!(
1336                vote_accounts.get_delegated_stake(&vote_pubkey),
1337                initial_expected_stake
1338            );
1339        }
1340        let thread_pool = ThreadPoolBuilder::new().num_threads(1).build().unwrap();
1341        let next_epoch = 3;
1342        let EpochBoundaryStakes {
1343            stake_history,
1344            vote_accounts,
1345            delegated_stakes,
1346            reward_epoch_delegated_stakes: effective_delegated_stakes,
1347            inert_stake_delegations,
1348        } = {
1349            let stakes = stakes_cache.stakes();
1350            let stake_delegations = stakes.stake_delegations_vec();
1351            stakes.calculate_activated_stake(
1352                next_epoch,
1353                &thread_pool,
1354                None,
1355                &stake_delegations,
1356                true,
1357            )
1358        };
1359        assert_eq!(
1360            effective_delegated_stakes
1361                .delegated_stakes
1362                .get(&vote_pubkey)
1363                .copied(),
1364            Some(initial_expected_stake)
1365        );
1366        stakes_cache.activate_epoch(
1367            next_epoch,
1368            stake_history,
1369            vote_accounts,
1370            delegated_stakes,
1371            &inert_stake_delegations,
1372        );
1373        {
1374            let stakes = stakes_cache.stakes();
1375            let vote_accounts = stakes.vote_accounts();
1376            let expected_stake = stake.stake_v2(stakes.epoch, &stakes.stake_history, None);
1377            assert_eq!(
1378                vote_accounts.get_delegated_stake(&vote_pubkey),
1379                expected_stake
1380            );
1381        }
1382    }
1383
1384    #[test]
1385    fn test_stakes_not_delegate() {
1386        let stakes_cache = StakesCache::new(Stakes {
1387            epoch: 4,
1388            ..Stakes::default()
1389        });
1390        let rent = Rent::default();
1391
1392        let ((vote_pubkey, vote_account), (stake_pubkey, stake_account)) =
1393            create_staked_node_accounts(10, &rent);
1394
1395        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None, false, false);
1396        stakes_cache.check_and_store(&stake_pubkey, &stake_account, None, false, false);
1397
1398        {
1399            let stakes = stakes_cache.stakes();
1400            let vote_accounts = stakes.vote_accounts();
1401            assert!(vote_accounts.get(&vote_pubkey).is_some());
1402            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 10);
1403        }
1404
1405        // not a stake account, and whacks above entry
1406        stakes_cache.check_and_store(
1407            &stake_pubkey,
1408            &AccountSharedData::new(1, 0, &stake::program::id()),
1409            None,
1410            false,
1411            false,
1412        );
1413        {
1414            let stakes = stakes_cache.stakes();
1415            let vote_accounts = stakes.vote_accounts();
1416            assert!(vote_accounts.get(&vote_pubkey).is_some());
1417            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 0);
1418        }
1419    }
1420
1421    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
1422    enum InertTestState {
1423        Activating,
1424        Active,
1425        Deactivating,
1426        OldDeactivated,
1427        NewDeactivated,
1428        ActivatedAndDeactivated,
1429    }
1430
1431    impl InertTestState {
1432        fn delegation_epochs(&self, current: Epoch) -> (Epoch, Epoch) {
1433            match self {
1434                InertTestState::Activating => (current, Epoch::MAX),
1435                InertTestState::Active => (0, Epoch::MAX),
1436                InertTestState::Deactivating => (0, current),
1437                InertTestState::OldDeactivated => (0, current - 2),
1438                InertTestState::NewDeactivated => (0, current - 1),
1439                InertTestState::ActivatedAndDeactivated => (current, current),
1440            }
1441        }
1442
1443        fn is_droppable(&self, in_rewards: bool) -> bool {
1444            match self {
1445                InertTestState::Activating
1446                | InertTestState::Active
1447                | InertTestState::Deactivating => false,
1448                InertTestState::OldDeactivated | InertTestState::ActivatedAndDeactivated => true,
1449                InertTestState::NewDeactivated => !in_rewards,
1450            }
1451        }
1452    }
1453
1454    #[test_matrix(
1455        [InertTestState::Activating, InertTestState::Active, InertTestState::Deactivating,
1456         InertTestState::OldDeactivated, InertTestState::NewDeactivated, InertTestState::ActivatedAndDeactivated],
1457        [false, true],
1458        [false, true]
1459    )]
1460    fn test_inert_delegation_is_removed_from_cache(
1461        case: InertTestState,
1462        in_epoch_rewards_period: bool,
1463        pre_insert_stake: bool,
1464    ) {
1465        const CURRENT_EPOCH: Epoch = 4;
1466        let stakes_cache = StakesCache::new(Stakes {
1467            epoch: CURRENT_EPOCH,
1468            ..Stakes::default()
1469        });
1470
1471        let (activation_epoch, deactivation_epoch) = case.delegation_epochs(CURRENT_EPOCH);
1472        let ((vote_pubkey, vote_account), (stake_pubkey, mut stake_account)) =
1473            create_staked_node_accounts(10, &Rent::default());
1474        let StakeStateV2::Stake(meta, mut stake, flags) = stake_account.state().unwrap() else {
1475            unreachable!()
1476        };
1477        stake.delegation.activation_epoch = activation_epoch;
1478        stake.delegation.deactivation_epoch = deactivation_epoch;
1479        stake_account
1480            .set_state(&StakeStateV2::Stake(meta, stake, flags))
1481            .unwrap();
1482
1483        stakes_cache.check_and_store(
1484            &vote_pubkey,
1485            &vote_account,
1486            None,
1487            in_epoch_rewards_period,
1488            true,
1489        );
1490
1491        // Possibly insert, called with the feature off so the eviction check is bypassed
1492        if pre_insert_stake {
1493            stakes_cache.check_and_store(&stake_pubkey, &stake_account, None, false, false);
1494            assert!(
1495                stakes_cache
1496                    .stakes()
1497                    .stake_delegations()
1498                    .contains_key(&stake_pubkey)
1499            );
1500        }
1501
1502        // Check and store with SIMD-0599 on. If we hold the stake, it should drop from cache,
1503        // and if we don't, it should fail to insert, if the stake is in a droppable state.
1504        // Stakes in a cacheable state should be inserted regardless.
1505        stakes_cache.check_and_store(
1506            &stake_pubkey,
1507            &stake_account,
1508            None,
1509            in_epoch_rewards_period,
1510            true,
1511        );
1512
1513        assert_eq!(
1514            !stakes_cache
1515                .stakes()
1516                .stake_delegations()
1517                .contains_key(&stake_pubkey),
1518            case.is_droppable(in_epoch_rewards_period)
1519        );
1520    }
1521
1522    #[test_matrix(
1523        [InertTestState::Activating, InertTestState::Active, InertTestState::Deactivating,
1524         InertTestState::OldDeactivated, InertTestState::NewDeactivated, InertTestState::ActivatedAndDeactivated],
1525        [false, true]
1526    )]
1527    fn test_inert_delegation_is_pruned_on_load(
1528        case: InertTestState,
1529        in_epoch_rewards_period: bool,
1530    ) {
1531        const CURRENT_EPOCH: Epoch = 4;
1532        let stakes_cache = StakesCache::new(Stakes {
1533            epoch: CURRENT_EPOCH,
1534            ..Stakes::default()
1535        });
1536
1537        let (activation_epoch, deactivation_epoch) = case.delegation_epochs(CURRENT_EPOCH);
1538        let ((vote_pubkey, vote_account), (stake_pubkey, mut stake_account)) =
1539            create_staked_node_accounts(10, &Rent::default());
1540        let StakeStateV2::Stake(meta, mut stake, flags) = stake_account.state().unwrap() else {
1541            unreachable!()
1542        };
1543        stake.delegation.activation_epoch = activation_epoch;
1544        stake.delegation.deactivation_epoch = deactivation_epoch;
1545        stake_account
1546            .set_state(&StakeStateV2::Stake(meta, stake, flags))
1547            .unwrap();
1548
1549        // Insert with the feature off, as a snapshot from before activation would
1550        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None, false, false);
1551        stakes_cache.check_and_store(&stake_pubkey, &stake_account, None, false, false);
1552        stakes_cache.remove_inert_stake_delegations(None, in_epoch_rewards_period);
1553
1554        let stakes = stakes_cache.stakes();
1555        assert_eq!(
1556            !stakes.stake_delegations().contains_key(&stake_pubkey),
1557            case.is_droppable(in_epoch_rewards_period)
1558        );
1559    }
1560}