1use {
2 crate::stakes::SerdeStakesToStakeFormat,
3 serde::{Deserialize, Serialize},
4 solana_bls_signatures::{Pubkey as BLSPubkey, PubkeyCompressed as BLSPubkeyCompressed},
5 solana_clock::Epoch,
6 solana_pubkey::Pubkey,
7 solana_vote::vote_account::VoteAccountsHashMap,
8 std::{
9 collections::HashMap,
10 sync::{Arc, OnceLock},
11 },
12};
13
14pub type NodeIdToVoteAccounts = HashMap<Pubkey, NodeVoteAccounts>;
15pub type EpochAuthorizedVoters = HashMap<Pubkey, Pubkey>;
16
17#[derive(Clone, Debug, Default)]
18#[cfg_attr(feature = "frozen-abi", derive(AbiExample))]
19#[cfg_attr(feature = "dev-context-only-utils", derive(PartialEq))]
20pub struct BLSPubkeyToRankMap {
21 rank_map: HashMap<BLSPubkey, u16>,
22 sorted_pubkeys: Vec<(Pubkey, BLSPubkey)>,
25}
26
27impl BLSPubkeyToRankMap {
28 pub fn new(epoch_vote_accounts_hash_map: &VoteAccountsHashMap) -> Self {
29 let mut pubkey_stake_pair_vec: Vec<(Pubkey, BLSPubkey, u64)> = epoch_vote_accounts_hash_map
30 .iter()
31 .filter_map(|(pubkey, (stake, account))| {
32 if *stake > 0 {
33 account
34 .vote_state_view()
35 .bls_pubkey_compressed()
36 .and_then(|bls_pubkey_compressed_bytes| {
37 let bls_pubkey_compressed =
38 BLSPubkeyCompressed(bls_pubkey_compressed_bytes);
39 BLSPubkey::try_from(bls_pubkey_compressed).ok()
40 })
41 .map(|bls_pubkey| (*pubkey, bls_pubkey, *stake))
42 } else {
43 None
44 }
45 })
46 .collect();
47 pubkey_stake_pair_vec.sort_by(|(_, a_pubkey, a_stake), (_, b_pubkey, b_stake)| {
48 b_stake.cmp(a_stake).then(a_pubkey.cmp(b_pubkey))
49 });
50 let mut sorted_pubkeys = Vec::new();
51 let mut bls_pubkey_to_rank_map = HashMap::new();
52 for (rank, (pubkey, bls_pubkey, _stake)) in pubkey_stake_pair_vec.into_iter().enumerate() {
53 sorted_pubkeys.push((pubkey, bls_pubkey));
54 bls_pubkey_to_rank_map.insert(bls_pubkey, rank as u16);
55 }
56 Self {
57 rank_map: bls_pubkey_to_rank_map,
58 sorted_pubkeys,
59 }
60 }
61
62 pub fn is_empty(&self) -> bool {
63 self.rank_map.is_empty()
64 }
65
66 pub fn len(&self) -> usize {
67 self.rank_map.len()
68 }
69
70 pub fn get_rank(&self, bls_pubkey: &BLSPubkey) -> Option<&u16> {
71 self.rank_map.get(bls_pubkey)
72 }
73
74 pub fn get_pubkey(&self, index: usize) -> Option<&(Pubkey, BLSPubkey)> {
75 self.sorted_pubkeys.get(index)
76 }
77}
78
79#[cfg_attr(feature = "frozen-abi", derive(AbiExample))]
80#[derive(Clone, Serialize, Debug, Deserialize, Default, PartialEq, Eq)]
81pub struct NodeVoteAccounts {
82 pub vote_accounts: Vec<Pubkey>,
83 pub total_stake: u64,
84}
85
86#[derive(Clone, Debug, Serialize, Deserialize)]
87#[cfg_attr(feature = "frozen-abi", derive(AbiExample, AbiEnumVisitor))]
88#[cfg_attr(feature = "dev-context-only-utils", derive(PartialEq))]
89pub enum VersionedEpochStakes {
90 Current {
91 stakes: SerdeStakesToStakeFormat,
92 total_stake: u64,
93 node_id_to_vote_accounts: Arc<NodeIdToVoteAccounts>,
94 epoch_authorized_voters: Arc<EpochAuthorizedVoters>,
95 #[serde(skip)]
96 bls_pubkey_to_rank_map: OnceLock<Arc<BLSPubkeyToRankMap>>,
97 },
98}
99
100impl VersionedEpochStakes {
101 pub(crate) fn new(stakes: SerdeStakesToStakeFormat, leader_schedule_epoch: Epoch) -> Self {
102 let epoch_vote_accounts = stakes.vote_accounts();
103 let (total_stake, node_id_to_vote_accounts, epoch_authorized_voters) =
104 Self::parse_epoch_vote_accounts(epoch_vote_accounts.as_ref(), leader_schedule_epoch);
105 Self::Current {
106 stakes,
107 total_stake,
108 node_id_to_vote_accounts: Arc::new(node_id_to_vote_accounts),
109 epoch_authorized_voters: Arc::new(epoch_authorized_voters),
110 bls_pubkey_to_rank_map: OnceLock::new(),
111 }
112 }
113
114 #[cfg(feature = "dev-context-only-utils")]
115 pub fn new_for_tests(
116 vote_accounts_hash_map: VoteAccountsHashMap,
117 leader_schedule_epoch: Epoch,
118 ) -> Self {
119 Self::new(
120 SerdeStakesToStakeFormat::Account(crate::stakes::Stakes::new_for_tests(
121 0,
122 solana_vote::vote_account::VoteAccounts::from(Arc::new(vote_accounts_hash_map)),
123 im::HashMap::default(),
124 )),
125 leader_schedule_epoch,
126 )
127 }
128
129 pub fn stakes(&self) -> &SerdeStakesToStakeFormat {
130 match self {
131 Self::Current { stakes, .. } => stakes,
132 }
133 }
134
135 pub fn total_stake(&self) -> u64 {
136 match self {
137 Self::Current { total_stake, .. } => *total_stake,
138 }
139 }
140
141 #[cfg(feature = "dev-context-only-utils")]
142 pub fn set_total_stake(&mut self, total_stake: u64) {
143 match self {
144 Self::Current {
145 total_stake: total_stake_field,
146 ..
147 } => {
148 *total_stake_field = total_stake;
149 }
150 }
151 }
152
153 pub fn node_id_to_vote_accounts(&self) -> &Arc<NodeIdToVoteAccounts> {
154 match self {
155 Self::Current {
156 node_id_to_vote_accounts,
157 ..
158 } => node_id_to_vote_accounts,
159 }
160 }
161
162 pub fn node_id_to_stake(&self, node_id: &Pubkey) -> Option<u64> {
163 self.node_id_to_vote_accounts()
164 .get(node_id)
165 .map(|x| x.total_stake)
166 }
167
168 pub fn epoch_authorized_voters(&self) -> &Arc<EpochAuthorizedVoters> {
169 match self {
170 Self::Current {
171 epoch_authorized_voters,
172 ..
173 } => epoch_authorized_voters,
174 }
175 }
176
177 pub fn bls_pubkey_to_rank_map(&self) -> &Arc<BLSPubkeyToRankMap> {
178 match self {
179 Self::Current {
180 bls_pubkey_to_rank_map,
181 ..
182 } => bls_pubkey_to_rank_map.get_or_init(|| {
183 Arc::new(BLSPubkeyToRankMap::new(
184 self.stakes().vote_accounts().as_ref(),
185 ))
186 }),
187 }
188 }
189
190 pub fn vote_account_stake(&self, vote_account: &Pubkey) -> u64 {
191 self.stakes()
192 .vote_accounts()
193 .get_delegated_stake(vote_account)
194 }
195
196 fn parse_epoch_vote_accounts(
197 epoch_vote_accounts: &VoteAccountsHashMap,
198 leader_schedule_epoch: Epoch,
199 ) -> (u64, NodeIdToVoteAccounts, EpochAuthorizedVoters) {
200 let mut node_id_to_vote_accounts: NodeIdToVoteAccounts = HashMap::new();
201 let total_stake = epoch_vote_accounts
202 .iter()
203 .map(|(_, (stake, _))| stake)
204 .sum();
205 let epoch_authorized_voters = epoch_vote_accounts
206 .iter()
207 .filter_map(|(key, (stake, account))| {
208 let vote_state = account.vote_state_view();
209
210 if *stake > 0 {
211 if let Some(authorized_voter) =
212 vote_state.get_authorized_voter(leader_schedule_epoch)
213 {
214 let node_vote_accounts = node_id_to_vote_accounts
215 .entry(*vote_state.node_pubkey())
216 .or_default();
217
218 node_vote_accounts.total_stake += stake;
219 node_vote_accounts.vote_accounts.push(*key);
220
221 Some((*key, *authorized_voter))
222 } else {
223 None
224 }
225 } else {
226 None
227 }
228 })
229 .collect();
230 (
231 total_stake,
232 node_id_to_vote_accounts,
233 epoch_authorized_voters,
234 )
235 }
236}
237
238#[cfg(test)]
239pub(crate) mod tests {
240 use {
241 super::*, solana_account::AccountSharedData,
242 solana_bls_signatures::keypair::Keypair as BLSKeypair,
243 solana_vote::vote_account::VoteAccount,
244 solana_vote_program::vote_state::create_v4_account_with_authorized, std::iter,
245 test_case::test_case,
246 };
247
248 struct VoteAccountInfo {
249 vote_account: Pubkey,
250 account: AccountSharedData,
251 authorized_voter: Pubkey,
252 }
253
254 fn new_vote_accounts(
255 num_nodes: usize,
256 num_vote_accounts_per_node: usize,
257 is_alpenglow: bool,
258 ) -> HashMap<Pubkey, Vec<VoteAccountInfo>> {
259 (0..num_nodes)
261 .map(|_| {
262 let node_id = solana_pubkey::new_rand();
263 (
264 node_id,
265 iter::repeat_with(|| {
266 let authorized_voter = solana_pubkey::new_rand();
267 let bls_pubkey_compressed: BLSPubkeyCompressed =
268 BLSKeypair::new().public.try_into().unwrap();
269 let bls_pubkey_compressed_serialized =
270 bincode::serialize(&bls_pubkey_compressed)
271 .unwrap()
272 .try_into()
273 .unwrap();
274
275 let account = if is_alpenglow {
276 create_v4_account_with_authorized(
277 &node_id,
278 &authorized_voter,
279 &node_id,
280 Some(bls_pubkey_compressed_serialized),
281 0,
282 100,
283 )
284 } else {
285 create_v4_account_with_authorized(
286 &node_id,
287 &authorized_voter,
288 &node_id,
289 None,
290 0,
291 100,
292 )
293 };
294 VoteAccountInfo {
295 vote_account: solana_pubkey::new_rand(),
296 account,
297 authorized_voter,
298 }
299 })
300 .take(num_vote_accounts_per_node)
301 .collect(),
302 )
303 })
304 .collect()
305 }
306
307 fn new_epoch_vote_accounts(
308 vote_accounts_map: &HashMap<Pubkey, Vec<VoteAccountInfo>>,
309 node_id_to_stake_fn: impl Fn(&Pubkey) -> u64,
310 ) -> VoteAccountsHashMap {
311 vote_accounts_map
313 .iter()
314 .flat_map(|(node_id, vote_accounts)| {
315 vote_accounts.iter().map(|v| {
316 let vote_account = VoteAccount::try_from(v.account.clone()).unwrap();
317 (v.vote_account, (node_id_to_stake_fn(node_id), vote_account))
318 })
319 })
320 .collect()
321 }
322
323 #[test_case(true; "alpenglow")]
324 #[test_case(false; "towerbft")]
325 fn test_parse_epoch_vote_accounts(is_alpenglow: bool) {
326 let stake_per_account = 100;
327 let num_vote_accounts_per_node = 2;
328 let num_nodes = 10;
329
330 let vote_accounts_map =
331 new_vote_accounts(num_nodes, num_vote_accounts_per_node, is_alpenglow);
332
333 let expected_authorized_voters: HashMap<_, _> = vote_accounts_map
334 .iter()
335 .flat_map(|(_, vote_accounts)| {
336 vote_accounts
337 .iter()
338 .map(|v| (v.vote_account, v.authorized_voter))
339 })
340 .collect();
341
342 let expected_node_id_to_vote_accounts: HashMap<_, _> = vote_accounts_map
343 .iter()
344 .map(|(node_pubkey, vote_accounts)| {
345 let mut vote_accounts = vote_accounts
346 .iter()
347 .map(|v| v.vote_account)
348 .collect::<Vec<_>>();
349 vote_accounts.sort();
350 let node_vote_accounts = NodeVoteAccounts {
351 vote_accounts,
352 total_stake: stake_per_account * num_vote_accounts_per_node as u64,
353 };
354 (*node_pubkey, node_vote_accounts)
355 })
356 .collect();
357
358 let epoch_vote_accounts =
359 new_epoch_vote_accounts(&vote_accounts_map, |_| stake_per_account);
360
361 let (total_stake, mut node_id_to_vote_accounts, epoch_authorized_voters) =
362 VersionedEpochStakes::parse_epoch_vote_accounts(&epoch_vote_accounts, 0);
363
364 node_id_to_vote_accounts
366 .iter_mut()
367 .for_each(|(_, node_vote_accounts)| node_vote_accounts.vote_accounts.sort());
368
369 assert!(
370 node_id_to_vote_accounts.len() == expected_node_id_to_vote_accounts.len()
371 && node_id_to_vote_accounts
372 .iter()
373 .all(|(k, v)| expected_node_id_to_vote_accounts.get(k).unwrap() == v)
374 );
375 assert!(
376 epoch_authorized_voters.len() == expected_authorized_voters.len()
377 && epoch_authorized_voters
378 .iter()
379 .all(|(k, v)| expected_authorized_voters.get(k).unwrap() == v)
380 );
381 assert_eq!(
382 total_stake,
383 num_nodes as u64 * num_vote_accounts_per_node as u64 * 100
384 );
385 }
386
387 #[test_case(true; "alpenglow")]
388 #[test_case(false; "towerbft")]
389 fn test_node_id_to_stake(is_alpenglow: bool) {
390 let num_nodes = 10;
391 let num_vote_accounts_per_node = 2;
392
393 let vote_accounts_map =
394 new_vote_accounts(num_nodes, num_vote_accounts_per_node, is_alpenglow);
395 let node_id_to_stake_map = vote_accounts_map
396 .keys()
397 .enumerate()
398 .map(|(index, node_id)| (*node_id, ((index + 1) * 100) as u64))
399 .collect::<HashMap<_, _>>();
400 let epoch_vote_accounts = new_epoch_vote_accounts(&vote_accounts_map, |node_id| {
401 *node_id_to_stake_map.get(node_id).unwrap()
402 });
403 let epoch_stakes = VersionedEpochStakes::new_for_tests(epoch_vote_accounts, 0);
404
405 assert_eq!(epoch_stakes.total_stake(), 11000);
406 for (node_id, stake) in node_id_to_stake_map.iter() {
407 assert_eq!(
408 epoch_stakes.node_id_to_stake(node_id),
409 Some(*stake * num_vote_accounts_per_node as u64)
410 );
411 }
412 }
413
414 #[test_case(1; "single_vote_account")]
415 #[test_case(2; "multiple_vote_accounts")]
416 fn test_bls_pubkey_rank_map(num_vote_accounts_per_node: usize) {
417 agave_logger::setup();
418 let num_nodes = 10;
419 let num_vote_accounts = num_nodes * num_vote_accounts_per_node;
420
421 let vote_accounts_map = new_vote_accounts(num_nodes, num_vote_accounts_per_node, true);
422 let node_id_to_stake_map = vote_accounts_map
423 .keys()
424 .enumerate()
425 .map(|(index, node_id)| (*node_id, ((index + 1) * 100) as u64))
426 .collect::<HashMap<_, _>>();
427 let epoch_vote_accounts = new_epoch_vote_accounts(&vote_accounts_map, |node_id| {
428 *node_id_to_stake_map.get(node_id).unwrap()
429 });
430 let epoch_stakes = VersionedEpochStakes::new_for_tests(epoch_vote_accounts.clone(), 0);
431 let bls_pubkey_to_rank_map = epoch_stakes.bls_pubkey_to_rank_map();
432 assert_eq!(bls_pubkey_to_rank_map.len(), num_vote_accounts);
433 for (pubkey, (_, vote_account)) in epoch_vote_accounts {
434 let vote_state_view = vote_account.vote_state_view();
435 let bls_pubkey_compressed = bincode::deserialize::<BLSPubkeyCompressed>(
436 &vote_state_view.bls_pubkey_compressed().unwrap(),
437 )
438 .unwrap();
439 let bls_pubkey = BLSPubkey::try_from(bls_pubkey_compressed).unwrap();
440 let index = bls_pubkey_to_rank_map.get_rank(&bls_pubkey).unwrap();
441 assert!(index >= &0 && index < &(num_vote_accounts as u16));
442 assert_eq!(
443 bls_pubkey_to_rank_map.get_pubkey(*index as usize),
444 Some(&(pubkey, bls_pubkey))
445 );
446 }
447
448 let mut bank_epoch_stakes = HashMap::new();
450 bank_epoch_stakes.insert(0, epoch_stakes.clone());
451 let epoch_stakes = bank_epoch_stakes
452 .get(&0)
453 .expect("Epoch stakes should exist");
454 let bls_pubkey_to_rank_map2 = epoch_stakes.bls_pubkey_to_rank_map();
455 assert_eq!(bls_pubkey_to_rank_map2, bls_pubkey_to_rank_map);
456 }
457}