solana_vote/
vote_account.rs

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(/*owner:*/ 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, (/*stake:*/ 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    // Inner Arc is meant to implement copy-on-write semantics.
51    #[serde(skip)]
52    staked_nodes: OnceLock<
53        Arc<
54            HashMap<
55                Pubkey, // VoteAccount.vote_state.node_pubkey.
56                u64,    // Total stake across all vote-accounts.
57            >,
58        >,
59    >,
60}
61
62impl Clone for VoteAccounts {
63    fn clone(&self) -> Self {
64        Self {
65            vote_accounts: Arc::clone(&self.vote_accounts),
66            // Reset this so that if the previous bank did compute `staked_nodes`, the new bank
67            // won't copy-on-write and keep updating the map if the staked nodes on this bank are
68            // never accessed. See [`VoteAccounts::add_stake`] [`VoteAccounts::sub_stake`] and
69            // [`VoteAccounts::staked_nodes`].
70            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    /// VoteState.node_pubkey of this vote-account.
93    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(), // lamports
125            &VoteStateVersions::new_current(vote_state.clone()),
126            &solana_sdk::vote::program::id(), // owner
127        )
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</*node_pubkey:*/ Pubkey, /*stake:*/ 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                // This is an upsert, we need to update the vote state and move the stake if needed.
200                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                        // The node keys have changed, we move the stake from the old node to the
207                        // new one
208                        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                // Update the vote state
214                Some(mem::replace(old_vote_account, new_vote_account))
215            }
216            Entry::Vacant(entry) => {
217                // This is a new vote account. We don't know the stake yet, so we need to compute it.
218                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        // Allocate as Arc<MaybeUninit<VoteAccountInner>> so we can initialize in place.
333        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        // Safety:
339        // - All the addr_of_mut!(...).write(...) calls are valid since we just allocated and so
340        // the field pointers are valid.
341        // - We use write() so that the old values aren't dropped since they're still
342        // uninitialized.
343        unsafe {
344            let vote_state = addr_of_mut!((*inner_ptr).vote_state);
345            // Safety:
346            // - vote_state is non-null and MaybeUninit<VoteState> is guaranteed to have same layout
347            // and alignment as VoteState.
348            // - Here it is safe to create a reference to MaybeUninit<VoteState> since the value is
349            // aligned and MaybeUninit<T> is valid for all possible bit values.
350            let vote_state = &mut *(vote_state as *mut MaybeUninit<VoteState>);
351
352            // Try to deserialize in place
353            if let Err(e) = VoteState::deserialize_into_uninit(account.data(), vote_state) {
354                // Safety:
355                // - Deserialization failed so at this point vote_state is uninitialized and must
356                // not be dropped. We're ok since `vote_state` is a subfield of `inner`  which is
357                // still MaybeUninit - which isn't dropped by definition - and so neither are its
358                // subfields.
359                return Err(e.into());
360            }
361
362            // Write the account field which completes the initialization of VoteAccountInner.
363            addr_of_mut!((*inner_ptr).account).write(account);
364
365            // Safety:
366            // - At this point both `inner.vote_state` and `inner.account`` are initialized, so it's safe to
367            // transmute the MaybeUninit<VoteAccountInner> to VoteAccountInner.
368            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, (/*stake:*/ 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
435// This custom deserializer is needed to ensure compatibility at snapshot loading with versions
436// before https://github.com/anza-xyz/agave/pull/2659 which would theoretically allow invalid vote
437// accounts in VoteAccounts.
438//
439// In the (near) future we should remove this custom deserializer and make it a hard error when we
440// find invalid vote accounts in snapshots.
441fn 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(), // lamports
517            &VoteStateVersions::new_current(vote_state.clone()),
518            &solana_sdk::vote::program::id(), // owner
519        )
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, (/*stake:*/ 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 that VoteAccount has the same wire format as Account.
590        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        // we'll populate the map with 1 valid and 2 invalid accounts, then ensure that we only get
635        // the valid one after deserialiation
636        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        // bad data
642        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        // wrong owner
648        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        // Add vote accounts.
671        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        // Remove some of the vote accounts.
681        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        // Modify the stakes for some of the accounts.
690        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        // Remove everything.
705        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        // first insert
727        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        // update with unchanged state
733        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        // stake is unchanged
739        assert_eq!(vote_accounts.get_delegated_stake(&pubkey), 42);
740        assert_eq!(vote_accounts.staked_nodes().get(&node_pubkey), Some(&42));
741
742        // update with changed state, same node pubkey
743        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        // stake is unchanged
751        assert_eq!(vote_accounts.get_delegated_stake(&pubkey), 42);
752        assert_eq!(vote_accounts.staked_nodes().get(&node_pubkey), Some(&42));
753
754        // update with new node pubkey, stake must be moved
755        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        // we call this here to initialize VoteAccounts::staked_nodes which is a OnceLock
780        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        // ensure that we didn't add a 0 stake entry to staked_nodes
785        assert_eq!(vote_accounts.staked_nodes().get(&node_pubkey), None);
786
787        // update with new node pubkey, stake is 0 and should remain 0
788        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    // Asserts that returned staked-nodes are copy-on-write references.
801    #[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        // Add vote accounts.
806        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    // Asserts that returned vote-accounts are copy-on-write references.
833    #[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        // Add vote accounts.
838        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}