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