1use {
2 itertools::Itertools,
3 serde::{
4 de::{MapAccess, Visitor},
5 ser::{Serialize, Serializer},
6 },
7 solana_sdk::{
8 account::{AccountSharedData, ReadableAccount},
9 instruction::InstructionError,
10 pubkey::Pubkey,
11 vote::state::VoteState,
12 },
13 std::{
14 cmp::Ordering,
15 collections::{hash_map::Entry, HashMap},
16 fmt,
17 iter::FromIterator,
18 mem::{self, MaybeUninit},
19 ptr::addr_of_mut,
20 sync::{Arc, OnceLock},
21 },
22 thiserror::Error,
23};
24
25#[cfg_attr(feature = "frozen-abi", derive(AbiExample))]
26#[derive(Clone, Debug, PartialEq)]
27pub struct VoteAccount(Arc<VoteAccountInner>);
28
29#[derive(Debug, Error)]
30pub enum Error {
31 #[error(transparent)]
32 InstructionError(#[from] InstructionError),
33 #[error("Invalid vote account owner: {0}")]
34 InvalidOwner(Pubkey),
35}
36
37#[cfg_attr(feature = "frozen-abi", derive(AbiExample))]
38#[derive(Debug)]
39struct VoteAccountInner {
40 account: AccountSharedData,
41 vote_state: VoteState,
42}
43
44pub type VoteAccountsHashMap = HashMap<Pubkey, (u64, VoteAccount)>;
45#[cfg_attr(feature = "frozen-abi", derive(AbiExample))]
46#[derive(Debug, Serialize, Deserialize)]
47pub struct VoteAccounts {
48 #[serde(deserialize_with = "deserialize_accounts_hash_map")]
49 vote_accounts: Arc<VoteAccountsHashMap>,
50 #[serde(skip)]
52 staked_nodes: OnceLock<
53 Arc<
54 HashMap<
55 Pubkey, u64, >,
58 >,
59 >,
60}
61
62impl Clone for VoteAccounts {
63 fn clone(&self) -> Self {
64 Self {
65 vote_accounts: Arc::clone(&self.vote_accounts),
66 staked_nodes: OnceLock::new(),
71 }
72 }
73}
74
75impl VoteAccount {
76 pub fn account(&self) -> &AccountSharedData {
77 &self.0.account
78 }
79
80 pub fn lamports(&self) -> u64 {
81 self.0.account.lamports()
82 }
83
84 pub fn owner(&self) -> &Pubkey {
85 self.0.account.owner()
86 }
87
88 pub fn vote_state(&self) -> &VoteState {
89 &self.0.vote_state
90 }
91
92 pub fn node_pubkey(&self) -> &Pubkey {
94 &self.0.vote_state.node_pubkey
95 }
96
97 #[cfg(feature = "dev-context-only-utils")]
98 pub fn new_random() -> VoteAccount {
99 use {
100 rand::Rng as _,
101 solana_sdk::{
102 clock::Clock,
103 vote::state::{VoteInit, VoteStateVersions},
104 },
105 };
106
107 let mut rng = rand::thread_rng();
108
109 let vote_init = VoteInit {
110 node_pubkey: Pubkey::new_unique(),
111 authorized_voter: Pubkey::new_unique(),
112 authorized_withdrawer: Pubkey::new_unique(),
113 commission: rng.gen(),
114 };
115 let clock = Clock {
116 slot: rng.gen(),
117 epoch_start_timestamp: rng.gen(),
118 epoch: rng.gen(),
119 leader_schedule_epoch: rng.gen(),
120 unix_timestamp: rng.gen(),
121 };
122 let vote_state = VoteState::new(&vote_init, &clock);
123 let account = AccountSharedData::new_data(
124 rng.gen(), &VoteStateVersions::new_current(vote_state.clone()),
126 &solana_sdk::vote::program::id(), )
128 .unwrap();
129
130 VoteAccount::try_from(account).unwrap()
131 }
132}
133
134impl VoteAccounts {
135 pub fn len(&self) -> usize {
136 self.vote_accounts.len()
137 }
138
139 pub fn is_empty(&self) -> bool {
140 self.vote_accounts.is_empty()
141 }
142
143 pub fn staked_nodes(&self) -> Arc<HashMap<Pubkey, u64>> {
144 self.staked_nodes
145 .get_or_init(|| {
146 Arc::new(
147 self.vote_accounts
148 .values()
149 .filter(|(stake, _)| *stake != 0u64)
150 .map(|(stake, vote_account)| (*vote_account.node_pubkey(), stake))
151 .into_grouping_map()
152 .aggregate(|acc, _node_pubkey, stake| {
153 Some(acc.unwrap_or_default() + stake)
154 }),
155 )
156 })
157 .clone()
158 }
159
160 pub fn get(&self, pubkey: &Pubkey) -> Option<&VoteAccount> {
161 let (_stake, vote_account) = self.vote_accounts.get(pubkey)?;
162 Some(vote_account)
163 }
164
165 pub fn get_delegated_stake(&self, pubkey: &Pubkey) -> u64 {
166 self.vote_accounts
167 .get(pubkey)
168 .map(|(stake, _vote_account)| *stake)
169 .unwrap_or_default()
170 }
171
172 pub fn iter(&self) -> impl Iterator<Item = (&Pubkey, &VoteAccount)> {
173 self.vote_accounts
174 .iter()
175 .map(|(vote_pubkey, (_stake, vote_account))| (vote_pubkey, vote_account))
176 }
177
178 pub fn delegated_stakes(&self) -> impl Iterator<Item = (&Pubkey, u64)> {
179 self.vote_accounts
180 .iter()
181 .map(|(vote_pubkey, (stake, _vote_account))| (vote_pubkey, *stake))
182 }
183
184 pub fn find_max_by_delegated_stake(&self) -> Option<&VoteAccount> {
185 let key = |(_pubkey, (stake, _vote_account)): &(_, &(u64, _))| *stake;
186 let (_pubkey, (_stake, vote_account)) = self.vote_accounts.iter().max_by_key(key)?;
187 Some(vote_account)
188 }
189
190 pub fn insert(
191 &mut self,
192 pubkey: Pubkey,
193 new_vote_account: VoteAccount,
194 calculate_stake: impl FnOnce() -> u64,
195 ) -> Option<VoteAccount> {
196 let vote_accounts = Arc::make_mut(&mut self.vote_accounts);
197 match vote_accounts.entry(pubkey) {
198 Entry::Occupied(mut entry) => {
199 let (stake, old_vote_account) = entry.get_mut();
201
202 if let Some(staked_nodes) = self.staked_nodes.get_mut() {
203 let old_node_pubkey = old_vote_account.node_pubkey();
204 let new_node_pubkey = new_vote_account.node_pubkey();
205 if new_node_pubkey != old_node_pubkey {
206 Self::do_sub_node_stake(staked_nodes, *stake, old_node_pubkey);
209 Self::do_add_node_stake(staked_nodes, *stake, *new_node_pubkey);
210 }
211 }
212
213 Some(mem::replace(old_vote_account, new_vote_account))
215 }
216 Entry::Vacant(entry) => {
217 let (stake, vote_account) = entry.insert((calculate_stake(), new_vote_account));
219 if let Some(staked_nodes) = self.staked_nodes.get_mut() {
220 Self::do_add_node_stake(staked_nodes, *stake, *vote_account.node_pubkey());
221 }
222 None
223 }
224 }
225 }
226
227 pub fn remove(&mut self, pubkey: &Pubkey) -> Option<(u64, VoteAccount)> {
228 let vote_accounts = Arc::make_mut(&mut self.vote_accounts);
229 let entry = vote_accounts.remove(pubkey);
230 if let Some((stake, ref vote_account)) = entry {
231 self.sub_node_stake(stake, vote_account);
232 }
233 entry
234 }
235
236 pub fn add_stake(&mut self, pubkey: &Pubkey, delta: u64) {
237 let vote_accounts = Arc::make_mut(&mut self.vote_accounts);
238 if let Some((stake, vote_account)) = vote_accounts.get_mut(pubkey) {
239 *stake += delta;
240 let vote_account = vote_account.clone();
241 self.add_node_stake(delta, &vote_account);
242 }
243 }
244
245 pub fn sub_stake(&mut self, pubkey: &Pubkey, delta: u64) {
246 let vote_accounts = Arc::make_mut(&mut self.vote_accounts);
247 if let Some((stake, vote_account)) = vote_accounts.get_mut(pubkey) {
248 *stake = stake
249 .checked_sub(delta)
250 .expect("subtraction value exceeds account's stake");
251 let vote_account = vote_account.clone();
252 self.sub_node_stake(delta, &vote_account);
253 }
254 }
255
256 fn add_node_stake(&mut self, stake: u64, vote_account: &VoteAccount) {
257 let Some(staked_nodes) = self.staked_nodes.get_mut() else {
258 return;
259 };
260
261 VoteAccounts::do_add_node_stake(staked_nodes, stake, *vote_account.node_pubkey());
262 }
263
264 fn do_add_node_stake(
265 staked_nodes: &mut Arc<HashMap<Pubkey, u64>>,
266 stake: u64,
267 node_pubkey: Pubkey,
268 ) {
269 if stake == 0u64 {
270 return;
271 }
272
273 Arc::make_mut(staked_nodes)
274 .entry(node_pubkey)
275 .and_modify(|s| *s += stake)
276 .or_insert(stake);
277 }
278
279 fn sub_node_stake(&mut self, stake: u64, vote_account: &VoteAccount) {
280 let Some(staked_nodes) = self.staked_nodes.get_mut() else {
281 return;
282 };
283
284 VoteAccounts::do_sub_node_stake(staked_nodes, stake, vote_account.node_pubkey());
285 }
286
287 fn do_sub_node_stake(
288 staked_nodes: &mut Arc<HashMap<Pubkey, u64>>,
289 stake: u64,
290 node_pubkey: &Pubkey,
291 ) {
292 if stake == 0u64 {
293 return;
294 }
295
296 let staked_nodes = Arc::make_mut(staked_nodes);
297 let current_stake = staked_nodes
298 .get_mut(node_pubkey)
299 .expect("this should not happen");
300 match (*current_stake).cmp(&stake) {
301 Ordering::Less => panic!("subtraction value exceeds node's stake"),
302 Ordering::Equal => {
303 staked_nodes.remove(node_pubkey);
304 }
305 Ordering::Greater => *current_stake -= stake,
306 }
307 }
308}
309
310impl Serialize for VoteAccount {
311 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
312 where
313 S: Serializer,
314 {
315 self.0.account.serialize(serializer)
316 }
317}
318
319impl From<VoteAccount> for AccountSharedData {
320 fn from(account: VoteAccount) -> Self {
321 account.0.account.clone()
322 }
323}
324
325impl TryFrom<AccountSharedData> for VoteAccount {
326 type Error = Error;
327 fn try_from(account: AccountSharedData) -> Result<Self, Self::Error> {
328 if !solana_sdk::vote::program::check_id(account.owner()) {
329 return Err(Error::InvalidOwner(*account.owner()));
330 }
331
332 let mut inner = Arc::new(MaybeUninit::<VoteAccountInner>::uninit());
334 let inner_ptr = Arc::get_mut(&mut inner)
335 .expect("we're the only ref")
336 .as_mut_ptr();
337
338 unsafe {
344 let vote_state = addr_of_mut!((*inner_ptr).vote_state);
345 let vote_state = &mut *(vote_state as *mut MaybeUninit<VoteState>);
351
352 if let Err(e) = VoteState::deserialize_into_uninit(account.data(), vote_state) {
354 return Err(e.into());
360 }
361
362 addr_of_mut!((*inner_ptr).account).write(account);
364
365 Ok(VoteAccount(mem::transmute::<
369 Arc<MaybeUninit<VoteAccountInner>>,
370 Arc<VoteAccountInner>,
371 >(inner)))
372 }
373 }
374}
375
376impl PartialEq<VoteAccountInner> for VoteAccountInner {
377 fn eq(&self, other: &Self) -> bool {
378 let Self {
379 account,
380 vote_state: _,
381 } = self;
382 account == &other.account
383 }
384}
385
386impl Default for VoteAccounts {
387 fn default() -> Self {
388 Self {
389 vote_accounts: Arc::default(),
390 staked_nodes: OnceLock::new(),
391 }
392 }
393}
394
395impl PartialEq<VoteAccounts> for VoteAccounts {
396 fn eq(&self, other: &Self) -> bool {
397 let Self {
398 vote_accounts,
399 staked_nodes: _,
400 } = self;
401 vote_accounts == &other.vote_accounts
402 }
403}
404
405impl From<Arc<VoteAccountsHashMap>> for VoteAccounts {
406 fn from(vote_accounts: Arc<VoteAccountsHashMap>) -> Self {
407 Self {
408 vote_accounts,
409 staked_nodes: OnceLock::new(),
410 }
411 }
412}
413
414impl AsRef<VoteAccountsHashMap> for VoteAccounts {
415 fn as_ref(&self) -> &VoteAccountsHashMap {
416 &self.vote_accounts
417 }
418}
419
420impl From<&VoteAccounts> for Arc<VoteAccountsHashMap> {
421 fn from(vote_accounts: &VoteAccounts) -> Self {
422 Arc::clone(&vote_accounts.vote_accounts)
423 }
424}
425
426impl FromIterator<(Pubkey, (u64, VoteAccount))> for VoteAccounts {
427 fn from_iter<I>(iter: I) -> Self
428 where
429 I: IntoIterator<Item = (Pubkey, (u64, VoteAccount))>,
430 {
431 Self::from(Arc::new(HashMap::from_iter(iter)))
432 }
433}
434
435fn deserialize_accounts_hash_map<'de, D>(
442 deserializer: D,
443) -> Result<Arc<VoteAccountsHashMap>, D::Error>
444where
445 D: serde::Deserializer<'de>,
446{
447 struct VoteAccountsVisitor;
448
449 impl<'de> Visitor<'de> for VoteAccountsVisitor {
450 type Value = Arc<VoteAccountsHashMap>;
451
452 fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
453 formatter.write_str("a map of vote accounts")
454 }
455
456 fn visit_map<M>(self, mut access: M) -> Result<Self::Value, M::Error>
457 where
458 M: MapAccess<'de>,
459 {
460 let mut accounts = HashMap::new();
461
462 while let Some((pubkey, (stake, account))) =
463 access.next_entry::<Pubkey, (u64, AccountSharedData)>()?
464 {
465 match VoteAccount::try_from(account) {
466 Ok(vote_account) => {
467 accounts.insert(pubkey, (stake, vote_account));
468 }
469 Err(e) => {
470 log::warn!("failed to deserialize vote account: {e}");
471 }
472 }
473 }
474
475 Ok(Arc::new(accounts))
476 }
477 }
478
479 deserializer.deserialize_map(VoteAccountsVisitor)
480}
481
482#[cfg(test)]
483mod tests {
484 use {
485 super::*,
486 bincode::Options,
487 rand::Rng,
488 solana_sdk::{
489 account::WritableAccount,
490 pubkey::Pubkey,
491 sysvar::clock::Clock,
492 vote::state::{VoteInit, VoteStateVersions},
493 },
494 std::iter::repeat_with,
495 };
496
497 fn new_rand_vote_account<R: Rng>(
498 rng: &mut R,
499 node_pubkey: Option<Pubkey>,
500 ) -> (AccountSharedData, VoteState) {
501 let vote_init = VoteInit {
502 node_pubkey: node_pubkey.unwrap_or_else(Pubkey::new_unique),
503 authorized_voter: Pubkey::new_unique(),
504 authorized_withdrawer: Pubkey::new_unique(),
505 commission: rng.gen(),
506 };
507 let clock = Clock {
508 slot: rng.gen(),
509 epoch_start_timestamp: rng.gen(),
510 epoch: rng.gen(),
511 leader_schedule_epoch: rng.gen(),
512 unix_timestamp: rng.gen(),
513 };
514 let vote_state = VoteState::new(&vote_init, &clock);
515 let account = AccountSharedData::new_data(
516 rng.gen(), &VoteStateVersions::new_current(vote_state.clone()),
518 &solana_sdk::vote::program::id(), )
520 .unwrap();
521 (account, vote_state)
522 }
523
524 fn new_rand_vote_accounts<R: Rng>(
525 rng: &mut R,
526 num_nodes: usize,
527 ) -> impl Iterator<Item = (Pubkey, (u64, VoteAccount))> + '_ {
528 let nodes: Vec<_> = repeat_with(Pubkey::new_unique).take(num_nodes).collect();
529 repeat_with(move || {
530 let node = nodes[rng.gen_range(0..nodes.len())];
531 let (account, _) = new_rand_vote_account(rng, Some(node));
532 let stake = rng.gen_range(0..997);
533 let vote_account = VoteAccount::try_from(account).unwrap();
534 (Pubkey::new_unique(), (stake, vote_account))
535 })
536 }
537
538 fn staked_nodes<'a, I>(vote_accounts: I) -> HashMap<Pubkey, u64>
539 where
540 I: IntoIterator<Item = &'a (Pubkey, (u64, VoteAccount))>,
541 {
542 let mut staked_nodes = HashMap::new();
543 for (_, (stake, vote_account)) in vote_accounts
544 .into_iter()
545 .filter(|(_, (stake, _))| *stake != 0)
546 {
547 staked_nodes
548 .entry(*vote_account.node_pubkey())
549 .and_modify(|s| *s += *stake)
550 .or_insert(*stake);
551 }
552 staked_nodes
553 }
554
555 #[test]
556 fn test_vote_account_try_from() {
557 let mut rng = rand::thread_rng();
558 let (account, vote_state) = new_rand_vote_account(&mut rng, None);
559 let lamports = account.lamports();
560 let vote_account = VoteAccount::try_from(account.clone()).unwrap();
561 assert_eq!(lamports, vote_account.lamports());
562 assert_eq!(vote_state, *vote_account.vote_state());
563 assert_eq!(&account, vote_account.account());
564 }
565
566 #[test]
567 #[should_panic(expected = "InvalidOwner")]
568 fn test_vote_account_try_from_invalid_owner() {
569 let mut rng = rand::thread_rng();
570 let (mut account, _) = new_rand_vote_account(&mut rng, None);
571 account.set_owner(Pubkey::new_unique());
572 VoteAccount::try_from(account).unwrap();
573 }
574
575 #[test]
576 #[should_panic(expected = "InvalidAccountData")]
577 fn test_vote_account_try_from_invalid_account() {
578 let mut account = AccountSharedData::default();
579 account.set_owner(solana_sdk::vote::program::id());
580 VoteAccount::try_from(account).unwrap();
581 }
582
583 #[test]
584 fn test_vote_account_serialize() {
585 let mut rng = rand::thread_rng();
586 let (account, vote_state) = new_rand_vote_account(&mut rng, None);
587 let vote_account = VoteAccount::try_from(account.clone()).unwrap();
588 assert_eq!(vote_state, *vote_account.vote_state());
589 assert_eq!(
591 bincode::serialize(&account).unwrap(),
592 bincode::serialize(&vote_account).unwrap()
593 );
594 }
595
596 #[test]
597 fn test_vote_accounts_serialize() {
598 let mut rng = rand::thread_rng();
599 let vote_accounts_hash_map: VoteAccountsHashMap =
600 new_rand_vote_accounts(&mut rng, 64).take(1024).collect();
601 let vote_accounts = VoteAccounts::from(Arc::new(vote_accounts_hash_map.clone()));
602 assert!(vote_accounts.staked_nodes().len() > 32);
603 assert_eq!(
604 bincode::serialize(&vote_accounts).unwrap(),
605 bincode::serialize(&vote_accounts_hash_map).unwrap(),
606 );
607 assert_eq!(
608 bincode::options().serialize(&vote_accounts).unwrap(),
609 bincode::options()
610 .serialize(&vote_accounts_hash_map)
611 .unwrap(),
612 )
613 }
614
615 #[test]
616 fn test_vote_accounts_deserialize() {
617 let mut rng = rand::thread_rng();
618 let vote_accounts_hash_map: VoteAccountsHashMap =
619 new_rand_vote_accounts(&mut rng, 64).take(1024).collect();
620 let data = bincode::serialize(&vote_accounts_hash_map).unwrap();
621 let vote_accounts: VoteAccounts = bincode::deserialize(&data).unwrap();
622 assert!(vote_accounts.staked_nodes().len() > 32);
623 assert_eq!(*vote_accounts.vote_accounts, vote_accounts_hash_map);
624 let data = bincode::options()
625 .serialize(&vote_accounts_hash_map)
626 .unwrap();
627 let vote_accounts: VoteAccounts = bincode::options().deserialize(&data).unwrap();
628 assert_eq!(*vote_accounts.vote_accounts, vote_accounts_hash_map);
629 }
630
631 #[test]
632 fn test_vote_accounts_deserialize_invalid_account() {
633 let mut rng = rand::thread_rng();
634 let mut vote_accounts_hash_map = HashMap::<Pubkey, (u64, AccountSharedData)>::new();
637
638 let (valid_account, _) = new_rand_vote_account(&mut rng, None);
639 vote_accounts_hash_map.insert(Pubkey::new_unique(), (0xAA, valid_account.clone()));
640
641 let invalid_account_data =
643 AccountSharedData::new_data(42, &vec![0xFF; 42], &solana_sdk::vote::program::id())
644 .unwrap();
645 vote_accounts_hash_map.insert(Pubkey::new_unique(), (0xBB, invalid_account_data));
646
647 let invalid_account_key =
649 AccountSharedData::new_data(42, &valid_account.data().to_vec(), &Pubkey::new_unique())
650 .unwrap();
651 vote_accounts_hash_map.insert(Pubkey::new_unique(), (0xCC, invalid_account_key));
652
653 let data = bincode::serialize(&vote_accounts_hash_map).unwrap();
654 let options = bincode::options()
655 .with_fixint_encoding()
656 .allow_trailing_bytes();
657 let mut deserializer = bincode::de::Deserializer::from_slice(&data, options);
658 let vote_accounts = deserialize_accounts_hash_map(&mut deserializer).unwrap();
659
660 assert_eq!(vote_accounts.len(), 1);
661 let (stake, _account) = vote_accounts.values().next().unwrap();
662 assert_eq!(*stake, 0xAA);
663 }
664
665 #[test]
666 fn test_staked_nodes() {
667 let mut rng = rand::thread_rng();
668 let mut accounts: Vec<_> = new_rand_vote_accounts(&mut rng, 64).take(1024).collect();
669 let mut vote_accounts = VoteAccounts::default();
670 for (k, (pubkey, (stake, vote_account))) in accounts.iter().enumerate() {
672 vote_accounts.insert(*pubkey, vote_account.clone(), || *stake);
673 if (k + 1) % 128 == 0 {
674 assert_eq!(
675 staked_nodes(&accounts[..k + 1]),
676 *vote_accounts.staked_nodes()
677 );
678 }
679 }
680 for k in 0..256 {
682 let index = rng.gen_range(0..accounts.len());
683 let (pubkey, (_, _)) = accounts.swap_remove(index);
684 vote_accounts.remove(&pubkey);
685 if (k + 1) % 32 == 0 {
686 assert_eq!(staked_nodes(&accounts), *vote_accounts.staked_nodes());
687 }
688 }
689 for k in 0..2048 {
691 let index = rng.gen_range(0..accounts.len());
692 let (pubkey, (stake, _)) = &mut accounts[index];
693 let new_stake = rng.gen_range(0..997);
694 if new_stake < *stake {
695 vote_accounts.sub_stake(pubkey, *stake - new_stake);
696 } else {
697 vote_accounts.add_stake(pubkey, new_stake - *stake);
698 }
699 *stake = new_stake;
700 if (k + 1) % 128 == 0 {
701 assert_eq!(staked_nodes(&accounts), *vote_accounts.staked_nodes());
702 }
703 }
704 while !accounts.is_empty() {
706 let index = rng.gen_range(0..accounts.len());
707 let (pubkey, (_, _)) = accounts.swap_remove(index);
708 vote_accounts.remove(&pubkey);
709 if accounts.len() % 32 == 0 {
710 assert_eq!(staked_nodes(&accounts), *vote_accounts.staked_nodes());
711 }
712 }
713 assert!(vote_accounts.staked_nodes.get().unwrap().is_empty());
714 }
715
716 #[test]
717 fn test_staked_nodes_update() {
718 let mut vote_accounts = VoteAccounts::default();
719
720 let mut rng = rand::thread_rng();
721 let pubkey = Pubkey::new_unique();
722 let node_pubkey = Pubkey::new_unique();
723 let (account1, _) = new_rand_vote_account(&mut rng, Some(node_pubkey));
724 let vote_account1 = VoteAccount::try_from(account1).unwrap();
725
726 let ret = vote_accounts.insert(pubkey, vote_account1.clone(), || 42);
728 assert_eq!(ret, None);
729 assert_eq!(vote_accounts.get_delegated_stake(&pubkey), 42);
730 assert_eq!(vote_accounts.staked_nodes().get(&node_pubkey), Some(&42));
731
732 let ret = vote_accounts.insert(pubkey, vote_account1.clone(), || {
734 panic!("should not be called")
735 });
736 assert_eq!(ret, Some(vote_account1.clone()));
737 assert_eq!(vote_accounts.get(&pubkey), Some(&vote_account1));
738 assert_eq!(vote_accounts.get_delegated_stake(&pubkey), 42);
740 assert_eq!(vote_accounts.staked_nodes().get(&node_pubkey), Some(&42));
741
742 let (account2, _) = new_rand_vote_account(&mut rng, Some(node_pubkey));
744 let vote_account2 = VoteAccount::try_from(account2).unwrap();
745 let ret = vote_accounts.insert(pubkey, vote_account2.clone(), || {
746 panic!("should not be called")
747 });
748 assert_eq!(ret, Some(vote_account1.clone()));
749 assert_eq!(vote_accounts.get(&pubkey), Some(&vote_account2));
750 assert_eq!(vote_accounts.get_delegated_stake(&pubkey), 42);
752 assert_eq!(vote_accounts.staked_nodes().get(&node_pubkey), Some(&42));
753
754 let new_node_pubkey = Pubkey::new_unique();
756 let (account3, _) = new_rand_vote_account(&mut rng, Some(new_node_pubkey));
757 let vote_account3 = VoteAccount::try_from(account3).unwrap();
758 let ret = vote_accounts.insert(pubkey, vote_account3.clone(), || {
759 panic!("should not be called")
760 });
761 assert_eq!(ret, Some(vote_account2.clone()));
762 assert_eq!(vote_accounts.staked_nodes().get(&node_pubkey), None);
763 assert_eq!(
764 vote_accounts.staked_nodes().get(&new_node_pubkey),
765 Some(&42)
766 );
767 }
768
769 #[test]
770 fn test_staked_nodes_zero_stake() {
771 let mut vote_accounts = VoteAccounts::default();
772
773 let mut rng = rand::thread_rng();
774 let pubkey = Pubkey::new_unique();
775 let node_pubkey = Pubkey::new_unique();
776 let (account1, _) = new_rand_vote_account(&mut rng, Some(node_pubkey));
777 let vote_account1 = VoteAccount::try_from(account1).unwrap();
778
779 assert!(vote_accounts.staked_nodes().is_empty());
781 let ret = vote_accounts.insert(pubkey, vote_account1.clone(), || 0);
782 assert_eq!(ret, None);
783 assert_eq!(vote_accounts.get_delegated_stake(&pubkey), 0);
784 assert_eq!(vote_accounts.staked_nodes().get(&node_pubkey), None);
786
787 let new_node_pubkey = Pubkey::new_unique();
789 let (account2, _) = new_rand_vote_account(&mut rng, Some(new_node_pubkey));
790 let vote_account2 = VoteAccount::try_from(account2).unwrap();
791 let ret = vote_accounts.insert(pubkey, vote_account2.clone(), || {
792 panic!("should not be called")
793 });
794 assert_eq!(ret, Some(vote_account1));
795 assert_eq!(vote_accounts.get_delegated_stake(&pubkey), 0);
796 assert_eq!(vote_accounts.staked_nodes().get(&node_pubkey), None);
797 assert_eq!(vote_accounts.staked_nodes().get(&new_node_pubkey), None);
798 }
799
800 #[test]
802 fn test_staked_nodes_cow() {
803 let mut rng = rand::thread_rng();
804 let mut accounts = new_rand_vote_accounts(&mut rng, 64);
805 let mut vote_accounts = VoteAccounts::default();
807 for (pubkey, (stake, vote_account)) in (&mut accounts).take(1024) {
808 vote_accounts.insert(pubkey, vote_account, || stake);
809 }
810 let staked_nodes = vote_accounts.staked_nodes();
811 let (pubkey, (more_stake, vote_account)) =
812 accounts.find(|(_, (stake, _))| *stake != 0).unwrap();
813 let node_pubkey = *vote_account.node_pubkey();
814 vote_accounts.insert(pubkey, vote_account, || more_stake);
815 assert_ne!(staked_nodes, vote_accounts.staked_nodes());
816 assert_eq!(
817 vote_accounts.staked_nodes()[&node_pubkey],
818 more_stake + staked_nodes.get(&node_pubkey).copied().unwrap_or_default()
819 );
820 for (pubkey, stake) in vote_accounts.staked_nodes().iter() {
821 if pubkey != &node_pubkey {
822 assert_eq!(*stake, staked_nodes[pubkey]);
823 } else {
824 assert_eq!(
825 *stake,
826 more_stake + staked_nodes.get(pubkey).copied().unwrap_or_default()
827 );
828 }
829 }
830 }
831
832 #[test]
834 fn test_vote_accounts_cow() {
835 let mut rng = rand::thread_rng();
836 let mut accounts = new_rand_vote_accounts(&mut rng, 64);
837 let mut vote_accounts = VoteAccounts::default();
839 for (pubkey, (stake, vote_account)) in (&mut accounts).take(1024) {
840 vote_accounts.insert(pubkey, vote_account, || stake);
841 }
842 let vote_accounts_hashmap = Arc::<VoteAccountsHashMap>::from(&vote_accounts);
843 assert_eq!(vote_accounts_hashmap, vote_accounts.vote_accounts);
844 assert!(Arc::ptr_eq(
845 &vote_accounts_hashmap,
846 &vote_accounts.vote_accounts
847 ));
848 let (pubkey, (more_stake, vote_account)) =
849 accounts.find(|(_, (stake, _))| *stake != 0).unwrap();
850 vote_accounts.insert(pubkey, vote_account.clone(), || more_stake);
851 assert!(!Arc::ptr_eq(
852 &vote_accounts_hashmap,
853 &vote_accounts.vote_accounts
854 ));
855 assert_ne!(vote_accounts_hashmap, vote_accounts.vote_accounts);
856 let other = (more_stake, vote_account);
857 for (pk, value) in vote_accounts.vote_accounts.iter() {
858 if *pk != pubkey {
859 assert_eq!(value, &vote_accounts_hashmap[pk]);
860 } else {
861 assert_eq!(value, &other);
862 }
863 }
864 }
865}