1use {
2 crate::{stakes::StakesEnum, vote_account::VoteAccountsHashMap},
3 serde::{Deserialize, Serialize},
4 solana_sdk::{clock::Epoch, pubkey::Pubkey},
5 std::{collections::HashMap, sync::Arc},
6};
7
8pub type NodeIdToVoteAccounts = HashMap<Pubkey, NodeVoteAccounts>;
9pub type EpochAuthorizedVoters = HashMap<Pubkey, Pubkey>;
10
11#[derive(Clone, Serialize, Debug, Deserialize, Default, PartialEq, Eq, AbiExample)]
12pub struct NodeVoteAccounts {
13 pub vote_accounts: Vec<Pubkey>,
14 pub total_stake: u64,
15}
16
17#[derive(Clone, Debug, Serialize, Deserialize, AbiExample, PartialEq)]
18pub struct EpochStakes {
19 #[serde(with = "crate::stakes::serde_stakes_enum_compat")]
20 stakes: Arc<StakesEnum>,
21 total_stake: u64,
22 node_id_to_vote_accounts: Arc<NodeIdToVoteAccounts>,
23 epoch_authorized_voters: Arc<EpochAuthorizedVoters>,
24}
25
26impl EpochStakes {
27 pub(crate) fn new(stakes: Arc<StakesEnum>, leader_schedule_epoch: Epoch) -> Self {
28 let epoch_vote_accounts = stakes.vote_accounts();
29 let (total_stake, node_id_to_vote_accounts, epoch_authorized_voters) =
30 Self::parse_epoch_vote_accounts(epoch_vote_accounts.as_ref(), leader_schedule_epoch);
31 Self {
32 stakes,
33 total_stake,
34 node_id_to_vote_accounts: Arc::new(node_id_to_vote_accounts),
35 epoch_authorized_voters: Arc::new(epoch_authorized_voters),
36 }
37 }
38
39 pub fn stakes(&self) -> &StakesEnum {
40 &self.stakes
41 }
42
43 pub fn total_stake(&self) -> u64 {
44 self.total_stake
45 }
46
47 pub fn node_id_to_vote_accounts(&self) -> &Arc<NodeIdToVoteAccounts> {
48 &self.node_id_to_vote_accounts
49 }
50
51 pub fn epoch_authorized_voters(&self) -> &Arc<EpochAuthorizedVoters> {
52 &self.epoch_authorized_voters
53 }
54
55 pub fn vote_account_stake(&self, vote_account: &Pubkey) -> u64 {
56 self.stakes
57 .vote_accounts()
58 .get_delegated_stake(vote_account)
59 }
60
61 fn parse_epoch_vote_accounts(
62 epoch_vote_accounts: &VoteAccountsHashMap,
63 leader_schedule_epoch: Epoch,
64 ) -> (u64, NodeIdToVoteAccounts, EpochAuthorizedVoters) {
65 let mut node_id_to_vote_accounts: NodeIdToVoteAccounts = HashMap::new();
66 let total_stake = epoch_vote_accounts
67 .iter()
68 .map(|(_, (stake, _))| stake)
69 .sum();
70 let epoch_authorized_voters = epoch_vote_accounts
71 .iter()
72 .filter_map(|(key, (stake, account))| {
73 let vote_state = account.vote_state();
74 let vote_state = match vote_state.as_ref() {
75 Err(_) => {
76 datapoint_warn!(
77 "parse_epoch_vote_accounts",
78 (
79 "warn",
80 format!("Unable to get vote_state from account {}", key),
81 String
82 ),
83 );
84 return None;
85 }
86 Ok(vote_state) => vote_state,
87 };
88
89 if *stake > 0 {
90 if let Some(authorized_voter) = vote_state
91 .authorized_voters()
92 .get_authorized_voter(leader_schedule_epoch)
93 {
94 let node_vote_accounts = node_id_to_vote_accounts
95 .entry(vote_state.node_pubkey)
96 .or_default();
97
98 node_vote_accounts.total_stake += stake;
99 node_vote_accounts.vote_accounts.push(*key);
100
101 Some((*key, authorized_voter))
102 } else {
103 None
104 }
105 } else {
106 None
107 }
108 })
109 .collect();
110 (
111 total_stake,
112 node_id_to_vote_accounts,
113 epoch_authorized_voters,
114 )
115 }
116}
117
118#[cfg(test)]
119pub(crate) mod tests {
120 use {
121 super::*, crate::vote_account::VoteAccount, solana_sdk::account::AccountSharedData,
122 solana_vote_program::vote_state::create_account_with_authorized, std::iter,
123 };
124
125 struct VoteAccountInfo {
126 vote_account: Pubkey,
127 account: AccountSharedData,
128 authorized_voter: Pubkey,
129 }
130
131 #[test]
132 fn test_parse_epoch_vote_accounts() {
133 let stake_per_account = 100;
134 let num_vote_accounts_per_node = 2;
135 let vote_accounts_map: HashMap<Pubkey, Vec<VoteAccountInfo>> = (0..10)
137 .map(|_| {
138 let node_id = solana_sdk::pubkey::new_rand();
139 (
140 node_id,
141 iter::repeat_with(|| {
142 let authorized_voter = solana_sdk::pubkey::new_rand();
143 VoteAccountInfo {
144 vote_account: solana_sdk::pubkey::new_rand(),
145 account: create_account_with_authorized(
146 &node_id,
147 &authorized_voter,
148 &node_id,
149 0,
150 100,
151 ),
152 authorized_voter,
153 }
154 })
155 .take(num_vote_accounts_per_node)
156 .collect(),
157 )
158 })
159 .collect();
160
161 let expected_authorized_voters: HashMap<_, _> = vote_accounts_map
162 .iter()
163 .flat_map(|(_, vote_accounts)| {
164 vote_accounts
165 .iter()
166 .map(|v| (v.vote_account, v.authorized_voter))
167 })
168 .collect();
169
170 let expected_node_id_to_vote_accounts: HashMap<_, _> = vote_accounts_map
171 .iter()
172 .map(|(node_pubkey, vote_accounts)| {
173 let mut vote_accounts = vote_accounts
174 .iter()
175 .map(|v| (v.vote_account))
176 .collect::<Vec<_>>();
177 vote_accounts.sort();
178 let node_vote_accounts = NodeVoteAccounts {
179 vote_accounts,
180 total_stake: stake_per_account * num_vote_accounts_per_node as u64,
181 };
182 (*node_pubkey, node_vote_accounts)
183 })
184 .collect();
185
186 let epoch_vote_accounts: HashMap<_, _> = vote_accounts_map
188 .iter()
189 .flat_map(|(_, vote_accounts)| {
190 vote_accounts.iter().map(|v| {
191 let vote_account = VoteAccount::try_from(v.account.clone()).unwrap();
192 (v.vote_account, (stake_per_account, vote_account))
193 })
194 })
195 .collect();
196
197 let (total_stake, mut node_id_to_vote_accounts, epoch_authorized_voters) =
198 EpochStakes::parse_epoch_vote_accounts(&epoch_vote_accounts, 0);
199
200 node_id_to_vote_accounts
202 .iter_mut()
203 .for_each(|(_, node_vote_accounts)| node_vote_accounts.vote_accounts.sort());
204
205 assert!(
206 node_id_to_vote_accounts.len() == expected_node_id_to_vote_accounts.len()
207 && node_id_to_vote_accounts
208 .iter()
209 .all(|(k, v)| expected_node_id_to_vote_accounts.get(k).unwrap() == v)
210 );
211 assert!(
212 epoch_authorized_voters.len() == expected_authorized_voters.len()
213 && epoch_authorized_voters
214 .iter()
215 .all(|(k, v)| expected_authorized_voters.get(k).unwrap() == v)
216 );
217 assert_eq!(
218 total_stake,
219 vote_accounts_map.len() as u64 * num_vote_accounts_per_node as u64 * 100
220 );
221 }
222}