1#[cfg(feature = "dev-context-only-utils")]
2use qualifier_attr::qualifiers;
3use {
4 crate::{stake_history::StakeHistory, stakes::SerdeStakesToStakeFormat},
5 serde::{
6 Deserialize, Deserializer, Serialize, Serializer,
7 de::{SeqAccess, Visitor},
8 },
9 solana_bls_signatures::pubkey::{
10 PopVerified, PubkeyAffine as BLSPubkeyAffine, PubkeyCompressed as BLSPubkeyCompressed,
11 },
12 solana_clock::Epoch,
13 solana_pubkey::Pubkey,
14 solana_stake_interface::state::Stake,
15 solana_vote::vote_account::{VoteAccounts, VoteAccountsHashMap},
16 solana_vote_interface::state::BLS_PUBLIC_KEY_COMPRESSED_SIZE,
17 std::{
18 collections::HashMap,
19 fmt,
20 mem::MaybeUninit,
21 num::NonZero,
22 sync::{Arc, OnceLock},
23 },
24 wincode::{ReadResult, SchemaRead, SchemaWrite, WriteResult, config::Config, io::Reader},
25};
26
27pub type NodeIdToVoteAccounts = HashMap<Pubkey, NodeVoteAccounts>;
28pub type EpochAuthorizedVoters = HashMap<Pubkey, Pubkey>;
29
30#[derive(Clone, Debug)]
33#[cfg_attr(feature = "dev-context-only-utils", derive(PartialEq))]
34pub struct BLSPubkeyStakeEntry {
35 pub vote_account_pubkey: Pubkey,
37 pub node_pubkey: Pubkey,
39 pub bls_pubkey: PopVerified<BLSPubkeyAffine>,
41 pub stake: NonZero<u64>,
43}
44
45#[derive(Clone, Debug)]
50#[cfg_attr(feature = "dev-context-only-utils", derive(PartialEq))]
51pub struct BLSPubkeyToRankMap {
52 vote_pubkey_to_rank: HashMap<Pubkey, u16>,
54 sorted_pubkeys: Vec<BLSPubkeyStakeEntry>,
56 node_pubkey_to_rank: HashMap<Pubkey, u16>,
58 total_stake: NonZero<u64>,
60}
61
62#[cfg(feature = "frozen-abi")]
65impl solana_frozen_abi::abi_example::AbiExample for BLSPubkeyToRankMap {
66 fn example() -> Self {
67 Self {
68 vote_pubkey_to_rank: HashMap::new(),
69 sorted_pubkeys: Vec::new(),
70 total_stake: NonZero::new(1).unwrap(),
71 node_pubkey_to_rank: HashMap::new(),
72 }
73 }
74}
75
76pub(crate) fn bls_pubkey_compressed_bytes_to_bls_pubkey(
77 bls_pubkey_compressed_bytes: [u8; BLS_PUBLIC_KEY_COMPRESSED_SIZE],
78) -> Option<(BLSPubkeyCompressed, PopVerified<BLSPubkeyAffine>)> {
79 let bls_pubkey_compressed: BLSPubkeyCompressed =
80 wincode::deserialize(&bls_pubkey_compressed_bytes).ok()?;
81 let bls_pubkey_affine = BLSPubkeyAffine::try_from(bls_pubkey_compressed).ok()?;
82 let bls_pubkey_pop_verified = unsafe { PopVerified::new_unchecked(bls_pubkey_affine) };
85 Some((bls_pubkey_compressed, bls_pubkey_pop_verified))
86}
87
88impl BLSPubkeyToRankMap {
89 pub fn new(epoch_vote_accounts_hash_map: &VoteAccountsHashMap) -> Self {
90 let mut candidates = Vec::with_capacity(epoch_vote_accounts_hash_map.len());
91 let mut bls_pubkey_counts = HashMap::new();
92 let mut node_pubkey_counts = HashMap::new();
93 for (&vote_account_pubkey, (stake, account)) in epoch_vote_accounts_hash_map {
94 let Some(stake) = NonZero::new(*stake) else {
95 continue;
96 };
97 let node_pubkey = *account.vote_state_view().node_pubkey();
98 let Some((bls_pubkey_compressed, bls_pubkey)) = account
99 .vote_state_view()
100 .bls_pubkey_compressed()
101 .and_then(bls_pubkey_compressed_bytes_to_bls_pubkey)
102 else {
103 continue;
104 };
105 let entry = BLSPubkeyStakeEntry {
106 vote_account_pubkey,
107 node_pubkey,
108 bls_pubkey,
109 stake,
110 };
111 *bls_pubkey_counts.entry(bls_pubkey_compressed).or_insert(0) += 1;
112 *node_pubkey_counts.entry(node_pubkey).or_insert(0) += 1;
113 candidates.push((entry, bls_pubkey_compressed));
114 }
115 let mut keys_stake_entry_with_compressed: Vec<(BLSPubkeyStakeEntry, BLSPubkeyCompressed)> =
116 candidates
117 .into_iter()
118 .filter_map(|(entry, bls_pubkey_compressed)| {
119 (bls_pubkey_counts[&bls_pubkey_compressed] == 1
120 && node_pubkey_counts[&entry.node_pubkey] == 1)
121 .then_some((entry, bls_pubkey_compressed))
122 })
123 .collect();
124 let total_stake = keys_stake_entry_with_compressed
125 .iter()
126 .fold(0u64, |stake, (entry, _)| {
127 stake.saturating_add(entry.stake.get())
128 });
129 let total_stake = NonZero::new(total_stake).expect("total stakes should not be 0");
130 keys_stake_entry_with_compressed.sort_by(
131 |(a_entry, a_pubkey_compressed), (b_entry, b_pubkey_compressed)| {
132 b_entry
133 .stake
134 .cmp(&a_entry.stake)
135 .then(a_pubkey_compressed.cmp(b_pubkey_compressed))
136 },
137 );
138 let mut sorted_pubkeys = Vec::with_capacity(keys_stake_entry_with_compressed.len());
139 let mut vote_pubkey_to_rank_map =
140 HashMap::with_capacity(keys_stake_entry_with_compressed.len());
141 let mut node_pubkey_to_rank =
142 HashMap::with_capacity(keys_stake_entry_with_compressed.len());
143 for (rank, (entry, _bls_pubkey_compressed)) in
144 keys_stake_entry_with_compressed.into_iter().enumerate()
145 {
146 let rank = u16::try_from(rank).expect("BLS validator rank must fit in u16");
147 vote_pubkey_to_rank_map.insert(entry.vote_account_pubkey, rank);
148 node_pubkey_to_rank.insert(entry.node_pubkey, rank);
149 sorted_pubkeys.push(entry);
150 }
151 Self {
152 vote_pubkey_to_rank: vote_pubkey_to_rank_map,
153 sorted_pubkeys,
154 total_stake,
155 node_pubkey_to_rank,
156 }
157 }
158
159 pub fn is_empty(&self) -> bool {
160 self.sorted_pubkeys.is_empty()
161 }
162
163 pub fn len(&self) -> usize {
164 self.sorted_pubkeys.len()
165 }
166
167 pub fn total_stake(&self) -> NonZero<u64> {
168 self.total_stake
169 }
170
171 pub fn get_rank_for_vote_pubkey(&self, vote_pubkey: &Pubkey) -> Option<&u16> {
172 self.vote_pubkey_to_rank.get(vote_pubkey)
173 }
174
175 pub fn get_pubkey_stake_entry(&self, index: usize) -> Option<&BLSPubkeyStakeEntry> {
176 self.sorted_pubkeys.get(index)
177 }
178
179 #[inline]
181 pub fn get_ranked_entry_for_node(
182 &self,
183 node_pubkey: &Pubkey,
184 ) -> Option<(u16, &BLSPubkeyStakeEntry)> {
185 let rank = *self.node_pubkey_to_rank.get(node_pubkey)?;
186 let entry = self.sorted_pubkeys.get(usize::from(rank))?;
187 Some((rank, entry))
188 }
189}
190
191#[cfg_attr(feature = "frozen-abi", derive(AbiExample, StableAbi, StableAbiSample))]
192#[derive(Clone, Serialize, Debug, Deserialize, Default, PartialEq, Eq, SchemaRead, SchemaWrite)]
193pub struct NodeVoteAccounts {
194 pub vote_accounts: Vec<Pubkey>,
195 pub total_stake: u64,
196}
197
198#[cfg_attr(
203 feature = "frozen-abi",
204 derive(Serialize, SchemaWrite, AbiEnumVisitor, StableAbi, StableAbiSample)
205)]
206#[derive(Clone, Debug, Deserialize, SchemaRead)]
207#[cfg_attr(feature = "dev-context-only-utils", qualifiers(pub))]
208pub(crate) enum DeserializableVersionedEpochStakes {
209 Current {
210 #[cfg_attr(
211 feature = "frozen-abi",
212 stable_abi_sample(with = "stable_abi_sample_deserializable_epoch_stakes(rng)")
213 )]
214 stakes: DeserializableEpochStakes,
215 total_stake: u64,
216 node_id_to_vote_accounts: NodeIdToVoteAccounts,
217 epoch_authorized_voters: EpochAuthorizedVoters,
218 },
219}
220
221#[cfg(feature = "frozen-abi")]
224fn stable_abi_sample_deserializable_epoch_stakes(
225 rng: &mut (impl solana_frozen_abi::rand::RngCore + ?Sized),
226) -> DeserializableEpochStakes {
227 use solana_frozen_abi::stable_abi::StableAbi;
228 let epoch = Epoch::random(rng);
229 let vote_accounts = VoteAccounts::random(rng);
230 let stake_history = StakeHistory::random(rng);
231 DeserializableEpochStakes {
232 vote_accounts,
233 _stake_delegations: Vec::new(),
234 _unused: 0,
235 epoch,
236 stake_history,
237 }
238}
239
240#[derive(Clone, Debug, Serialize, SchemaWrite)]
241#[cfg_attr(
242 feature = "frozen-abi",
243 derive(AbiExample, AbiEnumVisitor, StableAbi, StableAbiSample)
244)]
245#[cfg_attr(feature = "dev-context-only-utils", derive(PartialEq))]
246pub enum VersionedEpochStakes {
247 Current {
248 stakes: EpochStakes,
249 total_stake: u64,
251 node_id_to_vote_accounts: Arc<NodeIdToVoteAccounts>,
252 epoch_authorized_voters: Arc<EpochAuthorizedVoters>,
253 #[cfg_attr(feature = "frozen-abi", stable_abi_sample(with = "Default::default()"))]
254 #[serde(skip)]
255 #[wincode(skip)]
256 bls_pubkey_to_rank_map: OnceLock<Arc<BLSPubkeyToRankMap>>,
257 },
258}
259
260impl From<DeserializableVersionedEpochStakes> for VersionedEpochStakes {
261 fn from(epoch_stakes: DeserializableVersionedEpochStakes) -> Self {
262 let DeserializableVersionedEpochStakes::Current {
263 stakes,
264 total_stake,
265 node_id_to_vote_accounts,
266 epoch_authorized_voters,
267 } = epoch_stakes;
268 Self::Current {
269 stakes: stakes.into(),
270 total_stake,
271 node_id_to_vote_accounts: Arc::new(node_id_to_vote_accounts),
272 epoch_authorized_voters: Arc::new(epoch_authorized_voters),
273 bls_pubkey_to_rank_map: OnceLock::new(),
274 }
275 }
276}
277
278impl VersionedEpochStakes {
279 #[cfg_attr(feature = "dev-context-only-utils", qualifiers(pub))]
280 pub(crate) fn new(stakes: SerdeStakesToStakeFormat, leader_schedule_epoch: Epoch) -> Self {
281 let stakes = EpochStakes::from(stakes);
282 let epoch_vote_accounts = stakes.vote_accounts();
283 let (total_stake, node_id_to_vote_accounts, epoch_authorized_voters) =
284 Self::parse_epoch_vote_accounts(epoch_vote_accounts.as_ref(), leader_schedule_epoch);
285 Self::Current {
286 stakes,
287 total_stake,
288 node_id_to_vote_accounts: Arc::new(node_id_to_vote_accounts),
289 epoch_authorized_voters: Arc::new(epoch_authorized_voters),
290 bls_pubkey_to_rank_map: OnceLock::new(),
291 }
292 }
293
294 #[cfg(feature = "dev-context-only-utils")]
295 pub fn new_for_tests(
296 vote_accounts_hash_map: VoteAccountsHashMap,
297 leader_schedule_epoch: Epoch,
298 ) -> Self {
299 Self::new(
300 SerdeStakesToStakeFormat::Account(crate::stakes::Stakes::new_for_tests(
301 0,
302 solana_vote::vote_account::VoteAccounts::from(Arc::new(vote_accounts_hash_map)),
303 imbl::HashMap::default(),
304 )),
305 leader_schedule_epoch,
306 )
307 }
308
309 pub fn stakes(&self) -> &EpochStakes {
310 match self {
311 Self::Current { stakes, .. } => stakes,
312 }
313 }
314
315 pub fn total_stake(&self) -> u64 {
317 match self {
318 Self::Current { total_stake, .. } => *total_stake,
319 }
320 }
321
322 #[cfg(feature = "dev-context-only-utils")]
323 pub fn set_total_stake(&mut self, total_stake: u64) {
324 match self {
325 Self::Current {
326 total_stake: total_stake_field,
327 ..
328 } => {
329 *total_stake_field = total_stake;
330 }
331 }
332 }
333
334 pub fn node_id_to_vote_accounts(&self) -> &Arc<NodeIdToVoteAccounts> {
335 match self {
336 Self::Current {
337 node_id_to_vote_accounts,
338 ..
339 } => node_id_to_vote_accounts,
340 }
341 }
342
343 pub fn node_id_to_stake(&self, node_id: &Pubkey) -> Option<u64> {
344 self.node_id_to_vote_accounts()
345 .get(node_id)
346 .map(|x| x.total_stake)
347 }
348
349 pub fn epoch_authorized_voters(&self) -> &Arc<EpochAuthorizedVoters> {
350 match self {
351 Self::Current {
352 epoch_authorized_voters,
353 ..
354 } => epoch_authorized_voters,
355 }
356 }
357
358 pub fn bls_pubkey_to_rank_map(&self) -> &Arc<BLSPubkeyToRankMap> {
359 match self {
360 Self::Current {
361 bls_pubkey_to_rank_map,
362 ..
363 } => bls_pubkey_to_rank_map.get_or_init(|| {
364 Arc::new(BLSPubkeyToRankMap::new(
365 self.stakes().vote_accounts().as_ref(),
366 ))
367 }),
368 }
369 }
370
371 pub fn vote_account_stake(&self, vote_account: &Pubkey) -> u64 {
373 self.stakes()
374 .vote_accounts()
375 .get_delegated_stake(vote_account)
376 }
377
378 fn parse_epoch_vote_accounts(
379 epoch_vote_accounts: &VoteAccountsHashMap,
380 leader_schedule_epoch: Epoch,
381 ) -> (u64, NodeIdToVoteAccounts, EpochAuthorizedVoters) {
382 let mut node_id_to_vote_accounts: NodeIdToVoteAccounts = HashMap::new();
383 let mut epoch_authorized_voters: EpochAuthorizedVoters = HashMap::new();
384 let mut total_stake: u64 = 0;
385
386 for (key, (stake, account)) in epoch_vote_accounts.iter() {
387 total_stake += *stake;
388
389 if *stake == 0 {
390 continue;
391 }
392
393 let vote_state = account.vote_state_view();
394
395 if let Some(authorized_voter) = vote_state.get_authorized_voter(leader_schedule_epoch) {
396 let node_vote_accounts = node_id_to_vote_accounts
397 .entry(*vote_state.node_pubkey())
398 .or_default();
399
400 node_vote_accounts.total_stake += stake;
401 node_vote_accounts.vote_accounts.push(*key);
402
403 epoch_authorized_voters.insert(*key, *authorized_voter);
404 }
405 }
406
407 (
408 total_stake,
409 node_id_to_vote_accounts,
410 epoch_authorized_voters,
411 )
412 }
413}
414
415#[derive(Clone, Debug, Default)]
417#[cfg_attr(feature = "frozen-abi", derive(AbiExample, StableAbi, StableAbiSample))]
418#[cfg_attr(feature = "dev-context-only-utils", derive(PartialEq))]
419pub struct EpochStakes {
420 epoch: Epoch,
421 vote_accounts: VoteAccounts,
422 stake_history: StakeHistory,
423}
424
425impl EpochStakes {
426 pub fn vote_accounts(&self) -> &VoteAccounts {
427 &self.vote_accounts
428 }
429 pub fn staked_nodes(&self) -> Arc<HashMap<Pubkey, u64>> {
430 self.vote_accounts.staked_nodes()
431 }
432}
433
434#[derive(Serialize, SchemaWrite)]
440struct SerializableEpochStakes<'a> {
441 vote_accounts: &'a VoteAccounts,
442 stake_delegations: Vec<(Pubkey, Stake)>,
443 unused: u64,
444 epoch: Epoch,
445 stake_history: &'a StakeHistory,
446}
447
448impl<'a> From<&'a EpochStakes> for SerializableEpochStakes<'a> {
449 fn from(epoch_stakes: &'a EpochStakes) -> Self {
450 Self {
451 vote_accounts: &epoch_stakes.vote_accounts,
452 stake_delegations: Vec::new(), unused: 0,
454 epoch: epoch_stakes.epoch,
455 stake_history: &epoch_stakes.stake_history,
456 }
457 }
458}
459
460impl Serialize for EpochStakes {
461 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
462 where
463 S: Serializer,
464 {
465 SerializableEpochStakes::from(self).serialize(serializer)
466 }
467}
468
469unsafe impl<C: wincode::config::Config> SchemaWrite<C> for EpochStakes {
472 type Src = Self;
473
474 fn size_of(src: &Self::Src) -> WriteResult<usize> {
475 <SerializableEpochStakes<'_> as SchemaWrite<C>>::size_of(&src.into())
476 }
477
478 fn write(writer: impl wincode::io::Writer, src: &Self::Src) -> WriteResult<()> {
479 <SerializableEpochStakes<'_> as SchemaWrite<C>>::write(writer, &src.into())
480 }
481}
482
483impl From<SerdeStakesToStakeFormat> for EpochStakes {
484 fn from(stakes: SerdeStakesToStakeFormat) -> Self {
485 let (epoch, vote_accounts, stake_history) = match stakes {
486 SerdeStakesToStakeFormat::Stake(stakes) => stakes.into_epoch_stakes_fields(),
487 SerdeStakesToStakeFormat::Account(stakes) => stakes.into_epoch_stakes_fields(),
488 };
489 Self {
490 epoch,
491 vote_accounts,
492 stake_history,
493 }
494 }
495}
496
497#[cfg_attr(feature = "frozen-abi", derive(Serialize, SchemaWrite))]
503#[derive(Clone, Debug, Deserialize, SchemaRead)]
504#[cfg_attr(feature = "dev-context-only-utils", qualifiers(pub))]
505pub(crate) struct DeserializableEpochStakes {
506 vote_accounts: VoteAccounts,
507 #[serde(deserialize_with = "deserialize_and_ignore_stake_delegations")]
509 #[wincode(with = "IgnoredStakeDelegations")]
510 _stake_delegations: Vec<(Pubkey, Stake)>,
511 _unused: u64,
512 epoch: Epoch,
513 stake_history: StakeHistory,
514}
515
516struct IgnoredStakeDelegations;
520
521unsafe impl<'de, C: Config> SchemaRead<'de, C> for IgnoredStakeDelegations {
522 type Dst = Vec<(Pubkey, Stake)>;
523
524 fn read(reader: impl Reader<'de>, dst: &mut MaybeUninit<Self::Dst>) -> ReadResult<()> {
525 <Vec<(Pubkey, Stake)> as SchemaRead<'de, C>>::get(reader)?;
526 dst.write(Vec::new());
527 Ok(())
528 }
529}
530
531#[cfg(feature = "frozen-abi")]
532unsafe impl<C: Config> SchemaWrite<C> for IgnoredStakeDelegations {
533 type Src = Vec<(Pubkey, Stake)>;
534
535 const TYPE_META: wincode::TypeMeta = <Vec<(Pubkey, Stake)> as SchemaWrite<C>>::TYPE_META;
536
537 fn size_of(src: &Self::Src) -> WriteResult<usize> {
538 <Vec<(Pubkey, Stake)> as SchemaWrite<C>>::size_of(src)
539 }
540
541 fn write(writer: impl wincode::io::Writer, src: &Self::Src) -> WriteResult<()> {
542 <Vec<(Pubkey, Stake)> as SchemaWrite<C>>::write(writer, src)
543 }
544}
545
546impl From<DeserializableEpochStakes> for EpochStakes {
547 fn from(stakes: DeserializableEpochStakes) -> Self {
548 let DeserializableEpochStakes {
549 vote_accounts,
550 _stake_delegations: _,
551 _unused: _,
552 epoch,
553 stake_history,
554 } = stakes;
555 Self {
556 epoch,
557 vote_accounts,
558 stake_history,
559 }
560 }
561}
562
563fn deserialize_and_ignore_stake_delegations<'de, D>(
567 deserializer: D,
568) -> Result<Vec<(Pubkey, Stake)>, D::Error>
569where
570 D: Deserializer<'de>,
571{
572 struct IgnoredStakeDelegationsVisitor;
573
574 impl<'de> Visitor<'de> for IgnoredStakeDelegationsVisitor {
575 type Value = Vec<(Pubkey, Stake)>;
576
577 fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
578 formatter.write_str("a sequence of serialized stake delegations")
579 }
580
581 fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
582 where
583 A: SeqAccess<'de>,
584 {
585 while seq.next_element::<(Pubkey, Stake)>()?.is_some() {
586 }
588 Ok(Vec::new())
589 }
590 }
591
592 deserializer.deserialize_seq(IgnoredStakeDelegationsVisitor)
593}
594
595#[cfg(test)]
596pub(crate) mod tests {
597 use {
598 super::*,
599 crate::{
600 serde_snapshot::deserialize_wincode_from, stake_account::StakeAccount, stakes::Stakes,
601 },
602 solana_account::AccountSharedData,
603 solana_bls_signatures::keypair::Keypair as BLSKeypair,
604 solana_rent::Rent,
605 solana_vote::vote_account::VoteAccount,
606 solana_vote_program::vote_state::create_v4_account_with_authorized,
607 std::iter,
608 test_case::test_case,
609 };
610
611 struct VoteAccountInfo {
612 vote_account: Pubkey,
613 account: AccountSharedData,
614 authorized_voter: Pubkey,
615 }
616
617 fn new_vote_accounts(
618 num_nodes: usize,
619 num_vote_accounts_per_node: usize,
620 is_alpenglow: bool,
621 ) -> HashMap<Pubkey, Vec<VoteAccountInfo>> {
622 (0..num_nodes)
624 .map(|_| {
625 let node_id = solana_pubkey::new_rand();
626 (
627 node_id,
628 iter::repeat_with(|| {
629 let authorized_voter = solana_pubkey::new_rand();
630 let bls_pubkey_compressed: BLSPubkeyCompressed =
631 (*BLSKeypair::new().public).into();
632 let bls_pubkey_compressed_serialized =
633 wincode::serialize(&bls_pubkey_compressed)
634 .unwrap()
635 .try_into()
636 .unwrap();
637
638 let bls_pubkey = if is_alpenglow {
639 bls_pubkey_compressed_serialized
640 } else {
641 [0u8; BLS_PUBLIC_KEY_COMPRESSED_SIZE]
642 };
643 let account = create_v4_account_with_authorized(
644 &node_id,
645 &authorized_voter,
646 bls_pubkey,
647 &node_id,
648 0,
649 &node_id,
650 0,
651 &node_id,
652 100,
653 );
654 VoteAccountInfo {
655 vote_account: solana_pubkey::new_rand(),
656 account,
657 authorized_voter,
658 }
659 })
660 .take(num_vote_accounts_per_node)
661 .collect(),
662 )
663 })
664 .collect()
665 }
666
667 fn new_epoch_vote_accounts(
668 vote_accounts_map: &HashMap<Pubkey, Vec<VoteAccountInfo>>,
669 node_id_to_stake_fn: impl Fn(&Pubkey) -> u64,
670 ) -> VoteAccountsHashMap {
671 vote_accounts_map
673 .iter()
674 .flat_map(|(node_id, vote_accounts)| {
675 vote_accounts.iter().map(|v| {
676 let vote_account = VoteAccount::try_from(v.account.clone()).unwrap();
677 (v.vote_account, (node_id_to_stake_fn(node_id), vote_account))
678 })
679 })
680 .collect()
681 }
682
683 #[test_case(true; "alpenglow")]
684 #[test_case(false; "towerbft")]
685 fn test_parse_epoch_vote_accounts(is_alpenglow: bool) {
686 let stake_per_account = 100;
687 let num_vote_accounts_per_node = 2;
688 let num_nodes = 10;
689
690 let vote_accounts_map =
691 new_vote_accounts(num_nodes, num_vote_accounts_per_node, is_alpenglow);
692
693 let expected_authorized_voters: HashMap<_, _> = vote_accounts_map
694 .values()
695 .flat_map(|vote_accounts| {
696 vote_accounts
697 .iter()
698 .map(|v| (v.vote_account, v.authorized_voter))
699 })
700 .collect();
701
702 let expected_node_id_to_vote_accounts: HashMap<_, _> = vote_accounts_map
703 .iter()
704 .map(|(node_pubkey, vote_accounts)| {
705 let mut vote_accounts = vote_accounts
706 .iter()
707 .map(|v| v.vote_account)
708 .collect::<Vec<_>>();
709 vote_accounts.sort();
710 let node_vote_accounts = NodeVoteAccounts {
711 vote_accounts,
712 total_stake: stake_per_account * num_vote_accounts_per_node as u64,
713 };
714 (*node_pubkey, node_vote_accounts)
715 })
716 .collect();
717
718 let epoch_vote_accounts =
719 new_epoch_vote_accounts(&vote_accounts_map, |_| stake_per_account);
720
721 let (total_stake, mut node_id_to_vote_accounts, epoch_authorized_voters) =
722 VersionedEpochStakes::parse_epoch_vote_accounts(&epoch_vote_accounts, 0);
723
724 node_id_to_vote_accounts
726 .iter_mut()
727 .for_each(|(_, node_vote_accounts)| node_vote_accounts.vote_accounts.sort());
728
729 assert!(
730 node_id_to_vote_accounts.len() == expected_node_id_to_vote_accounts.len()
731 && node_id_to_vote_accounts
732 .iter()
733 .all(|(k, v)| expected_node_id_to_vote_accounts.get(k).unwrap() == v)
734 );
735 assert!(
736 epoch_authorized_voters.len() == expected_authorized_voters.len()
737 && epoch_authorized_voters
738 .iter()
739 .all(|(k, v)| expected_authorized_voters.get(k).unwrap() == v)
740 );
741 assert_eq!(
742 total_stake,
743 num_nodes as u64 * num_vote_accounts_per_node as u64 * 100
744 );
745 }
746
747 #[test_case(true; "alpenglow")]
748 #[test_case(false; "towerbft")]
749 fn test_node_id_to_stake(is_alpenglow: bool) {
750 let num_nodes = 10;
751 let num_vote_accounts_per_node = 2;
752
753 let vote_accounts_map =
754 new_vote_accounts(num_nodes, num_vote_accounts_per_node, is_alpenglow);
755 let node_id_to_stake_map = vote_accounts_map
756 .keys()
757 .enumerate()
758 .map(|(index, node_id)| (*node_id, ((index + 1) * 100) as u64))
759 .collect::<HashMap<_, _>>();
760 let epoch_vote_accounts = new_epoch_vote_accounts(&vote_accounts_map, |node_id| {
761 *node_id_to_stake_map.get(node_id).unwrap()
762 });
763 let epoch_stakes = VersionedEpochStakes::new_for_tests(epoch_vote_accounts, 0);
764
765 assert_eq!(epoch_stakes.total_stake(), 11000);
766 for (node_id, stake) in node_id_to_stake_map.iter() {
767 assert_eq!(
768 epoch_stakes.node_id_to_stake(node_id),
769 Some(*stake * num_vote_accounts_per_node as u64)
770 );
771 }
772 }
773
774 #[test]
775 fn test_bls_pubkey_rank_map() {
776 agave_logger::setup();
777 let num_nodes = 10;
778 let num_vote_accounts = num_nodes;
779
780 let vote_accounts_map = new_vote_accounts(num_nodes, 1, true);
781 let node_id_to_stake_map = vote_accounts_map
782 .keys()
783 .enumerate()
784 .map(|(index, node_id)| (*node_id, ((index + 1) * 100) as u64))
785 .collect::<HashMap<_, _>>();
786 let epoch_vote_accounts = new_epoch_vote_accounts(&vote_accounts_map, |node_id| {
787 *node_id_to_stake_map.get(node_id).unwrap()
788 });
789 let epoch_stakes = VersionedEpochStakes::new_for_tests(epoch_vote_accounts.clone(), 0);
790 let bls_pubkey_to_rank_map = epoch_stakes.bls_pubkey_to_rank_map();
791 let expected_num_vote_accounts = num_vote_accounts;
792 assert_eq!(bls_pubkey_to_rank_map.len(), expected_num_vote_accounts);
793 let expected_total_stake = epoch_stakes.total_stake();
794 assert_eq!(
795 bls_pubkey_to_rank_map.total_stake().get(),
796 expected_total_stake
797 );
798 for expected_rank in 0..bls_pubkey_to_rank_map.len() {
799 let expected_entry = bls_pubkey_to_rank_map
800 .get_pubkey_stake_entry(expected_rank)
801 .unwrap();
802 let (rank, entry) = bls_pubkey_to_rank_map
803 .get_ranked_entry_for_node(&expected_entry.node_pubkey)
804 .unwrap();
805 assert_eq!(usize::from(rank), expected_rank);
806 assert!(std::ptr::eq(entry, expected_entry));
807 assert_eq!(
808 bls_pubkey_to_rank_map
809 .get_rank_for_vote_pubkey(&entry.vote_account_pubkey)
810 .copied(),
811 Some(rank)
812 );
813 }
814 assert!(
815 bls_pubkey_to_rank_map
816 .get_ranked_entry_for_node(&Pubkey::new_unique())
817 .is_none()
818 );
819
820 let mut bank_epoch_stakes = HashMap::new();
822 bank_epoch_stakes.insert(0, epoch_stakes.clone());
823 let epoch_stakes = bank_epoch_stakes
824 .get(&0)
825 .expect("Epoch stakes should exist");
826 let bls_pubkey_to_rank_map2 = epoch_stakes.bls_pubkey_to_rank_map();
827 assert_eq!(bls_pubkey_to_rank_map2, bls_pubkey_to_rank_map);
828 }
829
830 #[test]
831 #[should_panic(expected = "total stakes should not be 0")]
832 fn test_multiple_vote_accounts_panics() {
833 agave_logger::setup();
834 let num_nodes = 10;
835
836 let vote_accounts_map = new_vote_accounts(num_nodes, 2, true);
837 let node_id_to_stake_map = vote_accounts_map
838 .keys()
839 .enumerate()
840 .map(|(index, node_id)| (*node_id, ((index + 1) * 100) as u64))
841 .collect::<HashMap<_, _>>();
842 let epoch_vote_accounts = new_epoch_vote_accounts(&vote_accounts_map, |node_id| {
843 *node_id_to_stake_map.get(node_id).unwrap()
844 });
845 let epoch_stakes = VersionedEpochStakes::new_for_tests(epoch_vote_accounts.clone(), 0);
846 epoch_stakes.bls_pubkey_to_rank_map();
847 }
848
849 #[test]
850 fn test_bls_pubkey_rank_map_excludes_duplicate_bls_and_identity() {
851 let new_bls_pubkey = || {
852 let compressed: BLSPubkeyCompressed = (*BLSKeypair::new().public).into();
853 wincode::serialize(&compressed).unwrap().try_into().unwrap()
854 };
855
856 let duplicate_bls_pubkey_serialized = new_bls_pubkey();
857 let duplicate_node_bls_pubkey_serialized = new_bls_pubkey();
858 let duplicate_node_bls_pubkey_serialized_2 = new_bls_pubkey();
859 let shared_voter_bls_pubkey_serialized = new_bls_pubkey();
860 let shared_voter_bls_pubkey_serialized_2 = new_bls_pubkey();
861 let unique_bls_pubkey_serialized = new_bls_pubkey();
862
863 let duplicate_bls_vote_pubkey = Pubkey::new_unique();
864 let duplicate_bls_vote_pubkey_2 = Pubkey::new_unique();
865 let duplicate_node_vote_pubkey = Pubkey::new_unique();
866 let duplicate_node_vote_pubkey_2 = Pubkey::new_unique();
867 let shared_voter_vote_pubkey = Pubkey::new_unique();
868 let shared_voter_vote_pubkey_2 = Pubkey::new_unique();
869 let unique_vote_pubkey = Pubkey::new_unique();
870
871 let duplicate_node_pubkey = Pubkey::new_unique();
872 let shared_authorized_voter = Pubkey::new_unique();
873 let shared_voter_node_pubkey = Pubkey::new_unique();
874 let shared_voter_node_pubkey_2 = Pubkey::new_unique();
875 let unique_node_pubkey = Pubkey::new_unique();
876 let unique_voter = Pubkey::new_unique();
877
878 let account = |node_pubkey, authorized_voter, bls_pubkey| {
879 VoteAccount::try_from(create_v4_account_with_authorized(
880 &node_pubkey,
881 &authorized_voter,
882 bls_pubkey,
883 &node_pubkey,
884 0,
885 &node_pubkey,
886 0,
887 &node_pubkey,
888 100,
889 ))
890 .unwrap()
891 };
892 let epoch_vote_accounts = VoteAccountsHashMap::from([
893 (
894 duplicate_bls_vote_pubkey,
895 (
896 100,
897 account(
898 Pubkey::new_unique(),
899 Pubkey::new_unique(),
900 duplicate_bls_pubkey_serialized,
901 ),
902 ),
903 ),
904 (
905 duplicate_bls_vote_pubkey_2,
906 (
907 100,
908 account(
909 Pubkey::new_unique(),
910 Pubkey::new_unique(),
911 duplicate_bls_pubkey_serialized,
912 ),
913 ),
914 ),
915 (
916 duplicate_node_vote_pubkey,
917 (
918 100,
919 account(
920 duplicate_node_pubkey,
921 Pubkey::new_unique(),
922 duplicate_node_bls_pubkey_serialized,
923 ),
924 ),
925 ),
926 (
927 duplicate_node_vote_pubkey_2,
928 (
929 100,
930 account(
931 duplicate_node_pubkey,
932 Pubkey::new_unique(),
933 duplicate_node_bls_pubkey_serialized_2,
934 ),
935 ),
936 ),
937 (
938 shared_voter_vote_pubkey,
939 (
940 100,
941 account(
942 shared_voter_node_pubkey,
943 shared_authorized_voter,
944 shared_voter_bls_pubkey_serialized,
945 ),
946 ),
947 ),
948 (
949 shared_voter_vote_pubkey_2,
950 (
951 100,
952 account(
953 shared_voter_node_pubkey_2,
954 shared_authorized_voter,
955 shared_voter_bls_pubkey_serialized_2,
956 ),
957 ),
958 ),
959 (
960 unique_vote_pubkey,
961 (
962 50,
963 account(
964 unique_node_pubkey,
965 unique_voter,
966 unique_bls_pubkey_serialized,
967 ),
968 ),
969 ),
970 ]);
971
972 let rank_map = BLSPubkeyToRankMap::new(&epoch_vote_accounts);
973
974 assert_eq!(rank_map.len(), 3);
975 assert_eq!(rank_map.total_stake().get(), 250);
976 for vote_pubkey in [
977 duplicate_bls_vote_pubkey,
978 duplicate_bls_vote_pubkey_2,
979 duplicate_node_vote_pubkey,
980 duplicate_node_vote_pubkey_2,
981 ] {
982 assert!(rank_map.get_rank_for_vote_pubkey(&vote_pubkey).is_none());
983 }
984 assert!(
985 rank_map
986 .get_ranked_entry_for_node(&duplicate_node_pubkey)
987 .is_none()
988 );
989 for node_pubkey in [
990 shared_voter_node_pubkey,
991 shared_voter_node_pubkey_2,
992 unique_node_pubkey,
993 ] {
994 let (rank, entry) = rank_map.get_ranked_entry_for_node(&node_pubkey).unwrap();
995 assert_eq!(entry.node_pubkey, node_pubkey);
996 assert_eq!(
997 rank_map
998 .get_rank_for_vote_pubkey(&entry.vote_account_pubkey)
999 .copied(),
1000 Some(rank)
1001 );
1002 }
1003 assert!(rank_map.get_pubkey_stake_entry(rank_map.len()).is_none());
1004 }
1005
1006 #[test]
1007 fn test_versioned_epoch_stakes_does_not_serialize_delegations() {
1008 #[derive(Deserialize)]
1010 enum SerializedVersionedEpochStakes {
1011 Current {
1012 stakes: SerializedEpochStakes,
1013 total_stake: u64,
1014 node_id_to_vote_accounts: NodeIdToVoteAccounts,
1015 epoch_authorized_voters: EpochAuthorizedVoters,
1016 },
1017 }
1018 #[derive(Deserialize)]
1019 struct SerializedEpochStakes {
1020 vote_accounts: VoteAccounts,
1021 stake_delegations: Vec<(Pubkey, Stake)>,
1022 unused: u64,
1023 epoch: Epoch,
1024 stake_history: StakeHistory,
1025 }
1026
1027 let epoch = 42;
1028 let delegated_amount = 456_789;
1029 let rent = Rent::default();
1030 let ((vote_pubkey, vote_account), (stake_pubkey, stake_account)) =
1031 crate::stakes::tests::create_staked_node_accounts(123, &rent);
1032 let vote_account = VoteAccount::try_from(vote_account).unwrap();
1033 let vote_accounts = VoteAccounts::from(Arc::new(HashMap::from([(
1034 vote_pubkey,
1035 (delegated_amount, vote_account),
1036 )])));
1037 let stake_account = StakeAccount::try_from(stake_account).unwrap();
1038 let stakes = Stakes::new_for_tests(
1039 epoch,
1040 vote_accounts,
1041 imbl::HashMap::from_iter([(stake_pubkey, stake_account)]),
1042 );
1043
1044 assert!(!stakes.stake_delegations().is_empty());
1046
1047 let epoch_stakes = VersionedEpochStakes::new(SerdeStakesToStakeFormat::Account(stakes), 0);
1048
1049 assert_eq!(
1050 epoch_stakes
1051 .stakes()
1052 .vote_accounts()
1053 .get_delegated_stake(&vote_pubkey),
1054 delegated_amount,
1055 );
1056
1057 let serialized_bytes = bincode::serialize(&epoch_stakes).unwrap();
1058 let serialized_epoch_stakes: SerializedVersionedEpochStakes =
1059 bincode::deserialize(&serialized_bytes).unwrap();
1060 match serialized_epoch_stakes {
1061 SerializedVersionedEpochStakes::Current {
1062 stakes,
1063 total_stake,
1064 node_id_to_vote_accounts,
1065 epoch_authorized_voters,
1066 } => {
1067 assert_eq!(
1068 stakes.vote_accounts.get_delegated_stake(&vote_pubkey),
1069 delegated_amount,
1070 );
1071 assert!(stakes.stake_delegations.is_empty()); assert_eq!(stakes.unused, 0);
1073 assert_eq!(stakes.epoch, epoch);
1074 assert_eq!(stakes.stake_history, StakeHistory::default());
1075 assert_eq!(total_stake, delegated_amount);
1076 assert_eq!(node_id_to_vote_accounts.len(), 1);
1077 assert_eq!(epoch_authorized_voters.len(), 1);
1078 }
1079 }
1080
1081 let deserialized_epoch_stakes: VersionedEpochStakes =
1082 deserialize_wincode_from::<_, DeserializableVersionedEpochStakes>(
1083 std::io::Cursor::new(&serialized_bytes),
1084 )
1085 .unwrap()
1086 .into();
1087 assert_eq!(
1088 deserialized_epoch_stakes
1089 .stakes()
1090 .vote_accounts()
1091 .get_delegated_stake(&vote_pubkey),
1092 delegated_amount,
1093 );
1094 }
1095}