diff --git a/programs/vote/src/vote_state/mod.rs b/programs/vote/src/vote_state/mod.rs index 6dc54f338295b9..6d4436a8621a85 100644 --- a/programs/vote/src/vote_state/mod.rs +++ b/programs/vote/src/vote_state/mod.rs @@ -16,7 +16,7 @@ use { solana_rent::Rent, solana_slot_hashes::SlotHash, solana_transaction_context::{BorrowedInstructionAccount, IndexOfAccount, InstructionContext}, - solana_vote_interface::{error::VoteError, program::id}, + solana_vote_interface::{authorized_voters::AuthorizedVoters, error::VoteError, program::id}, std::{ cmp::Ordering, collections::{HashSet, VecDeque}, @@ -1040,6 +1040,51 @@ pub fn create_account_with_authorized( vote_account } +// TODO(wen): when we have VoteStateV4::new(), switch all users there. +pub fn new_v4_vote_state( + node_pubkey: &Pubkey, + authorized_voter: &Pubkey, + authorized_withdrawer: &Pubkey, + bls_pubkey_compressed: Option<[u8; BLS_PUBLIC_KEY_COMPRESSED_SIZE]>, + inflation_rewards_commission_bps: u16, +) -> VoteStateV4 { + VoteStateV4 { + node_pubkey: *node_pubkey, + authorized_voters: AuthorizedVoters::new(0, *authorized_voter), + authorized_withdrawer: *authorized_withdrawer, + bls_pubkey_compressed, + inflation_rewards_commission_bps, + ..VoteStateV4::default() + } +} + +pub fn create_v4_account_with_authorized( + node_pubkey: &Pubkey, + authorized_voter: &Pubkey, + authorized_withdrawer: &Pubkey, + bls_pubkey_compressed: Option<[u8; BLS_PUBLIC_KEY_COMPRESSED_SIZE]>, + inflation_rewards_commission_bps: u16, + lamports: u64, +) -> AccountSharedData { + let mut vote_account = AccountSharedData::new(lamports, VoteStateV4::size_of(), &id()); + + let vote_state = new_v4_vote_state( + node_pubkey, + authorized_voter, + authorized_withdrawer, + bls_pubkey_compressed, + inflation_rewards_commission_bps, + ); + + VoteStateV4::serialize( + &VoteStateVersions::V4(Box::new(vote_state)), + vote_account.data_as_mut_slice(), + ) + .unwrap(); + + vote_account +} + // create_account() should be removed, use create_account_with_authorized() instead pub fn create_account( vote_pubkey: &Pubkey, @@ -3531,4 +3576,40 @@ mod tests { expected_allowed ); } + + #[test] + fn test_create_v4_account_with_authorized() { + let node_pubkey = Pubkey::new_unique(); + let authorized_voter = Pubkey::new_unique(); + let authorized_withdrawer = Pubkey::new_unique(); + let bls_pubkey_compressed = [42; 48]; + let inflation_rewards_commission_bps = 10000; + let lamports = 100; + let vote_account = create_v4_account_with_authorized( + &node_pubkey, + &authorized_voter, + &authorized_withdrawer, + Some(bls_pubkey_compressed), + inflation_rewards_commission_bps, + lamports, + ); + assert_eq!(vote_account.lamports(), lamports); + assert_eq!(vote_account.owner(), &id()); + assert_eq!(vote_account.data().len(), VoteStateV4::size_of()); + let vote_state_v4 = VoteStateV4::deserialize(vote_account.data(), &node_pubkey).unwrap(); + assert_eq!(vote_state_v4.node_pubkey, node_pubkey); + assert_eq!( + vote_state_v4.authorized_voters, + AuthorizedVoters::new(0, authorized_voter) + ); + assert_eq!(vote_state_v4.authorized_withdrawer, authorized_withdrawer); + assert_eq!( + vote_state_v4.bls_pubkey_compressed, + Some(bls_pubkey_compressed) + ); + assert_eq!( + vote_state_v4.inflation_rewards_commission_bps, + inflation_rewards_commission_bps + ); + } } diff --git a/runtime/src/epoch_stakes.rs b/runtime/src/epoch_stakes.rs index 618d2ef0002782..25039bce66ee66 100644 --- a/runtime/src/epoch_stakes.rs +++ b/runtime/src/epoch_stakes.rs @@ -174,6 +174,19 @@ impl VersionedEpochStakes { } } + pub fn bls_pubkey_to_rank_map(&self) -> &Arc { + match self { + Self::Current { + bls_pubkey_to_rank_map, + .. + } => bls_pubkey_to_rank_map.get_or_init(|| { + Arc::new(BLSPubkeyToRankMap::new( + self.stakes().vote_accounts().as_ref(), + )) + }), + } + } + pub fn vote_account_stake(&self, vote_account: &Pubkey) -> u64 { self.stakes() .vote_accounts() @@ -225,8 +238,15 @@ impl VersionedEpochStakes { #[cfg(test)] pub(crate) mod tests { use { - super::*, solana_account::AccountSharedData, solana_vote::vote_account::VoteAccount, - solana_vote_program::vote_state::create_account_with_authorized, std::iter, + super::*, + solana_account::AccountSharedData, + solana_bls_signatures::keypair::Keypair as BLSKeypair, + solana_vote::vote_account::VoteAccount, + solana_vote_program::vote_state::{ + create_account_with_authorized, create_v4_account_with_authorized, + }, + std::iter, + test_case::test_case, }; struct VoteAccountInfo { @@ -238,6 +258,7 @@ pub(crate) mod tests { fn new_vote_accounts( num_nodes: usize, num_vote_accounts_per_node: usize, + is_alpenglow: bool, ) -> HashMap> { // Create some vote accounts for each pubkey (0..num_nodes) @@ -247,15 +268,35 @@ pub(crate) mod tests { node_id, iter::repeat_with(|| { let authorized_voter = solana_pubkey::new_rand(); - VoteAccountInfo { - vote_account: solana_pubkey::new_rand(), - account: create_account_with_authorized( + let bls_pubkey_compressed: BLSPubkeyCompressed = + BLSKeypair::new().public.try_into().unwrap(); + let bls_pubkey_compressed_serialized = + bincode::serialize(&bls_pubkey_compressed) + .unwrap() + .try_into() + .unwrap(); + + let account = if is_alpenglow { + create_v4_account_with_authorized( &node_id, &authorized_voter, &node_id, + Some(bls_pubkey_compressed_serialized), 0, 100, - ), + ) + } else { + create_account_with_authorized( + &node_id, + &authorized_voter, + &node_id, + 0, + 100, + ) + }; + VoteAccountInfo { + vote_account: solana_pubkey::new_rand(), + account, authorized_voter, } }) @@ -282,13 +323,15 @@ pub(crate) mod tests { .collect() } - #[test] - fn test_parse_epoch_vote_accounts() { + #[test_case(true; "alpenglow")] + #[test_case(false; "towerbft")] + fn test_parse_epoch_vote_accounts(is_alpenglow: bool) { let stake_per_account = 100; let num_vote_accounts_per_node = 2; let num_nodes = 10; - let vote_accounts_map = new_vote_accounts(num_nodes, num_vote_accounts_per_node); + let vote_accounts_map = + new_vote_accounts(num_nodes, num_vote_accounts_per_node, is_alpenglow); let expected_authorized_voters: HashMap<_, _> = vote_accounts_map .iter() @@ -344,12 +387,14 @@ pub(crate) mod tests { ); } - #[test] - fn test_node_id_to_stake() { + #[test_case(true; "alpenglow")] + #[test_case(false; "towerbft")] + fn test_node_id_to_stake(is_alpenglow: bool) { let num_nodes = 10; let num_vote_accounts_per_node = 2; - let vote_accounts_map = new_vote_accounts(num_nodes, num_vote_accounts_per_node); + let vote_accounts_map = + new_vote_accounts(num_nodes, num_vote_accounts_per_node, is_alpenglow); let node_id_to_stake_map = vote_accounts_map .keys() .enumerate() @@ -368,4 +413,48 @@ pub(crate) mod tests { ); } } + + #[test_case(1; "single_vote_account")] + #[test_case(2; "multiple_vote_accounts")] + fn test_bls_pubkey_rank_map(num_vote_accounts_per_node: usize) { + solana_logger::setup(); + let num_nodes = 10; + let num_vote_accounts = num_nodes * num_vote_accounts_per_node; + + let vote_accounts_map = new_vote_accounts(num_nodes, num_vote_accounts_per_node, true); + let node_id_to_stake_map = vote_accounts_map + .keys() + .enumerate() + .map(|(index, node_id)| (*node_id, ((index + 1) * 100) as u64)) + .collect::>(); + let epoch_vote_accounts = new_epoch_vote_accounts(&vote_accounts_map, |node_id| { + *node_id_to_stake_map.get(node_id).unwrap() + }); + let epoch_stakes = VersionedEpochStakes::new_for_tests(epoch_vote_accounts.clone(), 0); + let bls_pubkey_to_rank_map = epoch_stakes.bls_pubkey_to_rank_map(); + assert_eq!(bls_pubkey_to_rank_map.len(), num_vote_accounts); + for (pubkey, (_, vote_account)) in epoch_vote_accounts { + let vote_state_view = vote_account.vote_state_view(); + let bls_pubkey_compressed = bincode::deserialize::( + &vote_state_view.bls_pubkey_compressed().unwrap(), + ) + .unwrap(); + let bls_pubkey = BLSPubkey::try_from(bls_pubkey_compressed).unwrap(); + let index = bls_pubkey_to_rank_map.get_rank(&bls_pubkey).unwrap(); + assert!(index >= &0 && index < &(num_vote_accounts as u16)); + assert_eq!( + bls_pubkey_to_rank_map.get_pubkey(*index as usize), + Some(&(pubkey, bls_pubkey)) + ); + } + + // Convert it to versioned and back, we should get the same rank map + let mut bank_epoch_stakes = HashMap::new(); + bank_epoch_stakes.insert(0, epoch_stakes.clone()); + let epoch_stakes = bank_epoch_stakes + .get(&0) + .expect("Epoch stakes should exist"); + let bls_pubkey_to_rank_map2 = epoch_stakes.bls_pubkey_to_rank_map(); + assert_eq!(bls_pubkey_to_rank_map2, bls_pubkey_to_rank_map); + } }