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