solana_vote/
vote_account.rs

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

solana_vote/
vote_account.rs
