diff --git a/ledger/src/leader_schedule.rs b/ledger/src/leader_schedule.rs index 9fdf5b24683770..bf663ddcf7dd9b 100644 --- a/ledger/src/leader_schedule.rs +++ b/ledger/src/leader_schedule.rs @@ -26,7 +26,7 @@ pub trait LeaderScheduleVariant: std::fmt::Debug + Send + Sync + Index { fn get_slot_leaders(&self) -> &[Pubkey]; - fn get_leader_slots_map(&self) -> &HashMap>>; + fn get_leader_slots_map(&self) -> &HashMap>; /// Get the vote account address for the given epoch slot index. This is /// guaranteed to be Some if the leader schedule is keyed by vote account @@ -38,29 +38,32 @@ pub trait LeaderScheduleVariant: &self, pubkey: &Pubkey, offset: usize, // Starting index. - ) -> Box> { - let index = self - .get_leader_slots_map() - .get(pubkey) - .cloned() - .unwrap_or_default(); + ) -> Box + '_> { + let index = self.get_leader_slots_map().get(pubkey); let num_slots = self.num_slots(); - let size = index.len(); - #[allow(clippy::reversed_empty_ranges)] - let range = if index.is_empty() { - 1..=0 // Intentionally empty range of type RangeInclusive. - } else { - let offset = index - .binary_search(&(offset % num_slots)) - .unwrap_or_else(identity) - + offset / num_slots * size; - offset..=usize::MAX - }; - // The modular arithmetic here and above replicate Index implementation - // for LeaderSchedule, where the schedule keeps repeating endlessly. - // The '%' returns where in a cycle we are and the '/' returns how many - // times the schedule is repeated. - Box::new(range.map(move |k| index[k % size] + k / size * num_slots)) + + match index { + Some(index) if !index.is_empty() => { + let size = index.len(); + let start_offset = index + .binary_search(&(offset % num_slots)) + .unwrap_or_else(identity) + + offset / num_slots * size; + // The modular arithmetic here and above replicate Index implementation + // for LeaderSchedule, where the schedule keeps repeating endlessly. + // The '%' returns where in a cycle we are and the '/' returns how many + // times the schedule is repeated. + Box::new( + (start_offset..=usize::MAX) + .map(move |k| index[k % size] + k / size * num_slots), + ) + } + _ => { + // Empty iterator for pubkeys not in schedule + #[allow(clippy::reversed_empty_ranges)] + Box::new((1..=0).map(|_| 0)) + } + } } fn num_slots(&self) -> usize { diff --git a/ledger/src/leader_schedule/identity_keyed.rs b/ledger/src/leader_schedule/identity_keyed.rs index 2eef2471d3a7fd..5def3702f21ac2 100644 --- a/ledger/src/leader_schedule/identity_keyed.rs +++ b/ledger/src/leader_schedule/identity_keyed.rs @@ -3,14 +3,14 @@ use { itertools::Itertools, solana_clock::Epoch, solana_pubkey::Pubkey, - std::{collections::HashMap, ops::Index, sync::Arc}, + std::{collections::HashMap, ops::Index}, }; #[derive(Default, Debug, PartialEq, Eq, Clone)] pub struct LeaderSchedule { slot_leaders: Vec, // Inverted index from pubkeys to indices where they are the leader. - leader_slots_map: HashMap>>, + leader_slots_map: HashMap>, } impl LeaderSchedule { @@ -36,15 +36,12 @@ impl LeaderSchedule { } } - fn invert_slot_leaders(slot_leaders: &[Pubkey]) -> HashMap>> { + fn invert_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 get_slot_leaders(&self) -> &[Pubkey] { @@ -57,7 +54,7 @@ impl LeaderScheduleVariant for LeaderSchedule { &self.slot_leaders } - fn get_leader_slots_map(&self) -> &HashMap>> { + fn get_leader_slots_map(&self) -> &HashMap> { &self.leader_slots_map } } diff --git a/ledger/src/leader_schedule/vote_keyed.rs b/ledger/src/leader_schedule/vote_keyed.rs index 1314885f0b46e3..c3bc7a1164fd0b 100644 --- a/ledger/src/leader_schedule/vote_keyed.rs +++ b/ledger/src/leader_schedule/vote_keyed.rs @@ -3,7 +3,7 @@ use { solana_clock::Epoch, solana_pubkey::Pubkey, solana_vote::vote_account::VoteAccountsHashMap, - std::{collections::HashMap, ops::Index, sync::Arc}, + std::{collections::HashMap, ops::Index}, }; #[derive(Debug, PartialEq, Eq, Clone)] @@ -80,7 +80,7 @@ impl LeaderScheduleVariant for LeaderSchedule { self.identity_keyed_leader_schedule.get_slot_leaders() } - fn get_leader_slots_map(&self) -> &HashMap>> { + fn get_leader_slots_map(&self) -> &HashMap> { self.identity_keyed_leader_schedule.get_leader_slots_map() } diff --git a/ledger/src/leader_schedule_cache.rs b/ledger/src/leader_schedule_cache.rs index af63c76d26c86e..67e6c43d10eaba 100644 --- a/ledger/src/leader_schedule_cache.rs +++ b/ledger/src/leader_schedule_cache.rs @@ -122,15 +122,20 @@ impl LeaderScheduleCache { ); return None; } - // Slots after current_slot where pubkey is the leader. - let mut schedule = (epoch..=max_epoch) + // Collect leader schedules first so they stay alive for the iterator chain + let schedules: Vec<_> = (epoch..=max_epoch) .map(|epoch| self.get_epoch_schedule_else_compute(epoch, bank)) .while_some() .zip(epoch..) + .collect(); + + // Slots after current_slot where pubkey is the leader. + let mut schedule = schedules + .iter() .flat_map(|(leader_schedule, k)| { - let offset = if k == epoch { start_index as usize } else { 0 }; - let num_slots = bank.get_slots_in_epoch(k) as usize; - let first_slot = bank.epoch_schedule().get_first_slot_in_epoch(k); + let offset = if *k == epoch { start_index as usize } else { 0 }; + let num_slots = bank.get_slots_in_epoch(*k) as usize; + let first_slot = bank.epoch_schedule().get_first_slot_in_epoch(*k); leader_schedule .get_leader_upcoming_slots(pubkey, offset) .take_while(move |i| *i < num_slots)