1use {
2 crate::stakes::SerdeStakesToStakeFormat,
3 serde::{Deserialize, Serialize},
4 solana_clock::Epoch,
5 solana_pubkey::Pubkey,
6 solana_vote::vote_account::VoteAccountsHashMap,
7 std::{collections::HashMap, sync::Arc},
8};
9
10pub type NodeIdToVoteAccounts = HashMap<Pubkey, NodeVoteAccounts>;
11pub type EpochAuthorizedVoters = HashMap<Pubkey, Pubkey>;
12
13#[cfg_attr(feature = "frozen-abi", derive(AbiExample))]
14#[derive(Clone, Serialize, Debug, Deserialize, Default, PartialEq, Eq)]
15pub struct NodeVoteAccounts {
16 pub vote_accounts: Vec<Pubkey>,
17 pub total_stake: u64,
18}
19
20#[derive(Clone, Debug, Serialize, Deserialize)]
21#[cfg_attr(feature = "frozen-abi", derive(AbiExample, AbiEnumVisitor))]
22#[cfg_attr(feature = "dev-context-only-utils", derive(PartialEq))]
23pub enum VersionedEpochStakes {
24 Current {
25 stakes: SerdeStakesToStakeFormat,
26 total_stake: u64,
27 node_id_to_vote_accounts: Arc<NodeIdToVoteAccounts>,
28 epoch_authorized_voters: Arc<EpochAuthorizedVoters>,
29 },
30}
31
32impl VersionedEpochStakes {
33 pub(crate) fn new(stakes: SerdeStakesToStakeFormat, leader_schedule_epoch: Epoch) -> Self {
34 let epoch_vote_accounts = stakes.vote_accounts();
35 let (total_stake, node_id_to_vote_accounts, epoch_authorized_voters) =
36 Self::parse_epoch_vote_accounts(epoch_vote_accounts.as_ref(), leader_schedule_epoch);
37 Self::Current {
38 stakes,
39 total_stake,
40 node_id_to_vote_accounts: Arc::new(node_id_to_vote_accounts),
41 epoch_authorized_voters: Arc::new(epoch_authorized_voters),
42 }
43 }
44
45 #[cfg(feature = "dev-context-only-utils")]
46 pub fn new_for_tests(
47 vote_accounts_hash_map: VoteAccountsHashMap,
48 leader_schedule_epoch: Epoch,
49 ) -> Self {
50 Self::new(
51 SerdeStakesToStakeFormat::Account(crate::stakes::Stakes::new_for_tests(
52 0,
53 solana_vote::vote_account::VoteAccounts::from(Arc::new(vote_accounts_hash_map)),
54 im::HashMap::default(),
55 )),
56 leader_schedule_epoch,
57 )
58 }
59
60 pub fn stakes(&self) -> &SerdeStakesToStakeFormat {
61 match self {
62 Self::Current { stakes, .. } => stakes,
63 }
64 }
65
66 pub fn total_stake(&self) -> u64 {
67 match self {
68 Self::Current { total_stake, .. } => *total_stake,
69 }
70 }
71
72 #[cfg(feature = "dev-context-only-utils")]
73 pub fn set_total_stake(&mut self, total_stake: u64) {
74 match self {
75 Self::Current {
76 total_stake: total_stake_field,
77 ..
78 } => {
79 *total_stake_field = total_stake;
80 }
81 }
82 }
83
84 pub fn node_id_to_vote_accounts(&self) -> &Arc<NodeIdToVoteAccounts> {
85 match self {
86 Self::Current {
87 node_id_to_vote_accounts,
88 ..
89 } => node_id_to_vote_accounts,
90 }
91 }
92
93 pub fn node_id_to_stake(&self, node_id: &Pubkey) -> Option<u64> {
94 self.node_id_to_vote_accounts()
95 .get(node_id)
96 .map(|x| x.total_stake)
97 }
98
99 pub fn epoch_authorized_voters(&self) -> &Arc<EpochAuthorizedVoters> {
100 match self {
101 Self::Current {
102 epoch_authorized_voters,
103 ..
104 } => epoch_authorized_voters,
105 }
106 }
107
108 pub fn vote_account_stake(&self, vote_account: &Pubkey) -> u64 {
109 self.stakes()
110 .vote_accounts()
111 .get_delegated_stake(vote_account)
112 }
113
114 fn parse_epoch_vote_accounts(
115 epoch_vote_accounts: &VoteAccountsHashMap,
116 leader_schedule_epoch: Epoch,
117 ) -> (u64, NodeIdToVoteAccounts, EpochAuthorizedVoters) {
118 let mut node_id_to_vote_accounts: NodeIdToVoteAccounts = HashMap::new();
119 let total_stake = epoch_vote_accounts
120 .iter()
121 .map(|(_, (stake, _))| stake)
122 .sum();
123 let epoch_authorized_voters = epoch_vote_accounts
124 .iter()
125 .filter_map(|(key, (stake, account))| {
126 let vote_state = account.vote_state_view();
127
128 if *stake > 0 {
129 if let Some(authorized_voter) =
130 vote_state.get_authorized_voter(leader_schedule_epoch)
131 {
132 let node_vote_accounts = node_id_to_vote_accounts
133 .entry(*vote_state.node_pubkey())
134 .or_default();
135
136 node_vote_accounts.total_stake += stake;
137 node_vote_accounts.vote_accounts.push(*key);
138
139 Some((*key, *authorized_voter))
140 } else {
141 None
142 }
143 } else {
144 None
145 }
146 })
147 .collect();
148 (
149 total_stake,
150 node_id_to_vote_accounts,
151 epoch_authorized_voters,
152 )
153 }
154}
155
156#[cfg(test)]
157pub(crate) mod tests {
158 use {
159 super::*, solana_account::AccountSharedData, solana_vote::vote_account::VoteAccount,
160 solana_vote_program::vote_state::create_account_with_authorized, std::iter,
161 };
162
163 struct VoteAccountInfo {
164 vote_account: Pubkey,
165 account: AccountSharedData,
166 authorized_voter: Pubkey,
167 }
168
169 fn new_vote_accounts(
170 num_nodes: usize,
171 num_vote_accounts_per_node: usize,
172 ) -> HashMap<Pubkey, Vec<VoteAccountInfo>> {
173 (0..num_nodes)
175 .map(|_| {
176 let node_id = solana_pubkey::new_rand();
177 (
178 node_id,
179 iter::repeat_with(|| {
180 let authorized_voter = solana_pubkey::new_rand();
181 VoteAccountInfo {
182 vote_account: solana_pubkey::new_rand(),
183 account: create_account_with_authorized(
184 &node_id,
185 &authorized_voter,
186 &node_id,
187 0,
188 100,
189 ),
190 authorized_voter,
191 }
192 })
193 .take(num_vote_accounts_per_node)
194 .collect(),
195 )
196 })
197 .collect()
198 }
199
200 fn new_epoch_vote_accounts(
201 vote_accounts_map: &HashMap<Pubkey, Vec<VoteAccountInfo>>,
202 node_id_to_stake_fn: impl Fn(&Pubkey) -> u64,
203 ) -> VoteAccountsHashMap {
204 vote_accounts_map
206 .iter()
207 .flat_map(|(node_id, vote_accounts)| {
208 vote_accounts.iter().map(|v| {
209 let vote_account = VoteAccount::try_from(v.account.clone()).unwrap();
210 (v.vote_account, (node_id_to_stake_fn(node_id), vote_account))
211 })
212 })
213 .collect()
214 }
215
216 #[test]
217 fn test_parse_epoch_vote_accounts() {
218 let stake_per_account = 100;
219 let num_vote_accounts_per_node = 2;
220 let num_nodes = 10;
221
222 let vote_accounts_map = new_vote_accounts(num_nodes, num_vote_accounts_per_node);
223
224 let expected_authorized_voters: HashMap<_, _> = vote_accounts_map
225 .iter()
226 .flat_map(|(_, vote_accounts)| {
227 vote_accounts
228 .iter()
229 .map(|v| (v.vote_account, v.authorized_voter))
230 })
231 .collect();
232
233 let expected_node_id_to_vote_accounts: HashMap<_, _> = vote_accounts_map
234 .iter()
235 .map(|(node_pubkey, vote_accounts)| {
236 let mut vote_accounts = vote_accounts
237 .iter()
238 .map(|v| (v.vote_account))
239 .collect::<Vec<_>>();
240 vote_accounts.sort();
241 let node_vote_accounts = NodeVoteAccounts {
242 vote_accounts,
243 total_stake: stake_per_account * num_vote_accounts_per_node as u64,
244 };
245 (*node_pubkey, node_vote_accounts)
246 })
247 .collect();
248
249 let epoch_vote_accounts =
250 new_epoch_vote_accounts(&vote_accounts_map, |_| stake_per_account);
251
252 let (total_stake, mut node_id_to_vote_accounts, epoch_authorized_voters) =
253 VersionedEpochStakes::parse_epoch_vote_accounts(&epoch_vote_accounts, 0);
254
255 node_id_to_vote_accounts
257 .iter_mut()
258 .for_each(|(_, node_vote_accounts)| node_vote_accounts.vote_accounts.sort());
259
260 assert!(
261 node_id_to_vote_accounts.len() == expected_node_id_to_vote_accounts.len()
262 && node_id_to_vote_accounts
263 .iter()
264 .all(|(k, v)| expected_node_id_to_vote_accounts.get(k).unwrap() == v)
265 );
266 assert!(
267 epoch_authorized_voters.len() == expected_authorized_voters.len()
268 && epoch_authorized_voters
269 .iter()
270 .all(|(k, v)| expected_authorized_voters.get(k).unwrap() == v)
271 );
272 assert_eq!(
273 total_stake,
274 num_nodes as u64 * num_vote_accounts_per_node as u64 * 100
275 );
276 }
277
278 #[test]
279 fn test_node_id_to_stake() {
280 let num_nodes = 10;
281 let num_vote_accounts_per_node = 2;
282
283 let vote_accounts_map = new_vote_accounts(num_nodes, num_vote_accounts_per_node);
284 let node_id_to_stake_map = vote_accounts_map
285 .keys()
286 .enumerate()
287 .map(|(index, node_id)| (*node_id, ((index + 1) * 100) as u64))
288 .collect::<HashMap<_, _>>();
289 let epoch_vote_accounts = new_epoch_vote_accounts(&vote_accounts_map, |node_id| {
290 *node_id_to_stake_map.get(node_id).unwrap()
291 });
292 let epoch_stakes = VersionedEpochStakes::new_for_tests(epoch_vote_accounts, 0);
293
294 assert_eq!(epoch_stakes.total_stake(), 11000);
295 for (node_id, stake) in node_id_to_stake_map.iter() {
296 assert_eq!(
297 epoch_stakes.node_id_to_stake(node_id),
298 Some(*stake * num_vote_accounts_per_node as u64)
299 );
300 }
301 }
302}