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    serde::Serialize,
15    solana_account::{AccountSharedData, ReadableAccount},
16    solana_accounts_db::utils::create_account_shared_data,
17    solana_clock::Epoch,
18    solana_leader_schedule::SlotLeader,
19    solana_pubkey::Pubkey,
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#[cfg_attr(feature = "dev-context-only-utils", qualifiers(pub))]
44pub(crate) use serde_stakes::serialize_stake_accounts_to_delegation_format;
45
46#[derive(Debug, Error)]
47pub enum Error {
48    #[error("Invalid delegation: {0}")]
49    InvalidDelegation(Pubkey),
50    #[error(transparent)]
51    InvalidStakeAccount(#[from] stake_account::Error),
52    #[error("Stake account not found: {0}")]
53    StakeAccountNotFound(Pubkey),
54    #[error("Vote account mismatch: {0}")]
55    VoteAccountMismatch(Pubkey),
56    #[error("Vote account not cached: {0}")]
57    VoteAccountNotCached(Pubkey),
58    #[error("Vote account not found: {0}")]
59    VoteAccountNotFound(Pubkey),
60}
61
62#[derive(Debug, Clone, PartialEq, Eq, ToPrimitive)]
63pub enum InvalidCacheEntryReason {
64    Missing,
65    BadState,
66    WrongOwner,
67}
68
69type StakeAccount = stake_account::StakeAccount<Delegation>;
70pub(crate) type DelegatedStakes = ImblHashMap<Pubkey, u64>;
71
72#[cfg_attr(feature = "frozen-abi", derive(AbiExample))]
73#[derive(Default, Debug)]
74pub(crate) struct StakesCache(RwLock<Stakes<StakeAccount>>);
75
76impl StakesCache {
77    pub(crate) fn new(stakes: Stakes<StakeAccount>) -> Self {
78        Self(RwLock::new(stakes))
79    }
80
81    pub(crate) fn stakes(&self) -> RwLockReadGuard<'_, Stakes<StakeAccount>> {
82        self.0.read().unwrap()
83    }
84
85    pub(crate) fn check_and_store(
86        &self,
87        pubkey: &Pubkey,
88        account: &impl ReadableAccount,
89        new_rate_activation_epoch: Option<Epoch>,
90    ) {
91        // TODO: If the account is already cached as a vote or stake account
92        // but the owner changes, then this needs to evict the account from
93        // the cache. see:
94        // https://github.com/solana-labs/solana/pull/24200#discussion_r849935444
95        let owner = account.owner();
96        // Zero lamport accounts are not stored in accounts-db
97        // and so should be removed from cache as well.
98        if account.lamports() == 0 {
99            if solana_vote_program::check_id(owner) {
100                let _old_vote_account = {
101                    let mut stakes = self.0.write().unwrap();
102                    stakes.remove_vote_account(pubkey)
103                };
104            } else if stake_program::check_id(owner) {
105                let mut stakes = self.0.write().unwrap();
106                stakes.remove_stake_delegation(pubkey, new_rate_activation_epoch);
107            }
108            return;
109        }
110        debug_assert_ne!(account.lamports(), 0u64);
111        if solana_vote_program::check_id(owner) {
112            if VoteStateVersions::is_correct_size_and_initialized(account.data()) {
113                match VoteAccount::try_from(create_account_shared_data(account)) {
114                    Ok(vote_account) => {
115                        // drop the old account after releasing the lock
116                        let _old_vote_account = {
117                            let mut stakes = self.0.write().unwrap();
118                            stakes.upsert_vote_account(pubkey, vote_account)
119                        };
120                    }
121                    Err(_) => {
122                        // drop the old account after releasing the lock
123                        let _old_vote_account = {
124                            let mut stakes = self.0.write().unwrap();
125                            stakes.remove_vote_account(pubkey)
126                        };
127                    }
128                }
129            } else {
130                // drop the old account after releasing the lock
131                let _old_vote_account = {
132                    let mut stakes = self.0.write().unwrap();
133                    stakes.remove_vote_account(pubkey)
134                };
135            };
136        } else if stake_program::check_id(owner) {
137            match StakeAccount::try_from(create_account_shared_data(account)) {
138                Ok(stake_account) => {
139                    let mut stakes = self.0.write().unwrap();
140                    stakes.upsert_stake_delegation(
141                        *pubkey,
142                        stake_account,
143                        new_rate_activation_epoch,
144                    );
145                }
146                Err(_) => {
147                    let mut stakes = self.0.write().unwrap();
148                    stakes.remove_stake_delegation(pubkey, new_rate_activation_epoch);
149                }
150            }
151        }
152    }
153
154    pub(crate) fn activate_epoch(
155        &self,
156        next_epoch: Epoch,
157        stake_history: StakeHistory,
158        vote_accounts: VoteAccounts,
159        delegated_stakes: DelegatedStakes,
160    ) {
161        let mut stakes = self.0.write().unwrap();
162        stakes.activate_epoch(next_epoch, stake_history, vote_accounts, delegated_stakes)
163    }
164
165    pub(crate) fn refresh_delegated_stakes(&self, new_rate_activation_epoch: Option<Epoch>) {
166        let mut stakes = self.0.write().unwrap();
167        stakes.refresh_delegated_stakes(new_rate_activation_epoch);
168    }
169}
170
171/// The generic type T is either Delegation or StakeAccount.
172/// [`Stakes<Delegation>`] is equivalent to the old code and is used for backward
173/// compatibility in [`crate::bank::BankFieldsToDeserialize`].
174/// But banks cache [`Stakes<StakeAccount>`] which includes the entire stake
175/// account and StakeStateV2 deserialized from the account. Doing so, will remove
176/// the need to load the stake account from accounts-db when working with
177/// stake-delegations.
178#[cfg_attr(feature = "frozen-abi", derive(AbiExample, StableAbi, StableAbiSample))]
179#[derive(Default, Clone, PartialEq, Debug, Serialize, SchemaWrite)]
180#[cfg_attr(
181    feature = "dev-context-only-utils",
182    field_qualifiers(
183        vote_accounts(pub),
184        stake_delegations(pub),
185        delegated_stakes(pub),
186        unused(pub),
187        epoch(pub),
188        stake_history(pub),
189    )
190)]
191pub struct Stakes<T: Clone> {
192    /// vote accounts
193    vote_accounts: VoteAccounts,
194
195    /// stake_delegations
196    #[cfg_attr(
197        feature = "frozen-abi",
198        stable_abi_sample(with = "sample_collection_sized(rng, SequenceLenMax(1))")
199    )]
200    #[wincode(with = "FromIntoIterator<ImblHashMap<Pubkey, T>, BincodeLen>")]
201    stake_delegations: ImblHashMap<Pubkey, T>,
202
203    /// current effective stake delegated to each vote account pubkey
204    #[cfg_attr(feature = "frozen-abi", stable_abi_sample(with = "Default::default()"))]
205    #[serde(skip)]
206    #[wincode(skip)]
207    delegated_stakes: DelegatedStakes,
208
209    /// unused
210    unused: u64,
211
212    /// current epoch, used to calculate current stake
213    epoch: Epoch,
214
215    /// history of staking levels
216    stake_history: StakeHistory,
217}
218
219impl<T: Clone> Stakes<T> {
220    pub fn new(vote_accounts: VoteAccounts, epoch: Epoch) -> Stakes<T> {
221        Stakes {
222            vote_accounts,
223            epoch,
224            stake_delegations: ImblHashMap::new(),
225            delegated_stakes: DelegatedStakes::default(),
226            unused: 0,
227            stake_history: StakeHistory::default(),
228        }
229    }
230
231    pub fn clone_and_filter_for_vat(
232        &self,
233        max_vote_accounts: usize,
234        minimum_vote_account_balance: u64,
235    ) -> Stakes<T> {
236        Self::new(
237            self.vote_accounts
238                .clone_and_filter_for_vat(max_vote_accounts, minimum_vote_account_balance),
239            self.epoch,
240        )
241    }
242
243    pub fn vote_accounts(&self) -> &VoteAccounts {
244        &self.vote_accounts
245    }
246
247    pub(crate) fn staked_nodes(&self) -> Arc<HashMap<Pubkey, u64>> {
248        self.vote_accounts.staked_nodes()
249    }
250
251    /// Destructure self and return the fields needed by EpochStakes
252    pub(crate) fn into_epoch_stakes_fields(self) -> (Epoch, VoteAccounts, StakeHistory) {
253        let Self {
254            vote_accounts,
255            stake_delegations: _,
256            delegated_stakes: _,
257            unused: _,
258            epoch,
259            stake_history,
260        } = self;
261        (epoch, vote_accounts, stake_history)
262    }
263}
264
265impl Stakes<StakeAccount> {
266    pub(crate) fn new_from_accounts_for_genesis<'a, T: ReadableAccount + 'a>(
267        new_rate_activation_epoch: Option<Epoch>,
268        accounts: impl IntoIterator<Item = (&'a Pubkey, &'a T)>,
269    ) -> Self {
270        let stake_history = StakeHistory::default();
271        let mut vote_accounts = VoteAccountsHashMap::default();
272        let mut delegated_stakes = DelegatedStakes::default();
273        let mut stake_delegations = ImblHashMap::new();
274        let epoch = 0;
275
276        for (pubkey, account) in accounts {
277            if account.lamports() == 0 {
278                continue;
279            }
280
281            if solana_vote_program::check_id(account.owner()) {
282                if VoteStateVersions::is_correct_size_and_initialized(account.data())
283                    && let Ok(vote_account) =
284                        VoteAccount::try_from(create_account_shared_data(account))
285                {
286                    vote_accounts.insert(*pubkey, (0, vote_account));
287                }
288            } else if stake_program::check_id(account.owner())
289                && let Ok(stake_account) =
290                    StakeAccount::try_from(create_account_shared_data(account))
291            {
292                let delegation = stake_account.delegation();
293                let stake = delegation.stake_v2(epoch, &stake_history, new_rate_activation_epoch);
294                if stake != 0 {
295                    *delegated_stakes.entry(delegation.voter_pubkey).or_default() += stake;
296                }
297                stake_delegations.insert(*pubkey, stake_account);
298            }
299        }
300
301        let mut vote_accounts = VoteAccounts::from(Arc::new(vote_accounts));
302        for (vote_pubkey, stake) in &delegated_stakes {
303            vote_accounts.add_stake(vote_pubkey, *stake);
304        }
305
306        Self {
307            vote_accounts,
308            stake_delegations,
309            delegated_stakes,
310            unused: 0,
311            epoch,
312            stake_history,
313        }
314    }
315
316    /// Creates a Stake<StakeAccount> from DeserializableDelegationStakes by loading the
317    /// full account state for respective stake pubkeys. get_account function
318    /// should return the account at the respective slot where stakes where
319    /// cached.
320    #[cfg_attr(feature = "dev-context-only-utils", qualifiers(pub))]
321    pub(crate) fn load_from_deserialized_delegations<F>(
322        stakes: DeserializableDelegationStakes,
323        get_account: F,
324    ) -> Result<Self, Error>
325    where
326        F: Fn(&Pubkey) -> Option<AccountSharedData> + Sync,
327    {
328        let stake_delegations = stakes
329            .stake_delegations
330            .into_par_iter()
331            // We use fold/reduce to aggregate the results, which does a bit more work than calling
332            // collect()/collect_vec_list() and then imbl::HashMap::from_iter(collected.into_iter()),
333            // but it does it in background threads, so effectively it's faster.
334            .try_fold(ImblHashMap::new, |mut map, (pubkey, delegation)| {
335                let Some(stake_account) = get_account(&pubkey) else {
336                    return Err(Error::StakeAccountNotFound(pubkey));
337                };
338
339                // Assert that all valid vote-accounts referenced in stake delegations are already
340                // contained in `stakes.vote_account`.
341                let voter_pubkey = &delegation.voter_pubkey;
342                if stakes.vote_accounts.get(voter_pubkey).is_none()
343                    && let Some(account) = get_account(voter_pubkey)
344                    && VoteStateVersions::is_correct_size_and_initialized(account.data())
345                    && VoteAccount::try_from(account.clone()).is_ok()
346                {
347                    error!("vote account not cached: {voter_pubkey}, {account:?}");
348                    return Err(Error::VoteAccountNotCached(*voter_pubkey));
349                }
350
351                let stake_account = StakeAccount::try_from(stake_account)?;
352                // Sanity check that the delegation is consistent with what is
353                // stored in the account.
354                if stake_account.delegation() == &delegation {
355                    map.insert(pubkey, stake_account);
356                    Ok(map)
357                } else {
358                    Err(Error::InvalidDelegation(pubkey))
359                }
360            })
361            .try_reduce(ImblHashMap::new, |a, b| Ok(a.union(b)))?;
362
363        // Assert that cached vote accounts are consistent with accounts-db.
364        //
365        // This currently includes ~5500 accounts, parallelizing brings minor
366        // (sub 2s) improvements.
367        for (pubkey, vote_account) in stakes.vote_accounts.iter() {
368            let Some(account) = get_account(pubkey) else {
369                return Err(Error::VoteAccountNotFound(*pubkey));
370            };
371            let vote_account = vote_account.account();
372            if vote_account != &account {
373                error!("vote account mismatch: {pubkey}, {vote_account:?}, {account:?}");
374                return Err(Error::VoteAccountMismatch(*pubkey));
375            }
376        }
377
378        Ok(Self {
379            vote_accounts: stakes.vote_accounts.clone(),
380            stake_delegations,
381            delegated_stakes: DelegatedStakes::default(),
382            unused: stakes.unused,
383            epoch: stakes.epoch,
384            stake_history: stakes.stake_history,
385        })
386    }
387
388    #[cfg(feature = "dev-context-only-utils")]
389    pub fn new_for_tests(
390        epoch: Epoch,
391        vote_accounts: VoteAccounts,
392        stake_delegations: ImblHashMap<Pubkey, StakeAccount>,
393    ) -> Self {
394        let stake_history = StakeHistory::default();
395        let delegated_stakes =
396            Self::calculate_delegated_stakes(&stake_delegations, epoch, &stake_history, None);
397        Self {
398            vote_accounts,
399            stake_delegations,
400            delegated_stakes,
401            unused: 0,
402            epoch,
403            stake_history,
404        }
405    }
406
407    pub(crate) fn history(&self) -> &StakeHistory {
408        &self.stake_history
409    }
410
411    pub(crate) fn calculate_activated_stake(
412        &self,
413        next_epoch: Epoch,
414        thread_pool: &ThreadPool,
415        new_rate_activation_epoch: Option<Epoch>,
416        stake_delegations: &[(&Pubkey, &StakeAccount)],
417    ) -> (
418        StakeHistory,
419        VoteAccounts,
420        DelegatedStakes,
421        RewardEpochDelegatedStakes,
422    ) {
423        // Wrap up the prev epoch by adding new stake history entry for the
424        // prev epoch.
425        let (stake_history_entry, effective_delegated_stakes) = thread_pool.install(|| {
426            stake_delegations
427                .par_iter()
428                .fold(
429                    || (StakeActivationStatus::default(), HashMap::default()),
430                    |(acc, mut delegated_stakes), (_stake_pubkey, stake_account)| {
431                        let delegation = stake_account.delegation();
432                        let activation_status = delegation.stake_activating_and_deactivating_v2(
433                            self.epoch,
434                            &self.stake_history,
435                            new_rate_activation_epoch,
436                        );
437                        *delegated_stakes.entry(delegation.voter_pubkey).or_default() +=
438                            activation_status.effective;
439                        (acc + activation_status, delegated_stakes)
440                    },
441                )
442                .reduce(
443                    || (StakeActivationStatus::default(), HashMap::default()),
444                    |(activation_status_a, delegated_stakes_a),
445                     (activation_status_b, delegated_stakes_b)| {
446                        (
447                            activation_status_a + activation_status_b,
448                            merge_delegated_stakes(delegated_stakes_a, delegated_stakes_b),
449                        )
450                    },
451                )
452        });
453        let mut stake_history = self.stake_history.clone();
454        stake_history.add(self.epoch, stake_history_entry);
455        // Refresh the stake distribution of vote accounts for the next epoch,
456        // using new stake history.
457        let (vote_accounts, delegated_stakes) = refresh_vote_accounts(
458            thread_pool,
459            next_epoch,
460            &self.vote_accounts,
461            stake_delegations,
462            &stake_history,
463            new_rate_activation_epoch,
464        );
465        let reward_epoch_delegated_stakes = RewardEpochDelegatedStakes {
466            epoch: self.epoch,
467            delegated_stakes: effective_delegated_stakes,
468        };
469        (
470            stake_history,
471            vote_accounts,
472            delegated_stakes,
473            reward_epoch_delegated_stakes,
474        )
475    }
476
477    pub(crate) fn activate_epoch(
478        &mut self,
479        next_epoch: Epoch,
480        stake_history: StakeHistory,
481        vote_accounts: VoteAccounts,
482        delegated_stakes: DelegatedStakes,
483    ) {
484        self.epoch = next_epoch;
485        self.stake_history = stake_history;
486        self.vote_accounts = vote_accounts;
487        self.delegated_stakes = delegated_stakes;
488    }
489
490    fn calculate_delegated_stakes(
491        stake_delegations: &ImblHashMap<Pubkey, StakeAccount>,
492        epoch: Epoch,
493        stake_history: &StakeHistory,
494        new_rate_activation_epoch: Option<Epoch>,
495    ) -> DelegatedStakes {
496        let mut delegated_stakes = DelegatedStakes::new();
497        for stake_account in stake_delegations.values() {
498            let delegation = stake_account.delegation();
499            let stake = delegation.stake_v2(epoch, stake_history, new_rate_activation_epoch);
500            if stake != 0 {
501                *delegated_stakes.entry(delegation.voter_pubkey).or_default() += stake;
502            }
503        }
504        delegated_stakes
505    }
506
507    fn refresh_delegated_stakes(&mut self, new_rate_activation_epoch: Option<Epoch>) {
508        self.delegated_stakes = Self::calculate_delegated_stakes(
509            &self.stake_delegations,
510            self.epoch,
511            &self.stake_history,
512            new_rate_activation_epoch,
513        );
514    }
515
516    fn add_delegated_stake(&mut self, voter_pubkey: Pubkey, stake: u64) {
517        if stake == 0 {
518            return;
519        }
520        *self.delegated_stakes.entry(voter_pubkey).or_default() += stake;
521    }
522
523    fn sub_delegated_stake(&mut self, voter_pubkey: &Pubkey, stake: u64) {
524        if stake == 0 {
525            return;
526        }
527        let current_stake = self
528            .delegated_stakes
529            .get_mut(voter_pubkey)
530            .expect("subtraction from missing delegated stake");
531        *current_stake = current_stake
532            .checked_sub(stake)
533            .expect("subtraction value exceeds delegated stake");
534        if *current_stake == 0 {
535            self.delegated_stakes.remove(voter_pubkey);
536        }
537    }
538
539    fn remove_vote_account(&mut self, vote_pubkey: &Pubkey) -> Option<VoteAccount> {
540        self.vote_accounts.remove(vote_pubkey).map(|(_, a)| a)
541    }
542
543    fn remove_stake_delegation(
544        &mut self,
545        stake_pubkey: &Pubkey,
546        new_rate_activation_epoch: Option<Epoch>,
547    ) {
548        if let Some(stake_account) = self.stake_delegations.remove(stake_pubkey) {
549            let removed_delegation = stake_account.delegation();
550            let removed_stake = removed_delegation.stake_v2(
551                self.epoch,
552                &self.stake_history,
553                new_rate_activation_epoch,
554            );
555            self.sub_delegated_stake(&removed_delegation.voter_pubkey, removed_stake);
556            self.vote_accounts
557                .sub_stake(&removed_delegation.voter_pubkey, removed_stake);
558        }
559    }
560
561    fn upsert_vote_account(
562        &mut self,
563        vote_pubkey: &Pubkey,
564        vote_account: VoteAccount,
565    ) -> Option<VoteAccount> {
566        debug_assert_ne!(vote_account.lamports(), 0u64);
567
568        let calculate_delegated_stake = || {
569            self.delegated_stakes
570                .get(vote_pubkey)
571                .copied()
572                .unwrap_or_default()
573        };
574        self.vote_accounts
575            .insert(*vote_pubkey, vote_account, calculate_delegated_stake)
576    }
577
578    fn upsert_stake_delegation(
579        &mut self,
580        stake_pubkey: Pubkey,
581        stake_account: StakeAccount,
582        new_rate_activation_epoch: Option<Epoch>,
583    ) {
584        debug_assert_ne!(stake_account.lamports(), 0u64);
585        let delegation = stake_account.delegation();
586        let voter_pubkey = delegation.voter_pubkey;
587        let stake = delegation.stake_v2(self.epoch, &self.stake_history, new_rate_activation_epoch);
588        match self.stake_delegations.insert(stake_pubkey, stake_account) {
589            None => {
590                self.add_delegated_stake(voter_pubkey, stake);
591                self.vote_accounts.add_stake(&voter_pubkey, stake);
592            }
593            Some(old_stake_account) => {
594                let old_delegation = old_stake_account.delegation();
595                let old_voter_pubkey = old_delegation.voter_pubkey;
596                let old_stake = old_delegation.stake_v2(
597                    self.epoch,
598                    &self.stake_history,
599                    new_rate_activation_epoch,
600                );
601                if voter_pubkey != old_voter_pubkey || stake != old_stake {
602                    self.sub_delegated_stake(&old_voter_pubkey, old_stake);
603                    self.add_delegated_stake(voter_pubkey, stake);
604                    self.vote_accounts.sub_stake(&old_voter_pubkey, old_stake);
605                    self.vote_accounts.add_stake(&voter_pubkey, stake);
606                }
607            }
608        }
609    }
610
611    /// Returns a reference to the map of stake delegations.
612    ///
613    /// # Performance
614    ///
615    /// `[imbl::HashMap]` is a [hash array mapped trie (HAMT)][hamt], which means
616    /// that inserts, deletions and lookups are average-case O(1) and
617    /// worst-case O(log n). However, the performance of iterations is poor due
618    /// to depth-first traversal and jumps. Currently it's also impossible to
619    /// iterate over it with [`rayon`].
620    ///
621    /// [hamt]: https://en.wikipedia.org/wiki/Hash_array_mapped_trie
622    pub(crate) fn stake_delegations(&self) -> &ImblHashMap<Pubkey, StakeAccount> {
623        &self.stake_delegations
624    }
625
626    /// Collects stake delegations into a vector, which then can be used for
627    /// parallel iteration with [`rayon`].
628    ///
629    /// # Performance
630    ///
631    /// The execution of this method takes ~200ms and it collects elements of
632    /// the [`imbl::HashMap`], which is a [hash array mapped trie (HAMT)][hamt],
633    /// so that operation involves a depth-first traversal with jumps. However,
634    /// it's still a reasonable tradeoff if the caller iterates over these
635    /// elements.
636    ///
637    /// [hamt]: https://en.wikipedia.org/wiki/Hash_array_mapped_trie
638    pub(crate) fn stake_delegations_vec(&self) -> Vec<(&Pubkey, &StakeAccount)> {
639        self.stake_delegations.iter().collect()
640    }
641
642    pub(crate) fn highest_staked_node(&self) -> Option<SlotLeader> {
643        let (vote_address, vote_account) = self.vote_accounts.find_max_by_delegated_stake()?;
644        Some(SlotLeader {
645            id: *vote_account.node_pubkey(),
646            vote_address: *vote_address,
647        })
648    }
649}
650
651/// Macro to generate `From<Stakes<From>> for Stakes<To>` impls.
652#[cfg(feature = "dev-context-only-utils")]
653macro_rules! impl_stake_format_conversion {
654    ($from:ty, $to:ty, |$binding:ident| $expr:expr) => {
655        /// This conversion is memory intensive so should only be used in development contexts.
656        impl From<Stakes<$from>> for Stakes<$to> {
657            fn from(stakes: Stakes<$from>) -> Self {
658                let Stakes {
659                    vote_accounts,
660                    stake_delegations,
661                    delegated_stakes: _,
662                    unused,
663                    epoch,
664                    stake_history,
665                } = stakes;
666                let stake_delegations = stake_delegations
667                    .into_iter()
668                    .map(|(pubkey, $binding)| (pubkey, $expr))
669                    .collect();
670                Self {
671                    vote_accounts,
672                    stake_delegations,
673                    delegated_stakes: DelegatedStakes::default(),
674                    unused,
675                    epoch,
676                    stake_history,
677                }
678            }
679        }
680    };
681}
682
683#[cfg(feature = "dev-context-only-utils")]
684impl_stake_format_conversion!(StakeAccount, Delegation, |sa| *sa.delegation());
685
686#[cfg(feature = "dev-context-only-utils")]
687impl_stake_format_conversion!(StakeAccount, Stake, |sa| *sa.stake());
688
689#[cfg(feature = "dev-context-only-utils")]
690impl_stake_format_conversion!(Stake, Delegation, |stake| stake.delegation);
691
692fn merge_delegated_stakes(
693    mut stakes: HashMap</*voter:*/ Pubkey, /*stake:*/ u64>,
694    other: HashMap</*voter:*/ Pubkey, /*stake:*/ u64>,
695) -> HashMap</*voter:*/ Pubkey, /*stake:*/ u64> {
696    if stakes.len() < other.len() {
697        return merge_delegated_stakes(other, stakes);
698    }
699    for (pubkey, stake) in other {
700        *stakes.entry(pubkey).or_default() += stake;
701    }
702    stakes
703}
704
705fn refresh_vote_accounts(
706    thread_pool: &ThreadPool,
707    epoch: Epoch,
708    vote_accounts: &VoteAccounts,
709    stake_delegations: &[(&Pubkey, &StakeAccount)],
710    stake_history: &StakeHistory,
711    new_rate_activation_epoch: Option<Epoch>,
712) -> (VoteAccounts, DelegatedStakes) {
713    fn merge(mut stakes: DelegatedStakes, other: DelegatedStakes) -> DelegatedStakes {
714        if stakes.len() < other.len() {
715            return merge(other, stakes);
716        }
717        for (pubkey, stake) in other {
718            *stakes.entry(pubkey).or_default() += stake;
719        }
720        stakes
721    }
722    let delegated_stakes = thread_pool.install(|| {
723        stake_delegations
724            .par_iter()
725            .fold(
726                DelegatedStakes::default,
727                |mut delegated_stakes, (_stake_pubkey, stake_account)| {
728                    let delegation = stake_account.delegation();
729                    let stake =
730                        delegation.stake_v2(epoch, stake_history, new_rate_activation_epoch);
731                    if stake != 0 {
732                        *delegated_stakes.entry(delegation.voter_pubkey).or_default() += stake;
733                    }
734                    delegated_stakes
735                },
736            )
737            .reduce(DelegatedStakes::default, merge)
738    });
739    let vote_accounts = vote_accounts
740        .iter()
741        .map(|(&vote_pubkey, vote_account)| {
742            let delegated_stake = delegated_stakes
743                .get(&vote_pubkey)
744                .copied()
745                .unwrap_or_default();
746            (vote_pubkey, (delegated_stake, vote_account.clone()))
747        })
748        .collect();
749    (vote_accounts, delegated_stakes)
750}
751
752#[cfg(test)]
753pub(crate) mod tests {
754    use {
755        super::*,
756        crate::stake_utils,
757        rayon::ThreadPoolBuilder,
758        solana_account::{WritableAccount, state_traits::StateMutWincode as _},
759        solana_pubkey::Pubkey,
760        solana_rent::Rent,
761        solana_stake_interface::{self as stake, state::StakeStateV2},
762        solana_vote_interface::state::{BLS_PUBLIC_KEY_COMPRESSED_SIZE, VoteStateV4},
763        solana_vote_program::vote_state,
764    };
765
766    impl Stakes<Delegation> {
767        /// Convert deserialized stakes into runtime stakes representation
768        pub(crate) fn from_deserialized(stakes: DeserializableDelegationStakes) -> Self {
769            Self {
770                vote_accounts: stakes.vote_accounts,
771                stake_delegations: ImblHashMap::from_iter(stakes.stake_delegations),
772                delegated_stakes: DelegatedStakes::default(),
773                unused: stakes.unused,
774                epoch: stakes.epoch,
775                stake_history: stakes.stake_history,
776            }
777        }
778    }
779
780    //  set up some dummies for a staked node     ((     vote      )  (     stake     ))
781    pub(crate) fn create_staked_node_accounts(
782        stake: u64,
783        rent: &Rent,
784    ) -> ((Pubkey, AccountSharedData), (Pubkey, AccountSharedData)) {
785        let vote_pubkey = solana_pubkey::new_rand();
786        let node_pubkey = solana_pubkey::new_rand();
787        let vote_account = vote_state::create_v4_account_with_authorized(
788            &node_pubkey,
789            &vote_pubkey,
790            [0u8; BLS_PUBLIC_KEY_COMPRESSED_SIZE],
791            &vote_pubkey,
792            0,
793            &vote_pubkey,
794            0,
795            &node_pubkey,
796            1,
797        );
798        let stake_pubkey = solana_pubkey::new_rand();
799        (
800            (vote_pubkey, vote_account),
801            (
802                stake_pubkey,
803                create_stake_account(stake, &vote_pubkey, &stake_pubkey, rent),
804            ),
805        )
806    }
807
808    //   add stake to a vote_pubkey                               (   stake    )
809    pub(crate) fn create_stake_account(
810        stake: u64,
811        vote_pubkey: &Pubkey,
812        stake_pubkey: &Pubkey,
813        rent: &Rent,
814    ) -> AccountSharedData {
815        let node_pubkey = solana_pubkey::new_rand();
816        let lamports = rent.minimum_balance(StakeStateV2::size_of()) + stake;
817        stake_utils::create_stake_account(
818            stake_pubkey,
819            vote_pubkey,
820            &vote_state::create_v4_account_with_authorized(
821                &node_pubkey,
822                vote_pubkey,
823                [0u8; BLS_PUBLIC_KEY_COMPRESSED_SIZE],
824                vote_pubkey,
825                0,
826                vote_pubkey,
827                0,
828                &node_pubkey,
829                1,
830            ),
831            rent,
832            lamports,
833        )
834    }
835
836    #[test]
837    fn test_stakes_basic() {
838        for i in 0..4 {
839            let stakes_cache = StakesCache::new(Stakes {
840                epoch: i,
841                ..Stakes::default()
842            });
843            let rent = Rent::default();
844
845            let ((vote_pubkey, vote_account), (stake_pubkey, mut stake_account)) =
846                create_staked_node_accounts(10, &rent);
847
848            stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
849            stakes_cache.check_and_store(&stake_pubkey, &stake_account, None);
850            let stake_state: StakeStateV2 = stake_account.state().unwrap();
851            let stake = stake_state.stake().unwrap();
852            {
853                let stakes = stakes_cache.stakes();
854                let vote_accounts = stakes.vote_accounts();
855                assert!(vote_accounts.get(&vote_pubkey).is_some());
856                let expected_stake = stake.stake_v2(i, &StakeHistory::default(), None);
857                assert_eq!(
858                    vote_accounts.get_delegated_stake(&vote_pubkey),
859                    expected_stake
860                );
861            }
862
863            stake_account.set_lamports(42);
864            stakes_cache.check_and_store(&stake_pubkey, &stake_account, None);
865            {
866                let stakes = stakes_cache.stakes();
867                let vote_accounts = stakes.vote_accounts();
868                assert!(vote_accounts.get(&vote_pubkey).is_some());
869                let expected_stake = stake.stake_v2(i, &StakeHistory::default(), None);
870                assert_eq!(
871                    vote_accounts.get_delegated_stake(&vote_pubkey),
872                    expected_stake
873                ); // stays old stake, because only 10 is activated
874            }
875
876            // activate more
877            let mut stake_account =
878                create_stake_account(42, &vote_pubkey, &solana_pubkey::new_rand(), &rent);
879            stakes_cache.check_and_store(&stake_pubkey, &stake_account, None);
880            let stake_state: StakeStateV2 = stake_account.state().unwrap();
881            let stake = stake_state.stake().unwrap();
882            {
883                let stakes = stakes_cache.stakes();
884                let vote_accounts = stakes.vote_accounts();
885                assert!(vote_accounts.get(&vote_pubkey).is_some());
886                let expected_stake = stake.stake_v2(i, &StakeHistory::default(), None);
887                assert_eq!(
888                    vote_accounts.get_delegated_stake(&vote_pubkey),
889                    expected_stake
890                ); // now stake of 42 is activated
891            }
892
893            stake_account.set_lamports(0);
894            stakes_cache.check_and_store(&stake_pubkey, &stake_account, None);
895            {
896                let stakes = stakes_cache.stakes();
897                let vote_accounts = stakes.vote_accounts();
898                assert!(vote_accounts.get(&vote_pubkey).is_some());
899                assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 0);
900            }
901        }
902    }
903
904    #[test]
905    fn test_stakes_highest() {
906        let stakes_cache = StakesCache::default();
907        let rent = Rent::default();
908
909        assert_eq!(stakes_cache.stakes().highest_staked_node(), None);
910
911        let ((vote_pubkey, vote_account), (stake_pubkey, stake_account)) =
912            create_staked_node_accounts(10, &rent);
913
914        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
915        stakes_cache.check_and_store(&stake_pubkey, &stake_account, None);
916
917        let ((vote11_pubkey, vote11_account), (stake11_pubkey, stake11_account)) =
918            create_staked_node_accounts(20, &rent);
919
920        stakes_cache.check_and_store(&vote11_pubkey, &vote11_account, None);
921        stakes_cache.check_and_store(&stake11_pubkey, &stake11_account, None);
922
923        let vote11_node_pubkey = VoteStateV4::deserialize(vote11_account.data(), &vote11_pubkey)
924            .unwrap()
925            .node_pubkey;
926
927        let highest_staked_node = stakes_cache.stakes().highest_staked_node();
928        assert_eq!(
929            highest_staked_node,
930            Some(SlotLeader {
931                id: vote11_node_pubkey,
932                vote_address: vote11_pubkey,
933            })
934        );
935    }
936
937    #[test]
938    fn test_stakes_vote_account_disappear_reappear() {
939        let stakes_cache = StakesCache::new(Stakes {
940            epoch: 4,
941            ..Stakes::default()
942        });
943        let rent = Rent::default();
944
945        let ((vote_pubkey, mut vote_account), (stake_pubkey, stake_account)) =
946            create_staked_node_accounts(10, &rent);
947
948        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
949        stakes_cache.check_and_store(&stake_pubkey, &stake_account, None);
950
951        {
952            let stakes = stakes_cache.stakes();
953            let vote_accounts = stakes.vote_accounts();
954            assert!(vote_accounts.get(&vote_pubkey).is_some());
955            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 10);
956        }
957
958        vote_account.set_lamports(0);
959        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
960
961        {
962            let stakes = stakes_cache.stakes();
963            let vote_accounts = stakes.vote_accounts();
964            assert!(vote_accounts.get(&vote_pubkey).is_none());
965            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 0);
966        }
967
968        vote_account.set_lamports(1);
969        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
970
971        {
972            let stakes = stakes_cache.stakes();
973            let vote_accounts = stakes.vote_accounts();
974            assert!(vote_accounts.get(&vote_pubkey).is_some());
975            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 10);
976        }
977
978        // Vote account too big
979        let cache_data = vote_account.data().to_vec();
980        let mut pushed = vote_account.data().to_vec();
981        pushed.push(0);
982        vote_account.set_data_from_slice(&pushed);
983        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
984
985        {
986            let stakes = stakes_cache.stakes();
987            let vote_accounts = stakes.vote_accounts();
988            assert!(vote_accounts.get(&vote_pubkey).is_none());
989            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 0);
990        }
991
992        // Vote account uninitialized
993        vote_account.set_data_from_slice(&vec![0; VoteStateV4::size_of()]);
994        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
995
996        {
997            let stakes = stakes_cache.stakes();
998            let vote_accounts = stakes.vote_accounts();
999            assert!(vote_accounts.get(&vote_pubkey).is_none());
1000            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 0);
1001        }
1002
1003        vote_account.set_data_from_slice(&cache_data);
1004        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
1005
1006        {
1007            let stakes = stakes_cache.stakes();
1008            let vote_accounts = stakes.vote_accounts();
1009            assert!(vote_accounts.get(&vote_pubkey).is_some());
1010            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 10);
1011        }
1012    }
1013
1014    #[test]
1015    fn test_stakes_change_delegate() {
1016        let stakes_cache = StakesCache::new(Stakes {
1017            epoch: 4,
1018            ..Stakes::default()
1019        });
1020        let rent = Rent::default();
1021
1022        let ((vote_pubkey, vote_account), (stake_pubkey, stake_account)) =
1023            create_staked_node_accounts(10, &rent);
1024
1025        let ((vote_pubkey2, vote_account2), (_stake_pubkey2, stake_account2)) =
1026            create_staked_node_accounts(10, &rent);
1027
1028        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
1029        stakes_cache.check_and_store(&vote_pubkey2, &vote_account2, None);
1030
1031        // delegates to vote_pubkey
1032        stakes_cache.check_and_store(&stake_pubkey, &stake_account, None);
1033
1034        let stake_state: StakeStateV2 = stake_account.state().unwrap();
1035        let stake = stake_state.stake().unwrap();
1036
1037        {
1038            let stakes = stakes_cache.stakes();
1039            let vote_accounts = stakes.vote_accounts();
1040            assert!(vote_accounts.get(&vote_pubkey).is_some());
1041            let expected_stake = stake.stake_v2(stakes.epoch, &stakes.stake_history, None);
1042            assert_eq!(
1043                vote_accounts.get_delegated_stake(&vote_pubkey),
1044                expected_stake
1045            );
1046            assert!(vote_accounts.get(&vote_pubkey2).is_some());
1047            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey2), 0);
1048        }
1049
1050        // delegates to vote_pubkey2
1051        stakes_cache.check_and_store(&stake_pubkey, &stake_account2, None);
1052
1053        {
1054            let stakes = stakes_cache.stakes();
1055            let vote_accounts = stakes.vote_accounts();
1056            assert!(vote_accounts.get(&vote_pubkey).is_some());
1057            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 0);
1058            assert!(vote_accounts.get(&vote_pubkey2).is_some());
1059            let expected_stake = stake.stake_v2(stakes.epoch, &stakes.stake_history, None);
1060            assert_eq!(
1061                vote_accounts.get_delegated_stake(&vote_pubkey2),
1062                expected_stake
1063            );
1064        }
1065    }
1066    #[test]
1067    fn test_stakes_multiple_stakers() {
1068        let stakes_cache = StakesCache::new(Stakes {
1069            epoch: 4,
1070            ..Stakes::default()
1071        });
1072        let rent = Rent::default();
1073
1074        let ((vote_pubkey, vote_account), (stake_pubkey, stake_account)) =
1075            create_staked_node_accounts(10, &rent);
1076
1077        let stake_pubkey2 = solana_pubkey::new_rand();
1078        let stake_account2 = create_stake_account(10, &vote_pubkey, &stake_pubkey2, &rent);
1079
1080        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
1081
1082        // delegates to vote_pubkey
1083        stakes_cache.check_and_store(&stake_pubkey, &stake_account, None);
1084        stakes_cache.check_and_store(&stake_pubkey2, &stake_account2, None);
1085
1086        {
1087            let stakes = stakes_cache.stakes();
1088            let vote_accounts = stakes.vote_accounts();
1089            assert!(vote_accounts.get(&vote_pubkey).is_some());
1090            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 20);
1091        }
1092    }
1093
1094    #[test]
1095    fn test_activate_epoch() {
1096        let stakes_cache = StakesCache::default();
1097        let rent = Rent::default();
1098
1099        let ((vote_pubkey, vote_account), (stake_pubkey, stake_account)) =
1100            create_staked_node_accounts(10, &rent);
1101
1102        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
1103        stakes_cache.check_and_store(&stake_pubkey, &stake_account, None);
1104        let stake_state: StakeStateV2 = stake_account.state().unwrap();
1105        let stake = stake_state.stake().unwrap();
1106
1107        let initial_expected_stake = {
1108            let stakes = stakes_cache.stakes();
1109            stake.stake_v2(stakes.epoch, &stakes.stake_history, None)
1110        };
1111        {
1112            let stakes = stakes_cache.stakes();
1113            let vote_accounts = stakes.vote_accounts();
1114            assert_eq!(
1115                vote_accounts.get_delegated_stake(&vote_pubkey),
1116                initial_expected_stake
1117            );
1118        }
1119        let thread_pool = ThreadPoolBuilder::new().num_threads(1).build().unwrap();
1120        let next_epoch = 3;
1121        let (stake_history, vote_accounts, delegated_stakes, effective_delegated_stakes) = {
1122            let stakes = stakes_cache.stakes();
1123            let stake_delegations = stakes.stake_delegations_vec();
1124            stakes.calculate_activated_stake(next_epoch, &thread_pool, None, &stake_delegations)
1125        };
1126        assert_eq!(
1127            effective_delegated_stakes
1128                .delegated_stakes
1129                .get(&vote_pubkey)
1130                .copied(),
1131            Some(initial_expected_stake)
1132        );
1133        stakes_cache.activate_epoch(next_epoch, stake_history, vote_accounts, delegated_stakes);
1134        {
1135            let stakes = stakes_cache.stakes();
1136            let vote_accounts = stakes.vote_accounts();
1137            let expected_stake = stake.stake_v2(stakes.epoch, &stakes.stake_history, None);
1138            assert_eq!(
1139                vote_accounts.get_delegated_stake(&vote_pubkey),
1140                expected_stake
1141            );
1142        }
1143    }
1144
1145    #[test]
1146    fn test_stakes_not_delegate() {
1147        let stakes_cache = StakesCache::new(Stakes {
1148            epoch: 4,
1149            ..Stakes::default()
1150        });
1151        let rent = Rent::default();
1152
1153        let ((vote_pubkey, vote_account), (stake_pubkey, stake_account)) =
1154            create_staked_node_accounts(10, &rent);
1155
1156        stakes_cache.check_and_store(&vote_pubkey, &vote_account, None);
1157        stakes_cache.check_and_store(&stake_pubkey, &stake_account, None);
1158
1159        {
1160            let stakes = stakes_cache.stakes();
1161            let vote_accounts = stakes.vote_accounts();
1162            assert!(vote_accounts.get(&vote_pubkey).is_some());
1163            assert_eq!(vote_accounts.get_delegated_stake(&vote_pubkey), 10);
1164        }
1165
1166        // not a stake account, and whacks above entry
1167        stakes_cache.check_and_store(
1168            &stake_pubkey,
1169            &AccountSharedData::new(1, 0, &stake::program::id()),
1170            None,
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), 0);
1177        }
1178    }
1179}