solana_runtime/
stakes.rs

1//! Stakes serve as a cache of stake and vote accounts to derive
2//! node stakes
3use {
4    crate::{stake_account, stake_history::StakeHistory},
5    im::HashMap as ImHashMap,
6    log::error,
7    num_derive::ToPrimitive,
8    rayon::{prelude::*, ThreadPool},
9    solana_sdk::{
10        account::{AccountSharedData, ReadableAccount},
11        clock::Epoch,
12        pubkey::Pubkey,
13        stake::state::{Delegation, StakeActivationStatus},
14        vote::state::VoteStateVersions,
15    },
16    solana_stake_program::stake_state::Stake,
17    solana_vote::vote_account::{VoteAccount, VoteAccounts},
18    std::{
19        collections::HashMap,
20        ops::Add,
21        sync::{Arc, RwLock, RwLockReadGuard},
22    },
23    thiserror::Error,
24};
25
26mod serde_stakes;
27pub(crate) use serde_stakes::serde_stakes_to_delegation_format;
28pub use serde_stakes::SerdeStakesToStakeFormat;
29
30#[derive(Debug, Error)]
31pub enum Error {
32    #[error("Invalid delegation: {0}")]
33    InvalidDelegation(Pubkey),
34    #[error(transparent)]
35    InvalidStakeAccount(#[from] stake_account::Error),
36    #[error("Stake account not found: {0}")]
37    StakeAccountNotFound(Pubkey),
38    #[error("Vote account mismatch: {0}")]
39    VoteAccountMismatch(Pubkey),
40    #[error("Vote account not cached: {0}")]
41    VoteAccountNotCached(Pubkey),
42    #[error("Vote account not found: {0}")]
43    VoteAccountNotFound(Pubkey),
44}
45
46#[derive(Debug, Clone, PartialEq, Eq, ToPrimitive)]
47pub enum InvalidCacheEntryReason {
48    Missing,
49    BadState,
50    WrongOwner,
51}
52
53type StakeAccount = stake_account::StakeAccount<Delegation>;
54
55#[cfg_attr(feature = "frozen-abi", derive(AbiExample))]
56#[derive(Default, Debug)]
57pub(crate) struct StakesCache(RwLock<Stakes<StakeAccount>>);
58
59impl StakesCache {
60    pub(crate) fn new(stakes: Stakes<StakeAccount>) -> Self {
61        Self(RwLock::new(stakes))
62    }
63
64    pub(crate) fn stakes(&self) -> RwLockReadGuard<Stakes<StakeAccount>> {
65        self.0.read().unwrap()
66    }
67
68    pub(crate) fn check_and_store(
69        &self,
70        pubkey: &Pubkey,
71        account: &impl ReadableAccount,
72        new_rate_activation_epoch: Option<Epoch>,
73    ) {
74        // TODO: If the account is already cached as a vote or stake account
75        // but the owner changes, then this needs to evict the account from
76        // the cache. see:
77        // https://github.com/solana-labs/solana/pull/24200#discussion_r849935444
78        let owner = account.owner();
79        // Zero lamport accounts are not stored in accounts-db
80        // and so should be removed from cache as well.
81        if account.lamports() == 0 {
82            if solana_vote_program::check_id(owner) {
83                let _old_vote_account = {
84                    let mut stakes = self.0.write().unwrap();
85                    stakes.remove_vote_account(pubkey)
86                };
87            } else if solana_stake_program::check_id(owner) {
88                let mut stakes = self.0.write().unwrap();
89                stakes.remove_stake_delegation(pubkey, new_rate_activation_epoch);
90            }
91            return;
92        }
93        debug_assert_ne!(account.lamports(), 0u64);
94        if solana_vote_program::check_id(owner) {
95            if VoteStateVersions::is_correct_size_and_initialized(account.data()) {
96                match VoteAccount::try_from(account.to_account_shared_data()) {
97                    Ok(vote_account) => {
98                        // drop the old account after releasing the lock
99                        let _old_vote_account = {
100                            let mut stakes = self.0.write().unwrap();
101                            stakes.upsert_vote_account(
102                                pubkey,
103                                vote_account,
104                                new_rate_activation_epoch,
105                            )
106                        };
107                    }
108                    Err(_) => {
109                        // drop the old account after releasing the lock
110                        let _old_vote_account = {
111                            let mut stakes = self.0.write().unwrap();
112                            stakes.remove_vote_account(pubkey)
113                        };
114                    }
115                }
116            } else {
117                // drop the old account after releasing the lock
118                let _old_vote_account = {
119                    let mut stakes = self.0.write().unwrap();
120                    stakes.remove_vote_account(pubkey)
121                };
122            };
123        } else if solana_stake_program::check_id(owner) {
124            match StakeAccount::try_from(account.to_account_shared_data()) {
125                Ok(stake_account) => {
126                    let mut stakes = self.0.write().unwrap();
127                    stakes.upsert_stake_delegation(
128                        *pubkey,
129                        stake_account,
130                        new_rate_activation_epoch,
131                    );
132                }
133                Err(_) => {
134                    let mut stakes = self.0.write().unwrap();
135                    stakes.remove_stake_delegation(pubkey, new_rate_activation_epoch);
136                }
137            }
138        }
139    }
140
141    pub(crate) fn activate_epoch(
142        &self,
143        next_epoch: Epoch,
144        thread_pool: &ThreadPool,
145        new_rate_activation_epoch: Option<Epoch>,
146    ) {
147        let mut stakes = self.0.write().unwrap();
148        stakes.activate_epoch(next_epoch, thread_pool, new_rate_activation_epoch)
149    }
150}
151
152/// The generic type T is either Delegation or StakeAccount.
153/// [`Stakes<Delegation>`] is equivalent to the old code and is used for backward
154/// compatibility in [`crate::bank::BankFieldsToDeserialize`].
155/// But banks cache [`Stakes<StakeAccount>`] which includes the entire stake
156/// account and StakeStateV2 deserialized from the account. Doing so, will remove
157/// the need to load the stake account from accounts-db when working with
158/// stake-delegations.
159#[cfg_attr(feature = "frozen-abi", derive(AbiExample))]
160#[derive(Default, Clone, PartialEq, Debug, Deserialize, Serialize)]
161pub struct Stakes<T: Clone> {
162    /// vote accounts
163    vote_accounts: VoteAccounts,
164
165    /// stake_delegations
166    stake_delegations: ImHashMap<Pubkey, T>,
167
168    /// unused
169    unused: u64,
170
171    /// current epoch, used to calculate current stake
172    epoch: Epoch,
173
174    /// history of staking levels
175    stake_history: StakeHistory,
176}
177
178// For backward compatibility, we can only serialize and deserialize
179// Stakes<Delegation> in the old `epoch_stakes` bank snapshot field. However,
180// Stakes<StakeAccount> entries are added to the bank's epoch stakes hashmap
181// when crossing epoch boundaries and Stakes<Stake> entries are added when
182// starting up from bank snapshots that have the new epoch stakes field. By
183// using this enum, the cost of converting all entries to Stakes<Delegation> is
184// put off until serializing new snapshots. This helps avoid bogging down epoch
185// boundaries and startup with the conversion overhead.
186#[cfg_attr(feature = "frozen-abi", derive(AbiExample))]
187#[derive(Debug, Clone)]
188pub enum StakesEnum {
189    Accounts(Stakes<StakeAccount>),
190    Delegations(Stakes<Delegation>),
191    Stakes(Stakes<Stake>),
192}
193
194impl<T: Clone> Stakes<T> {
195    pub fn vote_accounts(&self) -> &VoteAccounts {
196        &self.vote_accounts
197    }
198
199    pub(crate) fn staked_nodes(&self) -> Arc<HashMap<Pubkey, u64>> {
200        self.vote_accounts.staked_nodes()
201    }
202}
203
204impl Stakes<StakeAccount> {
205    /// Creates a Stake<StakeAccount> from Stake<Delegation> by loading the
206    /// full account state for respective stake pubkeys. get_account function
207    /// should return the account at the respective slot where stakes where
208    /// cached.
209    pub(crate) fn new<F>(stakes: &Stakes<Delegation>, get_account: F) -> Result<Self, Error>
210    where
211        F: Fn(&Pubkey) -> Option<AccountSharedData> + Sync,
212    {
213        let stake_delegations = stakes
214            .stake_delegations
215            .iter()
216            // im::HashMap doesn't support rayon so we manually build a temporary vector. Note this is
217            // what std HashMap::par_iter() does internally too.
218            .collect::<Vec<_>>()
219            .into_par_iter()
220            // We use fold/reduce to aggregate the results, which does a bit more work than calling
221            // collect()/collect_vec_list() and then im::HashMap::from_iter(collected.into_iter()),
222            // but it does it in background threads, so effectively it's faster.
223            .try_fold(ImHashMap::new, |mut map, (pubkey, delegation)| {
224                let Some(stake_account) = get_account(pubkey) else {
225                    return Err(Error::StakeAccountNotFound(*pubkey));
226                };
227
228                // Assert that all valid vote-accounts referenced in stake delegations are already
229                // contained in `stakes.vote_account`.
230                let voter_pubkey = &delegation.voter_pubkey;
231                if stakes.vote_accounts.get(voter_pubkey).is_none() {
232                    if let Some(account) = get_account(voter_pubkey) {
233                        if VoteStateVersions::is_correct_size_and_initialized(account.data())
234                            && VoteAccount::try_from(account.clone()).is_ok()
235                        {
236                            error!("vote account not cached: {voter_pubkey}, {account:?}");
237                            return Err(Error::VoteAccountNotCached(*voter_pubkey));
238                        }
239                    }
240                }
241
242                let stake_account = StakeAccount::try_from(stake_account)?;
243                // Sanity check that the delegation is consistent with what is
244                // stored in the account.
245                if stake_account.delegation() == delegation {
246                    map.insert(*pubkey, stake_account);
247                    Ok(map)
248                } else {
249                    Err(Error::InvalidDelegation(*pubkey))
250                }
251            })
252            .try_reduce(ImHashMap::new, |a, b| Ok(a.union(b)))?;
253
254        // Assert that cached vote accounts are consistent with accounts-db.
255        //
256        // This currently includes ~5500 accounts, parallelizing brings minor
257        // (sub 2s) improvements.
258        for (pubkey, vote_account) in stakes.vote_accounts.iter() {
259            let Some(account) = get_account(pubkey) else {
260                return Err(Error::VoteAccountNotFound(*pubkey));
261            };
262            let vote_account = vote_account.account();
263            if vote_account != &account {
264                error!("vote account mismatch: {pubkey}, {vote_account:?}, {account:?}");
265                return Err(Error::VoteAccountMismatch(*pubkey));
266            }
267        }
268
269        Ok(Self {
270            vote_accounts: stakes.vote_accounts.clone(),
271            stake_delegations,
272            unused: stakes.unused,
273            epoch: stakes.epoch,
274            stake_history: stakes.stake_history.clone(),
275        })
276    }
277
278    #[cfg(feature = "dev-context-only-utils")]
279    pub fn new_for_tests(
280        epoch: Epoch,
281        vote_accounts: VoteAccounts,
282        stake_delegations: ImHashMap<Pubkey, StakeAccount>,
283    ) -> Self {
284        Self {
285            vote_accounts,
286            stake_delegations,
287            unused: 0,
288            epoch,
289            stake_history: StakeHistory::default(),
290        }
291    }
292
293    pub(crate) fn history(&self) -> &StakeHistory {
294        &self.stake_history
295    }
296
297    fn activate_epoch(
298        &mut self,
299        next_epoch: Epoch,
300        thread_pool: &ThreadPool,
301        new_rate_activation_epoch: Option<Epoch>,
302    ) {
303        let stake_delegations: Vec<_> = self.stake_delegations.values().collect();
304        // Wrap up the prev epoch by adding new stake history entry for the
305        // prev epoch.
306        let stake_history_entry = thread_pool.install(|| {
307            stake_delegations
308                .par_iter()
309                .fold(StakeActivationStatus::default, |acc, stake_account| {
310                    let delegation = stake_account.delegation();
311                    acc + delegation.stake_activating_and_deactivating(
312                        self.epoch,
313                        &self.stake_history,
314                        new_rate_activation_epoch,
315                    )
316                })
317                .reduce(StakeActivationStatus::default, Add::add)
318        });
319        self.stake_history.add(self.epoch, stake_history_entry);
320        self.epoch = next_epoch;
321        // Refresh the stake distribution of vote accounts for the next epoch,
322        // using new stake history.
323        self.vote_accounts = refresh_vote_accounts(
324            thread_pool,
325            self.epoch,
326            &self.vote_accounts,
327            &stake_delegations,
328            &self.stake_history,
329            new_rate_activation_epoch,
330        );
331    }
332
333    /// Sum the stakes that point to the given voter_pubkey
334    fn calculate_stake(
335        stake_delegations: &ImHashMap<Pubkey, StakeAccount>,
336        voter_pubkey: &Pubkey,
337        epoch: Epoch,
338        stake_history: &StakeHistory,
339        new_rate_activation_epoch: Option<Epoch>,
340    ) -> u64 {
341        stake_delegations
342            .values()
343            .map(StakeAccount::delegation)
344            .filter(|delegation| &delegation.voter_pubkey == voter_pubkey)
345            .map(|delegation| delegation.stake(epoch, stake_history, new_rate_activation_epoch))
346            .sum()
347    }
348
349    /// Sum the lamports of the vote accounts and the delegated stake
350    pub(crate) fn vote_balance_and_staked(&self) -> u64 {
351        let get_stake = |stake_account: &StakeAccount| stake_account.delegation().stake;
352        let get_lamports = |(_, vote_account): (_, &VoteAccount)| vote_account.lamports();
353
354        self.stake_delegations.values().map(get_stake).sum::<u64>()
355            + self.vote_accounts.iter().map(get_lamports).sum::<u64>()
356    }
357
358    fn remove_vote_account(&mut self, vote_pubkey: &Pubkey) -> Option<VoteAccount> {
359        self.vote_accounts.remove(vote_pubkey).map(|(_, a)| a)
360    }
361
362    fn remove_stake_delegation(
363        &mut self,
364        stake_pubkey: &Pubkey,
365        new_rate_activation_epoch: Option<Epoch>,
366    ) {
367        if let Some(stake_account) = self.stake_delegations.remove(stake_pubkey) {
368            let removed_delegation = stake_account.delegation();
369            let removed_stake = removed_delegation.stake(
370                self.epoch,
371                &self.stake_history,
372                new_rate_activation_epoch,
373            );
374            self.vote_accounts
375                .sub_stake(&removed_delegation.voter_pubkey, removed_stake);
376        }
377    }
378
379    fn upsert_vote_account(
380        &mut self,
381        vote_pubkey: &Pubkey,
382        vote_account: VoteAccount,
383        new_rate_activation_epoch: Option<Epoch>,
384    ) -> Option<VoteAccount> {
385        debug_assert_ne!(vote_account.lamports(), 0u64);
386
387        let stake_delegations = &self.stake_delegations;
388        self.vote_accounts.insert(*vote_pubkey, vote_account, || {
389            Self::calculate_stake(
390                stake_delegations,
391                vote_pubkey,
392                self.epoch,
393                &self.stake_history,
394                new_rate_activation_epoch,
395            )
396        })
397    }
398
399    fn upsert_stake_delegation(
400        &mut self,
401        stake_pubkey: Pubkey,
402        stake_account: StakeAccount,
403        new_rate_activation_epoch: Option<Epoch>,
404    ) {
405        debug_assert_ne!(stake_account.lamports(), 0u64);
406        let delegation = stake_account.delegation();
407        let voter_pubkey = delegation.voter_pubkey;
408        let stake = delegation.stake(self.epoch, &self.stake_history, new_rate_activation_epoch);
409        match self.stake_delegations.insert(stake_pubkey, stake_account) {
410            None => self.vote_accounts.add_stake(&voter_pubkey, stake),
411            Some(old_stake_account) => {
412                let old_delegation = old_stake_account.delegation();
413                let old_voter_pubkey = old_delegation.voter_pubkey;
414                let old_stake = old_delegation.stake(
415                    self.epoch,
416                    &self.stake_history,
417                    new_rate_activation_epoch,
418                );
419                if voter_pubkey != old_voter_pubkey || stake != old_stake {
420                    self.vote_accounts.sub_stake(&old_voter_pubkey, old_stake);
421                    self.vote_accounts.add_stake(&voter_pubkey, stake);
422                }
423            }
424        }
425    }
426
427    pub(crate) fn stake_delegations(&self) -> &ImHashMap<Pubkey, StakeAccount> {
428        &self.stake_delegations
429    }
430
431    pub(crate) fn highest_staked_node(&self) -> Option<&Pubkey> {
432        let vote_account = self.vote_accounts.find_max_by_delegated_stake()?;
433        Some(vote_account.node_pubkey())
434    }
435}
436
437impl StakesEnum {
438    pub fn vote_accounts(&self) -> &VoteAccounts {
439        match self {
440            StakesEnum::Accounts(stakes) => stakes.vote_accounts(),
441            StakesEnum::Delegations(stakes) => stakes.vote_accounts(),
442            StakesEnum::Stakes(stakes) => stakes.vote_accounts(),
443        }
444    }
445
446    pub(crate) fn staked_nodes(&self) -> Arc<HashMap<Pubkey, u64>> {
447        match self {
448            StakesEnum::Accounts(stakes) => stakes.staked_nodes(),
449            StakesEnum::Delegations(stakes) => stakes.staked_nodes(),
450            StakesEnum::Stakes(stakes) => stakes.staked_nodes(),
451        }
452    }
453}
454
455/// This conversion is very memory intensive so should only be used in
456/// development contexts.
457#[cfg(feature = "dev-context-only-utils")]
458impl From<Stakes<StakeAccount>> for Stakes<Delegation> {
459    fn from(stakes: Stakes<StakeAccount>) -> Self {
460        let stake_delegations = stakes
461            .stake_delegations
462            .into_iter()
463            .map(|(pubkey, stake_account)| (pubkey, *stake_account.delegation()))
464            .collect();
465        Self {
466            vote_accounts: stakes.vote_accounts,
467            stake_delegations,
468            unused: stakes.unused,
469            epoch: stakes.epoch,
470            stake_history: stakes.stake_history,
471        }
472    }
473}
474
475/// This conversion is very memory intensive so should only be used in
476/// development contexts.
477#[cfg(feature = "dev-context-only-utils")]
478impl From<Stakes<StakeAccount>> for Stakes<Stake> {
479    fn from(stakes: Stakes<StakeAccount>) -> Self {
480        let stake_delegations = stakes
481            .stake_delegations
482            .into_iter()
483            .map(|(pubkey, stake_account)| (pubkey, *stake_account.stake()))
484            .collect();
485        Self {
486            vote_accounts: stakes.vote_accounts,
487            stake_delegations,
488            unused: stakes.unused,
489            epoch: stakes.epoch,
490            stake_history: stakes.stake_history,
491        }
492    }
493}
494
495/// This conversion is memory intensive so should only be used in development
496/// contexts.
497#[cfg(feature = "dev-context-only-utils")]
498impl From<Stakes<Stake>> for Stakes<Delegation> {
499    fn from(stakes: Stakes<Stake>) -> Self {
500        let stake_delegations = stakes
501            .stake_delegations
502            .into_iter()
503            .map(|(pubkey, stake)| (pubkey, stake.delegation))
504            .collect();
505        Self {
506            vote_accounts: stakes.vote_accounts,
507            stake_delegations,
508            unused: stakes.unused,
509            epoch: stakes.epoch,
510            stake_history: stakes.stake_history,
511        }
512    }
513}
514
515/// This conversion is memory intensive so should only be used in development
516/// contexts.
517#[cfg(feature = "dev-context-only-utils")]
518impl From<StakesEnum> for Stakes<Delegation> {
519    fn from(stakes: StakesEnum) -> Self {
520        match stakes {
521            StakesEnum::Accounts(stakes) => stakes.into(),
522            StakesEnum::Delegations(stakes) => stakes,
523            StakesEnum::Stakes(stakes) => stakes.into(),
524        }
525    }
526}
527
528impl From<Stakes<StakeAccount>> for StakesEnum {
529    fn from(stakes: Stakes<StakeAccount>) -> Self {
530        Self::Accounts(stakes)
531    }
532}
533
534impl From<Stakes<Delegation>> for StakesEnum {
535    fn from(stakes: Stakes<Delegation>) -> Self {
536        Self::Delegations(stakes)
537    }
538}
539
540// Two StakesEnums are equal as long as they represent the same delegations;
541// whether these delegations are stored as StakeAccounts or Delegations.
542// Therefore, if one side is Stakes<StakeAccount> and the other is a
543// Stakes<Delegation> we convert the former one to Stakes<Delegation> before
544// comparing for equality.
545#[cfg(feature = "dev-context-only-utils")]
546impl PartialEq<StakesEnum> for StakesEnum {
547    fn eq(&self, other: &StakesEnum) -> bool {
548        match (self, other) {
549            (Self::Accounts(stakes), Self::Accounts(other)) => stakes == other,
550            (Self::Delegations(stakes), Self::Delegations(other)) => stakes == other,
551            (Self::Stakes(stakes), Self::Stakes(other)) => stakes == other,
552            (stakes, other) => {
553                let stakes = Stakes::<Delegation>::from(stakes.clone());
554                let other = Stakes::<Delegation>::from(other.clone());
555                stakes == other
556            }
557        }
558    }
559}
560
561fn refresh_vote_accounts(
562    thread_pool: &ThreadPool,
563    epoch: Epoch,
564    vote_accounts: &VoteAccounts,
565    stake_delegations: &[&StakeAccount],
566    stake_history: &StakeHistory,
567    new_rate_activation_epoch: Option<Epoch>,
568) -> VoteAccounts {
569    type StakesHashMap = HashMap</*voter:*/ Pubkey, /*stake:*/ u64>;
570    fn merge(mut stakes: StakesHashMap, other: StakesHashMap) -> StakesHashMap {
571        if stakes.len() < other.len() {
572            return merge(other, stakes);
573        }
574        for (pubkey, stake) in other {
575            *stakes.entry(pubkey).or_default() += stake;
576        }
577        stakes
578    }
579    let delegated_stakes = thread_pool.install(|| {
580        stake_delegations
581            .par_iter()
582            .fold(HashMap::default, |mut delegated_stakes, stake_account| {
583                let delegation = stake_account.delegation();
584                let entry = delegated_stakes.entry(delegation.voter_pubkey).or_default();
585                *entry += delegation.stake(epoch, stake_history, new_rate_activation_epoch);
586                delegated_stakes
587            })
588            .reduce(HashMap::default, merge)
589    });
590    vote_accounts
591        .iter()
592        .map(|(&vote_pubkey, vote_account)| {
593            let delegated_stake = delegated_stakes
594                .get(&vote_pubkey)
595                .copied()
596                .unwrap_or_default();
597            (vote_pubkey, (delegated_stake, vote_account.clone()))
598        })
599        .collect()
600}
601
602#[cfg(test)]
603pub(crate) mod tests {
604    use {
605        super::*,
606        rayon::ThreadPoolBuilder,
607        solana_sdk::{account::WritableAccount, pubkey::Pubkey, rent::Rent, stake},
608        solana_stake_program::stake_state,
609        solana_vote_program::vote_state::{self, VoteState, VoteStateVersions},
610    };
611
612    //  set up some dummies for a staked node     ((     vote      )  (     stake     ))
613    pub(crate) fn create_staked_node_accounts(
614        stake: u64,
615    ) -> ((Pubkey, AccountSharedData), (Pubkey, AccountSharedData)) {
616        let vote_pubkey = solana_pubkey::new_rand();
617        let vote_account =
618            vote_state::create_account(&vote_pubkey, &solana_pubkey::new_rand(), 0, 1);
619        let stake_pubkey = solana_pubkey::new_rand();
620        (
621            (vote_pubkey, vote_account),
622            (
623                stake_pubkey,
624                create_stake_account(stake, &vote_pubkey, &stake_pubkey),
625            ),
626        )
627    }
628
629    //   add stake to a vote_pubkey                               (   stake    )
630    pub(crate) fn create_stake_account(
631        stake: u64,
632        vote_pubkey: &Pubkey,
633        stake_pubkey: &Pubkey,
634    ) -> AccountSharedData {
635        stake_state::create_account(
636            stake_pubkey,
637            vote_pubkey,
638            &vote_state::create_account(vote_pubkey, &solana_pubkey::new_rand(), 0, 1),
639            &Rent::free(),
640            stake,
641        )
642    }
643
644    fn create_warming_staked_node_accounts(
645        stake: u64,
646        epoch: Epoch,
647    ) -> ((Pubkey, AccountSharedData), (Pubkey, AccountSharedData)) {
648        let vote_pubkey = solana_pubkey::new_rand();
649        let vote_account =
650            vote_state::create_account(&vote_pubkey, &solana_pubkey::new_rand(), 0, 1);
651        (
652            (vote_pubkey, vote_account),
653            create_warming_stake_account(stake, epoch, &vote_pubkey),
654        )
655    }
656
657    // add stake to a vote_pubkey                               (   stake    )
658    fn create_warming_stake_account(
659        stake: u64,
660        epoch: Epoch,
661        vote_pubkey: &Pubkey,
662    ) -> (Pubkey, AccountSharedData) {
663        let stake_pubkey = solana_pubkey::new_rand();
664        (
665            stake_pubkey,
666            stake_state::create_account_with_activation_epoch(
667                &stake_pubkey,
668                vote_pubkey,
669                &vote_state::create_account(vote_pubkey, &solana_pubkey::new_rand(), 0, 1),
670                &Rent::free(),
671                stake,
672                epoch,
673            ),
674        )
675    }
676
677    #[test]
678    fn test_stakes_basic() {
679        for i in 0..4 {
680            let stakes_cache = StakesCache::new(Stakes {
681                epoch: i,
682                ..Stakes::default()
683            });
684
685            let ((vote_pubkey, vote_account), (stake_pubkey, mut stake_account)) =
686                create_staked_node_accounts(10);
687
688            stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
689            stakes_cache.check_and_store(&stake_pubkey, &stake_account, None);
690            let stake = stake_state::stake_from(&stake_account).unwrap();
691            {
692                let stakes = stakes_cache.stakes();
693                let vote_accounts = stakes.vote_accounts();
694                assert!(vote_accounts.get(&vote_pubkey).is_some());
695                assert_eq!(
696                    vote_accounts.get_delegated_stake(&vote_pubkey),
697                    stake.stake(i, &StakeHistory::default(), None)
698                );
699            }
700
701            stake_account.set_lamports(42);
702            stakes_cache.check_and_store(&stake_pubkey, &stake_account, None);
703            {
704                let stakes = stakes_cache.stakes();
705                let vote_accounts = stakes.vote_accounts();
706                assert!(vote_accounts.get(&vote_pubkey).is_some());
707                assert_eq!(
708                    vote_accounts.get_delegated_stake(&vote_pubkey),
709                    stake.stake(i, &StakeHistory::default(), None)
710                ); // stays old stake, because only 10 is activated
711            }
712
713            // activate more
714            let mut stake_account =
715                create_stake_account(42, &vote_pubkey, &solana_pubkey::new_rand());
716            stakes_cache.check_and_store(&stake_pubkey, &stake_account, None);
717            let stake = stake_state::stake_from(&stake_account).unwrap();
718            {
719                let stakes = stakes_cache.stakes();
720                let vote_accounts = stakes.vote_accounts();
721                assert!(vote_accounts.get(&vote_pubkey).is_some());
722                assert_eq!(
723                    vote_accounts.get_delegated_stake(&vote_pubkey),
724                    stake.stake(i, &StakeHistory::default(), None)
725                ); // now stake of 42 is activated
726            }
727
728            stake_account.set_lamports(0);
729            stakes_cache.check_and_store(&stake_pubkey, &stake_account, None);
730            {
731                let stakes = stakes_cache.stakes();
732                let vote_accounts = stakes.vote_accounts();
733                assert!(vote_accounts.get(&vote_pubkey).is_some());
734                assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 0);
735            }
736        }
737    }
738
739    #[test]
740    fn test_stakes_highest() {
741        let stakes_cache = StakesCache::default();
742
743        assert_eq!(stakes_cache.stakes().highest_staked_node(), None);
744
745        let ((vote_pubkey, vote_account), (stake_pubkey, stake_account)) =
746            create_staked_node_accounts(10);
747
748        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
749        stakes_cache.check_and_store(&stake_pubkey, &stake_account, None);
750
751        let ((vote11_pubkey, vote11_account), (stake11_pubkey, stake11_account)) =
752            create_staked_node_accounts(20);
753
754        stakes_cache.check_and_store(&vote11_pubkey, &vote11_account, None);
755        stakes_cache.check_and_store(&stake11_pubkey, &stake11_account, None);
756
757        let vote11_node_pubkey = vote_state::from(&vote11_account).unwrap().node_pubkey;
758
759        let highest_staked_node = stakes_cache.stakes().highest_staked_node().copied();
760        assert_eq!(highest_staked_node, Some(vote11_node_pubkey));
761    }
762
763    #[test]
764    fn test_stakes_vote_account_disappear_reappear() {
765        let stakes_cache = StakesCache::new(Stakes {
766            epoch: 4,
767            ..Stakes::default()
768        });
769
770        let ((vote_pubkey, mut vote_account), (stake_pubkey, stake_account)) =
771            create_staked_node_accounts(10);
772
773        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
774        stakes_cache.check_and_store(&stake_pubkey, &stake_account, None);
775
776        {
777            let stakes = stakes_cache.stakes();
778            let vote_accounts = stakes.vote_accounts();
779            assert!(vote_accounts.get(&vote_pubkey).is_some());
780            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 10);
781        }
782
783        vote_account.set_lamports(0);
784        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
785
786        {
787            let stakes = stakes_cache.stakes();
788            let vote_accounts = stakes.vote_accounts();
789            assert!(vote_accounts.get(&vote_pubkey).is_none());
790            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 0);
791        }
792
793        vote_account.set_lamports(1);
794        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
795
796        {
797            let stakes = stakes_cache.stakes();
798            let vote_accounts = stakes.vote_accounts();
799            assert!(vote_accounts.get(&vote_pubkey).is_some());
800            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 10);
801        }
802
803        // Vote account too big
804        let cache_data = vote_account.data().to_vec();
805        let mut pushed = vote_account.data().to_vec();
806        pushed.push(0);
807        vote_account.set_data(pushed);
808        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
809
810        {
811            let stakes = stakes_cache.stakes();
812            let vote_accounts = stakes.vote_accounts();
813            assert!(vote_accounts.get(&vote_pubkey).is_none());
814            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 0);
815        }
816
817        // Vote account uninitialized
818        let default_vote_state = VoteState::default();
819        let versioned = VoteStateVersions::new_current(default_vote_state);
820        vote_state::to(&versioned, &mut vote_account).unwrap();
821        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
822
823        {
824            let stakes = stakes_cache.stakes();
825            let vote_accounts = stakes.vote_accounts();
826            assert!(vote_accounts.get(&vote_pubkey).is_none());
827            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 0);
828        }
829
830        vote_account.set_data(cache_data);
831        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
832
833        {
834            let stakes = stakes_cache.stakes();
835            let vote_accounts = stakes.vote_accounts();
836            assert!(vote_accounts.get(&vote_pubkey).is_some());
837            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 10);
838        }
839    }
840
841    #[test]
842    fn test_stakes_change_delegate() {
843        let stakes_cache = StakesCache::new(Stakes {
844            epoch: 4,
845            ..Stakes::default()
846        });
847
848        let ((vote_pubkey, vote_account), (stake_pubkey, stake_account)) =
849            create_staked_node_accounts(10);
850
851        let ((vote_pubkey2, vote_account2), (_stake_pubkey2, stake_account2)) =
852            create_staked_node_accounts(10);
853
854        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
855        stakes_cache.check_and_store(&vote_pubkey2, &vote_account2, None);
856
857        // delegates to vote_pubkey
858        stakes_cache.check_and_store(&stake_pubkey, &stake_account, None);
859
860        let stake = stake_state::stake_from(&stake_account).unwrap();
861
862        {
863            let stakes = stakes_cache.stakes();
864            let vote_accounts = stakes.vote_accounts();
865            assert!(vote_accounts.get(&vote_pubkey).is_some());
866            assert_eq!(
867                vote_accounts.get_delegated_stake(&vote_pubkey),
868                stake.stake(stakes.epoch, &stakes.stake_history, None)
869            );
870            assert!(vote_accounts.get(&vote_pubkey2).is_some());
871            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey2), 0);
872        }
873
874        // delegates to vote_pubkey2
875        stakes_cache.check_and_store(&stake_pubkey, &stake_account2, None);
876
877        {
878            let stakes = stakes_cache.stakes();
879            let vote_accounts = stakes.vote_accounts();
880            assert!(vote_accounts.get(&vote_pubkey).is_some());
881            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 0);
882            assert!(vote_accounts.get(&vote_pubkey2).is_some());
883            assert_eq!(
884                vote_accounts.get_delegated_stake(&vote_pubkey2),
885                stake.stake(stakes.epoch, &stakes.stake_history, None)
886            );
887        }
888    }
889    #[test]
890    fn test_stakes_multiple_stakers() {
891        let stakes_cache = StakesCache::new(Stakes {
892            epoch: 4,
893            ..Stakes::default()
894        });
895
896        let ((vote_pubkey, vote_account), (stake_pubkey, stake_account)) =
897            create_staked_node_accounts(10);
898
899        let stake_pubkey2 = solana_pubkey::new_rand();
900        let stake_account2 = create_stake_account(10, &vote_pubkey, &stake_pubkey2);
901
902        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
903
904        // delegates to vote_pubkey
905        stakes_cache.check_and_store(&stake_pubkey, &stake_account, None);
906        stakes_cache.check_and_store(&stake_pubkey2, &stake_account2, None);
907
908        {
909            let stakes = stakes_cache.stakes();
910            let vote_accounts = stakes.vote_accounts();
911            assert!(vote_accounts.get(&vote_pubkey).is_some());
912            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 20);
913        }
914    }
915
916    #[test]
917    fn test_activate_epoch() {
918        let stakes_cache = StakesCache::default();
919
920        let ((vote_pubkey, vote_account), (stake_pubkey, stake_account)) =
921            create_staked_node_accounts(10);
922
923        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
924        stakes_cache.check_and_store(&stake_pubkey, &stake_account, None);
925        let stake = stake_state::stake_from(&stake_account).unwrap();
926
927        {
928            let stakes = stakes_cache.stakes();
929            let vote_accounts = stakes.vote_accounts();
930            assert_eq!(
931                vote_accounts.get_delegated_stake(&vote_pubkey),
932                stake.stake(stakes.epoch, &stakes.stake_history, None)
933            );
934        }
935        let thread_pool = ThreadPoolBuilder::new().num_threads(1).build().unwrap();
936        stakes_cache.activate_epoch(3, &thread_pool, None);
937        {
938            let stakes = stakes_cache.stakes();
939            let vote_accounts = stakes.vote_accounts();
940            assert_eq!(
941                vote_accounts.get_delegated_stake(&vote_pubkey),
942                stake.stake(stakes.epoch, &stakes.stake_history, None)
943            );
944        }
945    }
946
947    #[test]
948    fn test_stakes_not_delegate() {
949        let stakes_cache = StakesCache::new(Stakes {
950            epoch: 4,
951            ..Stakes::default()
952        });
953
954        let ((vote_pubkey, vote_account), (stake_pubkey, stake_account)) =
955            create_staked_node_accounts(10);
956
957        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
958        stakes_cache.check_and_store(&stake_pubkey, &stake_account, None);
959
960        {
961            let stakes = stakes_cache.stakes();
962            let vote_accounts = stakes.vote_accounts();
963            assert!(vote_accounts.get(&vote_pubkey).is_some());
964            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 10);
965        }
966
967        // not a stake account, and whacks above entry
968        stakes_cache.check_and_store(
969            &stake_pubkey,
970            &AccountSharedData::new(1, 0, &stake::program::id()),
971            None,
972        );
973        {
974            let stakes = stakes_cache.stakes();
975            let vote_accounts = stakes.vote_accounts();
976            assert!(vote_accounts.get(&vote_pubkey).is_some());
977            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 0);
978        }
979    }
980
981    #[test]
982    fn test_vote_balance_and_staked_empty() {
983        let stakes = Stakes::<StakeAccount>::default();
984        assert_eq!(stakes.vote_balance_and_staked(), 0);
985    }
986
987    #[test]
988    fn test_vote_balance_and_staked_normal() {
989        let stakes_cache = StakesCache::default();
990        #[allow(non_local_definitions)]
991        impl Stakes<StakeAccount> {
992            fn vote_balance_and_warmed_staked(&self) -> u64 {
993                let vote_balance: u64 = self
994                    .vote_accounts
995                    .iter()
996                    .map(|(_pubkey, account)| account.lamports())
997                    .sum();
998                let warmed_stake: u64 = self
999                    .vote_accounts
1000                    .delegated_stakes()
1001                    .map(|(_pubkey, stake)| stake)
1002                    .sum();
1003                vote_balance + warmed_stake
1004            }
1005        }
1006
1007        let genesis_epoch = 0;
1008        let ((vote_pubkey, vote_account), (stake_pubkey, stake_account)) =
1009            create_warming_staked_node_accounts(10, genesis_epoch);
1010        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
1011        stakes_cache.check_and_store(&stake_pubkey, &stake_account, None);
1012
1013        {
1014            let stakes = stakes_cache.stakes();
1015            assert_eq!(stakes.vote_balance_and_staked(), 11);
1016            assert_eq!(stakes.vote_balance_and_warmed_staked(), 1);
1017        }
1018
1019        let thread_pool = ThreadPoolBuilder::new().num_threads(1).build().unwrap();
1020        for (epoch, expected_warmed_stake) in ((genesis_epoch + 1)..=3).zip(&[2, 3, 4]) {
1021            stakes_cache.activate_epoch(epoch, &thread_pool, None);
1022            // vote_balance_and_staked() always remain to return same lamports
1023            // while vote_balance_and_warmed_staked() gradually increases
1024            let stakes = stakes_cache.stakes();
1025            assert_eq!(stakes.vote_balance_and_staked(), 11);
1026            assert_eq!(
1027                stakes.vote_balance_and_warmed_staked(),
1028                *expected_warmed_stake
1029            );
1030        }
1031    }
1032}