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