From bd48e406293b08519509090940bd835c7fc91fae Mon Sep 17 00:00:00 2001 From: Justin Starry Date: Fri, 14 Feb 2025 15:15:53 +0000 Subject: [PATCH] use modules --- ledger/src/blockstore.rs | 4 +- ledger/src/leader_schedule.rs | 173 ++++--------------- ledger/src/leader_schedule/identity_keyed.rs | 49 ++++++ ledger/src/leader_schedule/vote_keyed.rs | 68 ++++++++ ledger/src/leader_schedule_utils.rs | 6 +- local-cluster/src/integration_tests.rs | 2 +- poh/src/poh_recorder.rs | 4 +- 7 files changed, 164 insertions(+), 142 deletions(-) create mode 100644 ledger/src/leader_schedule/identity_keyed.rs create mode 100644 ledger/src/leader_schedule/vote_keyed.rs diff --git a/ledger/src/blockstore.rs b/ledger/src/blockstore.rs index e224892376ff15..9a6f7ed99c7756 100644 --- a/ledger/src/blockstore.rs +++ b/ledger/src/blockstore.rs @@ -10486,8 +10486,8 @@ pub mod tests { let bank = Arc::new(Bank::new_for_tests(&genesis_config)); let mut leader_schedule_cache = LeaderScheduleCache::new_from_bank(&bank); let fixed_schedule = FixedSchedule { - leader_schedule: Arc::new(LeaderSchedule::new_from_schedule(vec![ - leader_keypair.pubkey() + leader_schedule: Arc::new(LeaderSchedule::new_identity_keyed_with_schedule(vec![ + leader_keypair.pubkey(), ])), }; leader_schedule_cache.set_fixed_leader_schedule(Some(fixed_schedule)); diff --git a/ledger/src/leader_schedule.rs b/ledger/src/leader_schedule.rs index 91ec660f1b095c..b217390c8657fb 100644 --- a/ledger/src/leader_schedule.rs +++ b/ledger/src/leader_schedule.rs @@ -1,5 +1,4 @@ use { - itertools::Itertools, rand::distributions::{Distribution, WeightedIndex}, rand_chacha::{rand_core::SeedableRng, ChaChaRng}, solana_pubkey::Pubkey, @@ -8,6 +7,9 @@ use { std::{collections::HashMap, convert::identity, ops::Index, sync::Arc}, }; +mod identity_keyed; +mod vote_keyed; + // Used for testing #[derive(Clone, Debug)] pub struct FixedSchedule { @@ -21,120 +23,48 @@ pub struct LeaderSchedule(LeaderScheduleVariants); #[cfg(feature = "dev-context-only-utils")] impl Default for LeaderSchedule { fn default() -> Self { - Self( - LeaderScheduleVariants::ValidatorIdentityKeyedLeaderSchedule( - ValidatorIdentityKeyedLeaderSchedule { - slot_leaders: vec![], - index: HashMap::new(), - }, - ), - ) + Self(LeaderScheduleVariants::IdentityKeyed( + identity_keyed::LeaderSchedule::default(), + )) } } -#[derive(Debug, PartialEq, Eq, Clone)] -struct VoteAccountKeyedLeaderSchedule { - slot_leader_vote_account_addresses: Vec, - // cached leader schedule keyed by validator identities created by mapping - // vote account addresses to the validator identity designated at the time - // of leader schedule generation. This is used to avoid the need to look up - // the validator identity address for each slot. - validator_identity_keyed_leader_schedule: ValidatorIdentityKeyedLeaderSchedule, -} - -#[derive(Debug, PartialEq, Eq, Clone)] -struct ValidatorIdentityKeyedLeaderSchedule { - slot_leaders: Vec, - // Inverted index from pubkeys to indices where they are the leader. - index: HashMap>>, -} - #[derive(Debug, PartialEq, Eq, Clone)] enum LeaderScheduleVariants { // Latest leader schedule algorithm which designates a specific vote account // to each slot so that the runtime can load vote state (e.g. commission and // fee collector accounts) for a given slot - VoteAccountKeyedLeaderSchedule(VoteAccountKeyedLeaderSchedule), + VoteKeyed(vote_keyed::LeaderSchedule), // Old leader schedule algorithm which designates a specific validator // identity to each slot. Since multiple vote accounts can be associated // with a single validator identity, it's not possible to use this to load // vote state for a given slot. - ValidatorIdentityKeyedLeaderSchedule(ValidatorIdentityKeyedLeaderSchedule), + IdentityKeyed(identity_keyed::LeaderSchedule), } impl LeaderSchedule { // Note: passing in zero vote accounts will cause a panic. - pub fn new_keyed_by_vote_account( + pub fn new_vote_keyed( vote_accounts_map: &VoteAccountsHashMap, epoch: Epoch, len: u64, repeat: u64, ) -> Self { - let keyed_stakes: Vec<_> = vote_accounts_map - .iter() - .map(|(vote_pubkey, (stake, _account))| (vote_pubkey, *stake)) - .collect(); - let slot_leader_vote_account_addresses = - Self::stake_weighted_slot_leaders(keyed_stakes, epoch, len, repeat); - - let validator_identity_keyed_leader_schedule = { - struct SlotLeaderInfo<'a> { - vote_account_address: &'a Pubkey, - validator_identity_address: &'a Pubkey, - } - - let default_pubkey = Pubkey::default(); - let mut current_slot_leader_info = SlotLeaderInfo { - vote_account_address: &default_pubkey, - validator_identity_address: &default_pubkey, - }; - - let slot_leaders: Vec = slot_leader_vote_account_addresses - .iter() - .map(|vote_account_address| { - if vote_account_address != current_slot_leader_info.vote_account_address { - let validator_identity_address = vote_accounts_map - .get(vote_account_address) - .unwrap() - .1 - .node_pubkey(); - current_slot_leader_info = SlotLeaderInfo { - vote_account_address, - validator_identity_address, - }; - } - *current_slot_leader_info.validator_identity_address - }) - .collect(); - - let index = Self::index_from_slot_leaders(&slot_leaders); - ValidatorIdentityKeyedLeaderSchedule { - slot_leaders, - index, - } - }; - - Self(LeaderScheduleVariants::VoteAccountKeyedLeaderSchedule( - VoteAccountKeyedLeaderSchedule { - slot_leader_vote_account_addresses, - validator_identity_keyed_leader_schedule, - }, + Self(LeaderScheduleVariants::VoteKeyed( + vote_keyed::LeaderSchedule::new(vote_accounts_map, epoch, len, repeat), )) } // Note: passing in zero stakers will cause a panic. - pub fn new_keyed_by_validator_identity( + pub fn new_identity_keyed( epoch_staked_nodes: &HashMap, epoch: Epoch, len: u64, repeat: u64, ) -> Self { - let keyed_stakes: Vec<_> = epoch_staked_nodes - .iter() - .map(|(pubkey, stake)| (pubkey, *stake)) - .collect(); - let slot_leaders = Self::stake_weighted_slot_leaders(keyed_stakes, epoch, len, repeat); - Self::new_from_schedule(slot_leaders) + Self(LeaderScheduleVariants::IdentityKeyed( + identity_keyed::LeaderSchedule::new(epoch_staked_nodes, epoch, len, repeat), + )) } // Note: passing in zero stakers will cause a panic. @@ -161,33 +91,14 @@ impl LeaderSchedule { .collect() } - pub fn new_from_schedule(slot_leaders: Vec) -> Self { - Self( - LeaderScheduleVariants::ValidatorIdentityKeyedLeaderSchedule( - ValidatorIdentityKeyedLeaderSchedule { - index: Self::index_from_slot_leaders(&slot_leaders), - slot_leaders, - }, - ), - ) - } - - fn index_from_slot_leaders(slot_leaders: &[Pubkey]) -> HashMap>> { - slot_leaders - .iter() - .enumerate() - .map(|(i, pk)| (*pk, i)) - .into_group_map() - .into_iter() - .map(|(k, v)| (k, Arc::new(v))) - .collect() + pub fn new_identity_keyed_with_schedule(slot_leaders: Vec) -> Self { + Self(LeaderScheduleVariants::IdentityKeyed( + identity_keyed::LeaderSchedule::new_from_schedule(slot_leaders), + )) } - pub fn is_keyed_by_vote_account(&self) -> bool { - matches!( - self.0, - LeaderScheduleVariants::VoteAccountKeyedLeaderSchedule(_) - ) + pub fn is_vote_keyed(&self) -> bool { + matches!(self.0, LeaderScheduleVariants::VoteKeyed(_)) } /// Get the vote account address for the given epoch slot index. This is @@ -198,23 +109,19 @@ impl LeaderSchedule { epoch_slot_index: usize, ) -> Option<&Pubkey> { match &self.0 { - LeaderScheduleVariants::VoteAccountKeyedLeaderSchedule(schedule) => schedule + LeaderScheduleVariants::VoteKeyed(schedule) => schedule .slot_leader_vote_account_addresses .get(epoch_slot_index), - LeaderScheduleVariants::ValidatorIdentityKeyedLeaderSchedule(_) => None, + LeaderScheduleVariants::IdentityKeyed(_) => None, } } pub fn get_slot_leaders(&self) -> &[Pubkey] { match self.0 { - LeaderScheduleVariants::VoteAccountKeyedLeaderSchedule(ref schedule) => { - &schedule - .validator_identity_keyed_leader_schedule - .slot_leaders - } - LeaderScheduleVariants::ValidatorIdentityKeyedLeaderSchedule(ref schedule) => { - &schedule.slot_leaders + LeaderScheduleVariants::VoteKeyed(ref schedule) => { + &schedule.identity_keyed_leader_schedule.slot_leaders } + LeaderScheduleVariants::IdentityKeyed(ref schedule) => &schedule.slot_leaders, } } @@ -224,12 +131,10 @@ impl LeaderSchedule { fn index(&self) -> &HashMap>> { match &self.0 { - LeaderScheduleVariants::VoteAccountKeyedLeaderSchedule(schedule) => { - &schedule.validator_identity_keyed_leader_schedule.index - } - LeaderScheduleVariants::ValidatorIdentityKeyedLeaderSchedule(schedule) => { - &schedule.index + LeaderScheduleVariants::VoteKeyed(schedule) => { + &schedule.identity_keyed_leader_schedule.index } + LeaderScheduleVariants::IdentityKeyed(schedule) => &schedule.index, } } @@ -287,13 +192,14 @@ fn sort_stakes(stakes: &mut Vec<(&Pubkey, u64)>) { #[cfg(test)] mod tests { - use {super::*, rand::Rng, std::iter::repeat_with}; + use {super::*, itertools::Itertools, rand::Rng, std::iter::repeat_with}; #[test] fn test_leader_schedule_index() { let pubkey0 = solana_pubkey::new_rand(); let pubkey1 = solana_pubkey::new_rand(); - let leader_schedule = LeaderSchedule::new_from_schedule(vec![pubkey0, pubkey1]); + let leader_schedule = + LeaderSchedule::new_identity_keyed_with_schedule(vec![pubkey0, pubkey1]); assert_eq!(leader_schedule[0], pubkey0); assert_eq!(leader_schedule[1], pubkey1); assert_eq!(leader_schedule[2], pubkey0); @@ -308,10 +214,8 @@ mod tests { let epoch: Epoch = rand::random(); let len = num_keys * 10; - let leader_schedule = - LeaderSchedule::new_keyed_by_validator_identity(&stakes, epoch, len, 1); - let leader_schedule2 = - LeaderSchedule::new_keyed_by_validator_identity(&stakes, epoch, len, 1); + let leader_schedule = LeaderSchedule::new_identity_keyed(&stakes, epoch, len, 1); + let leader_schedule2 = LeaderSchedule::new_identity_keyed(&stakes, epoch, len, 1); assert_eq!(leader_schedule.num_slots() as u64, len); // Check that the same schedule is reproducibly generated assert_eq!(leader_schedule, leader_schedule2); @@ -327,8 +231,7 @@ mod tests { let epoch = rand::random::(); let len = num_keys * 10; let repeat = 8; - let leader_schedule = - LeaderSchedule::new_keyed_by_validator_identity(&stakes, epoch, len, repeat); + let leader_schedule = LeaderSchedule::new_identity_keyed(&stakes, epoch, len, repeat); assert_eq!(leader_schedule.num_slots() as u64, len); let mut leader_node = Pubkey::default(); for (i, node) in leader_schedule.get_slot_leaders().iter().enumerate() { @@ -349,12 +252,12 @@ mod tests { let epoch = 0; let len = 8; // What the schedule looks like without any repeats - let leaders1 = LeaderSchedule::new_keyed_by_validator_identity(&stakes, epoch, len, 1) + let leaders1 = LeaderSchedule::new_identity_keyed(&stakes, epoch, len, 1) .get_slot_leaders() .to_vec(); // What the schedule looks like with repeats - let leaders2 = LeaderSchedule::new_keyed_by_validator_identity(&stakes, epoch, len, 2) + let leaders2 = LeaderSchedule::new_identity_keyed(&stakes, epoch, len, 2) .get_slot_leaders() .to_vec(); assert_eq!(leaders1.len(), leaders2.len()); @@ -392,7 +295,7 @@ mod tests { let schedule: Vec<_> = repeat_with(|| pubkeys[rng.gen_range(0..3)]) .take(19) .collect(); - let schedule = LeaderSchedule::new_from_schedule(schedule); + let schedule = LeaderSchedule::new_identity_keyed_with_schedule(schedule); let leaders = (0..NUM_SLOTS) .map(|i| (schedule[i as u64], i)) .into_group_map(); diff --git a/ledger/src/leader_schedule/identity_keyed.rs b/ledger/src/leader_schedule/identity_keyed.rs new file mode 100644 index 00000000000000..e575cfc1221e27 --- /dev/null +++ b/ledger/src/leader_schedule/identity_keyed.rs @@ -0,0 +1,49 @@ +use { + itertools::Itertools, + solana_pubkey::Pubkey, + solana_sdk::clock::Epoch, + std::{collections::HashMap, sync::Arc}, +}; + +#[derive(Default, Debug, PartialEq, Eq, Clone)] +pub(super) struct LeaderSchedule { + pub(super) slot_leaders: Vec, + // Inverted index from pubkeys to indices where they are the leader. + pub(super) index: HashMap>>, +} + +// Note: passing in zero stakers will cause a panic. +impl LeaderSchedule { + pub(super) fn new( + epoch_staked_nodes: &HashMap, + epoch: Epoch, + len: u64, + repeat: u64, + ) -> Self { + let keyed_stakes: Vec<_> = epoch_staked_nodes + .iter() + .map(|(pubkey, stake)| (pubkey, *stake)) + .collect(); + let slot_leaders = + super::LeaderSchedule::stake_weighted_slot_leaders(keyed_stakes, epoch, len, repeat); + Self::new_from_schedule(slot_leaders) + } + + pub(super) fn new_from_schedule(slot_leaders: Vec) -> Self { + Self { + index: Self::index_from_slot_leaders(&slot_leaders), + slot_leaders, + } + } + + fn index_from_slot_leaders(slot_leaders: &[Pubkey]) -> HashMap>> { + slot_leaders + .iter() + .enumerate() + .map(|(i, pk)| (*pk, i)) + .into_group_map() + .into_iter() + .map(|(k, v)| (k, Arc::new(v))) + .collect() + } +} diff --git a/ledger/src/leader_schedule/vote_keyed.rs b/ledger/src/leader_schedule/vote_keyed.rs new file mode 100644 index 00000000000000..814f3261885be6 --- /dev/null +++ b/ledger/src/leader_schedule/vote_keyed.rs @@ -0,0 +1,68 @@ +use { + super::identity_keyed, solana_pubkey::Pubkey, solana_sdk::clock::Epoch, + solana_vote::vote_account::VoteAccountsHashMap, +}; + +#[derive(Debug, PartialEq, Eq, Clone)] +pub(super) struct LeaderSchedule { + pub(super) slot_leader_vote_account_addresses: Vec, + // cached leader schedule keyed by validator identities created by mapping + // vote account addresses to the validator identity designated at the time + // of leader schedule generation. This is used to avoid the need to look up + // the validator identity address for each slot. + pub(super) identity_keyed_leader_schedule: identity_keyed::LeaderSchedule, +} + +impl LeaderSchedule { + pub(super) fn new( + vote_accounts_map: &VoteAccountsHashMap, + epoch: Epoch, + len: u64, + repeat: u64, + ) -> Self { + let keyed_stakes: Vec<_> = vote_accounts_map + .iter() + .map(|(vote_pubkey, (stake, _account))| (vote_pubkey, *stake)) + .collect(); + let slot_leader_vote_account_addresses = + super::LeaderSchedule::stake_weighted_slot_leaders(keyed_stakes, epoch, len, repeat); + + let identity_keyed_leader_schedule = { + struct SlotLeaderInfo<'a> { + vote_account_address: &'a Pubkey, + validator_identity_address: &'a Pubkey, + } + + let default_pubkey = Pubkey::default(); + let mut current_slot_leader_info = SlotLeaderInfo { + vote_account_address: &default_pubkey, + validator_identity_address: &default_pubkey, + }; + + let slot_leaders: Vec = slot_leader_vote_account_addresses + .iter() + .map(|vote_account_address| { + if vote_account_address != current_slot_leader_info.vote_account_address { + let validator_identity_address = vote_accounts_map + .get(vote_account_address) + .unwrap() + .1 + .node_pubkey(); + current_slot_leader_info = SlotLeaderInfo { + vote_account_address, + validator_identity_address, + }; + } + *current_slot_leader_info.validator_identity_address + }) + .collect(); + + identity_keyed::LeaderSchedule::new_from_schedule(slot_leaders) + }; + + Self { + slot_leader_vote_account_addresses, + identity_keyed_leader_schedule, + } + } +} diff --git a/ledger/src/leader_schedule_utils.rs b/ledger/src/leader_schedule_utils.rs index e39f179ee615a9..313e3b7f28ba0e 100644 --- a/ledger/src/leader_schedule_utils.rs +++ b/ledger/src/leader_schedule_utils.rs @@ -13,7 +13,7 @@ pub fn leader_schedule(epoch: Epoch, bank: &Bank) -> Option { let use_new_leader_schedule = bank.should_use_vote_address_leader_schedule(epoch)?; if use_new_leader_schedule { bank.epoch_vote_accounts(epoch).map(|vote_accounts_map| { - LeaderSchedule::new_keyed_by_vote_account( + LeaderSchedule::new_vote_keyed( vote_accounts_map, epoch, bank.get_slots_in_epoch(epoch), @@ -22,7 +22,7 @@ pub fn leader_schedule(epoch: Epoch, bank: &Bank) -> Option { }) } else { bank.epoch_staked_nodes(epoch).map(|stakes| { - LeaderSchedule::new_keyed_by_validator_identity( + LeaderSchedule::new_identity_keyed( &stakes, epoch, bank.get_slots_in_epoch(epoch), @@ -100,7 +100,7 @@ mod tests { let leader_schedule = leader_schedule(0, &bank).unwrap(); assert_eq!( - leader_schedule.is_keyed_by_vote_account(), + leader_schedule.is_vote_keyed(), use_vote_keyed_leader_schedule ); diff --git a/local-cluster/src/integration_tests.rs b/local-cluster/src/integration_tests.rs index e0f21d9efbe6d1..568ee1bc304f0a 100644 --- a/local-cluster/src/integration_tests.rs +++ b/local-cluster/src/integration_tests.rs @@ -273,7 +273,7 @@ pub fn create_custom_leader_schedule( } info!("leader_schedule: {}", leader_schedule.len()); - LeaderSchedule::new_from_schedule(leader_schedule) + LeaderSchedule::new_identity_keyed_with_schedule(leader_schedule) } pub fn create_custom_leader_schedule_with_random_keys( diff --git a/poh/src/poh_recorder.rs b/poh/src/poh_recorder.rs index 81dd20f0d4d802..7132350e132ab6 100644 --- a/poh/src/poh_recorder.rs +++ b/poh/src/poh_recorder.rs @@ -1884,7 +1884,9 @@ mod tests { let mut leader_schedule_cache = LeaderScheduleCache::new_from_bank(&bank); let fixed_schedule = solana_ledger::leader_schedule::FixedSchedule { leader_schedule: Arc::new( - solana_ledger::leader_schedule::LeaderSchedule::new_from_schedule(slot_leaders), + solana_ledger::leader_schedule::LeaderSchedule::new_identity_keyed_with_schedule( + slot_leaders, + ), ), }; leader_schedule_cache.set_fixed_leader_schedule(Some(fixed_schedule));