solana_vote/
vote_account.rs

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(/*owner:*/ 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, (/*stake:*/ 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    // Inner Arc is meant to implement copy-on-write semantics.
48    #[serde(skip)]
49    staked_nodes: OnceLock<
50        Arc<
51            HashMap<
52                Pubkey, // VoteAccount.vote_state.node_pubkey.
53                u64,    // Total stake across all vote-accounts.
54            >,
55        >,
56    >,
57}
58
59impl Clone for VoteAccounts {
60    fn clone(&self) -> Self {
61        Self {
62            vote_accounts: Arc::clone(&self.vote_accounts),
63            // Reset this so that if the previous bank did compute `staked_nodes`, the new bank
64            // won't copy-on-write and keep updating the map if the staked nodes on this bank are
65            // never accessed. See [`VoteAccounts::add_stake`] [`VoteAccounts::sub_stake`] and
66            // [`VoteAccounts::staked_nodes`].
67            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    /// VoteState.node_pubkey of this vote-account.
90    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(), // lamports
120            &VoteStateVersions::new_current(vote_state.clone()),
121            &solana_sdk_ids::vote::id(), // owner
122        )
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</*node_pubkey:*/ Pubkey, /*stake:*/ 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                // This is an upsert, we need to update the vote state and move the stake if needed.
195                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                        // The node keys have changed, we move the stake from the old node to the
202                        // new one
203                        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                // Update the vote state
209                Some(mem::replace(old_vote_account, new_vote_account))
210            }
211            Entry::Vacant(entry) => {
212                // This is a new vote account. We don't know the stake yet, so we need to compute it.
213                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, (/*stake:*/ 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
400// This custom deserializer is needed to ensure compatibility at snapshot loading with versions
401// before https://github.com/anza-xyz/agave/pull/2659 which would theoretically allow invalid vote
402// accounts in VoteAccounts.
403//
404// In the (near) future we should remove this custom deserializer and make it a hard error when we
405// find invalid vote accounts in snapshots.
406fn 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(), // lamports
480            &VoteStateVersions::new_current(vote_state.clone()),
481            &solana_sdk_ids::vote::id(), // owner
482        )
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, (/*stake:*/ 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 that VoteAccount has the same wire format as Account.
550        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        // we'll populate the map with 1 valid and 2 invalid accounts, then ensure that we only get
595        // the valid one after deserialiation
596        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        // bad data
602        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        // wrong owner
607        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        // Add vote accounts.
630        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        // Remove some of the vote accounts.
640        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        // Modify the stakes for some of the accounts.
649        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        // Remove everything.
664        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        // first insert
686        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        // update with unchanged state
692        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        // stake is unchanged
698        assert_eq!(vote_accounts.get_delegated_stake(&pubkey), 42);
699        assert_eq!(vote_accounts.staked_nodes().get(&node_pubkey), Some(&42));
700
701        // update with changed state, same node pubkey
702        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        // stake is unchanged
710        assert_eq!(vote_accounts.get_delegated_stake(&pubkey), 42);
711        assert_eq!(vote_accounts.staked_nodes().get(&node_pubkey), Some(&42));
712
713        // update with new node pubkey, stake must be moved
714        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        // we call this here to initialize VoteAccounts::staked_nodes which is a OnceLock
739        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        // ensure that we didn't add a 0 stake entry to staked_nodes
744        assert_eq!(vote_accounts.staked_nodes().get(&node_pubkey), None);
745
746        // update with new node pubkey, stake is 0 and should remain 0
747        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    // Asserts that returned staked-nodes are copy-on-write references.
760    #[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        // Add vote accounts.
765        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    // Asserts that returned vote-accounts are copy-on-write references.
792    #[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        // Add vote accounts.
797        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}