From 4f61bd4212a2ff7b8abcd2e988aa869907af8874 Mon Sep 17 00:00:00 2001 From: HaoranYi <219428+HaoranYi@users.noreply.github.com> Date: Wed, 15 Oct 2025 15:36:48 +0000 Subject: [PATCH 1/2] Remove Arc from leader_slots_map in LeaderSchedule Changed `HashMap>>` to `HashMap>` in the leader schedule implementation. The Arc was previously needed because `get_leader_upcoming_slots` returned a boxed iterator requiring ownership of the Vec. By modifying the method to return a borrowed iterator (`Box`), and collecting the Arc references upfront in `next_leader_slot`, the schedules stay alive for the entire iterator chain, eliminating the need for Arc wrapping and avoiding intermediate Vec allocations. Changes: - Updated trait method signature and implementations - Simplified `invert_slot_leaders` to return HashMap without Arc wrapping - Modified `get_leader_upcoming_slots` to return borrowed iterator - Collect leader schedules in `next_leader_slot` before creating iterators - Removed Arc imports from identity_keyed and vote_keyed modules --- ledger/src/leader_schedule.rs | 44 ++++++++++---------- ledger/src/leader_schedule/identity_keyed.rs | 11 ++--- ledger/src/leader_schedule/vote_keyed.rs | 4 +- ledger/src/leader_schedule_cache.rs | 15 ++++--- 4 files changed, 39 insertions(+), 35 deletions(-) diff --git a/ledger/src/leader_schedule.rs b/ledger/src/leader_schedule.rs index 9fdf5b24683..eb718103e7e 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,31 @@ pub trait LeaderScheduleVariant: &self, pubkey: &Pubkey, offset: usize, // Starting index. - ) -> Box> { + ) -> Box + '_> { let index = self .get_leader_slots_map() - .get(pubkey) - .cloned() - .unwrap_or_default(); + .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 2eef2471d3a..5def3702f21 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 1314885f0b4..c3bc7a1164f 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 af63c76d26c..67e6c43d10e 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) From 63cee943b85220ebc13fda87c095228ec2ba4039 Mon Sep 17 00:00:00 2001 From: HaoranYi <219428+HaoranYi@users.noreply.github.com> Date: Wed, 15 Oct 2025 16:08:58 +0000 Subject: [PATCH 2/2] fmt --- ledger/src/leader_schedule.rs | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/ledger/src/leader_schedule.rs b/ledger/src/leader_schedule.rs index eb718103e7e..bf663ddcf7d 100644 --- a/ledger/src/leader_schedule.rs +++ b/ledger/src/leader_schedule.rs @@ -39,9 +39,7 @@ pub trait LeaderScheduleVariant: pubkey: &Pubkey, offset: usize, // Starting index. ) -> Box + '_> { - let index = self - .get_leader_slots_map() - .get(pubkey); + let index = self.get_leader_slots_map().get(pubkey); let num_slots = self.num_slots(); match index { @@ -55,7 +53,10 @@ pub trait LeaderScheduleVariant: // 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)) + Box::new( + (start_offset..=usize::MAX) + .map(move |k| index[k % size] + k / size * num_slots), + ) } _ => { // Empty iterator for pubkeys not in schedule