diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs index 7195dde8f4..08e9242d47 100644 --- a/runtime/src/bank.rs +++ b/runtime/src/bank.rs @@ -663,48 +663,6 @@ impl Bank { .collect() } - fn sort_stakes(stakes: &mut Vec<(Pubkey, u64)>) { - // Sort first by stake. If stakes are the same, sort by pubkey to ensure a - // deterministic result. - // Note: Use unstable sort, because we dedup right after to remove the equal elements. - stakes.sort_unstable_by(|(pubkey0, stake0), (pubkey1, stake1)| { - if stake0 == stake1 { - pubkey0.cmp(&pubkey1) - } else { - stake0.cmp(&stake1) - } - }); - - // Now that it's sorted, we can do an O(n) dedup. - stakes.dedup(); - } - - /// Return a sorted, filtered list of staker/stake pairs. - pub fn filtered_stakes(&self, filter: F) -> Vec<(Pubkey, u64)> - where - F: Fn(&VoteState) -> bool, - { - let mut stakes: Vec<_> = self - .vote_states(filter) - .iter() - .filter_map(|vote_state| { - let pubkey = vote_state.staker_id; - let stake = self.get_balance(&pubkey); - if stake > 0 { - Some((pubkey, stake)) - } else { - None - } - }) - .collect(); - Self::sort_stakes(&mut stakes); - stakes - } - - pub fn stakes(&self) -> Vec<(Pubkey, u64)> { - self.filtered_stakes(|_| true) - } - /// Return the number of ticks per slot that should be used calls to slot_height(). pub fn ticks_per_slot(&self) -> u64 { self.ticks_per_slot @@ -1407,31 +1365,4 @@ mod tests { bank.merge_parents(); } } - - #[test] - fn test_sort_stakes_basic() { - let pubkey0 = Keypair::new().pubkey(); - let pubkey1 = Keypair::new().pubkey(); - let mut stakes = vec![(pubkey0, 2), (pubkey1, 1)]; - Bank::sort_stakes(&mut stakes); - assert_eq!(stakes, vec![(pubkey1, 1), (pubkey0, 2)]); - } - - #[test] - fn test_sort_stakes_with_dup() { - let pubkey0 = Keypair::new().pubkey(); - let pubkey1 = Keypair::new().pubkey(); - let mut stakes = vec![(pubkey0, 1), (pubkey1, 2), (pubkey0, 1)]; - Bank::sort_stakes(&mut stakes); - assert_eq!(stakes, vec![(pubkey0, 1), (pubkey1, 2)]); - } - - #[test] - fn test_sort_stakes_with_equal_stakes() { - let pubkey0 = Pubkey::default(); - let pubkey1 = Keypair::new().pubkey(); - let mut stakes = vec![(pubkey0, 1), (pubkey1, 1)]; - Bank::sort_stakes(&mut stakes); - assert_eq!(stakes, vec![(pubkey0, 1), (pubkey1, 1)]); - } } diff --git a/src/leader_schedule.rs b/src/leader_schedule.rs index 768e883528..0c884bd2d0 100644 --- a/src/leader_schedule.rs +++ b/src/leader_schedule.rs @@ -31,7 +31,8 @@ impl LeaderSchedule { pub fn new_with_bank(bank: &Bank) -> Self { let mut seed = [0u8; 32]; seed.copy_from_slice(bank.last_id().as_ref()); - Self::new(&bank.stakes(), &seed, bank.slots_per_epoch()) + let stakes: Vec<_> = bank.staked_nodes().into_iter().collect(); + Self::new(&stakes, &seed, bank.slots_per_epoch()) } } diff --git a/src/leader_scheduler.rs b/src/leader_scheduler.rs index 385e5922b8..44bf924c18 100644 --- a/src/leader_scheduler.rs +++ b/src/leader_scheduler.rs @@ -49,6 +49,22 @@ impl Default for LeaderSchedulerConfig { } } +fn sort_stakes(stakes: &mut Vec<(Pubkey, u64)>) { + // Sort first by stake. If stakes are the same, sort by pubkey to ensure a + // deterministic result. + // Note: Use unstable sort, because we dedup right after to remove the equal elements. + stakes.sort_unstable_by(|(pubkey0, stake0), (pubkey1, stake1)| { + if stake0 == stake1 { + pubkey0.cmp(&pubkey1) + } else { + stake0.cmp(&stake1) + } + }); + + // Now that it's sorted, we can do an O(n) dedup. + stakes.dedup(); +} + // Return true of the latest vote is between the lower and upper bounds (inclusive) fn is_active_staker(vote_state: &VoteState, lower_bound: u64, upper_bound: u64) -> bool { vote_state @@ -58,13 +74,27 @@ fn is_active_staker(vote_state: &VoteState, lower_bound: u64, upper_bound: u64) .is_some() } +/// Return a sorted, filtered list of node_id/stake pairs. fn get_active_stakes( bank: &Bank, active_window_num_slots: u64, upper_bound: u64, ) -> Vec<(Pubkey, u64)> { let lower_bound = upper_bound.saturating_sub(active_window_num_slots); - bank.filtered_stakes(|vote_state| is_active_staker(vote_state, lower_bound, upper_bound)) + let mut stakes: Vec<_> = bank + .vote_states(|vote_state| is_active_staker(vote_state, lower_bound, upper_bound)) + .iter() + .filter_map(|vote_state| { + let stake = bank.get_balance(&vote_state.staker_id); + if stake > 0 { + Some((vote_state.node_id, stake)) + } else { + None + } + }) + .collect(); + sort_stakes(&mut stakes); + stakes } #[derive(Clone, Debug)] @@ -1137,4 +1167,31 @@ pub mod tests { run_consecutive_leader_test(2, false); run_consecutive_leader_test(10, false); } + + #[test] + fn test_sort_stakes_basic() { + let pubkey0 = Keypair::new().pubkey(); + let pubkey1 = Keypair::new().pubkey(); + let mut stakes = vec![(pubkey0, 2), (pubkey1, 1)]; + sort_stakes(&mut stakes); + assert_eq!(stakes, vec![(pubkey1, 1), (pubkey0, 2)]); + } + + #[test] + fn test_sort_stakes_with_dup() { + let pubkey0 = Keypair::new().pubkey(); + let pubkey1 = Keypair::new().pubkey(); + let mut stakes = vec![(pubkey0, 1), (pubkey1, 2), (pubkey0, 1)]; + sort_stakes(&mut stakes); + assert_eq!(stakes, vec![(pubkey0, 1), (pubkey1, 2)]); + } + + #[test] + fn test_sort_stakes_with_equal_stakes() { + let pubkey0 = Pubkey::default(); + let pubkey1 = Keypair::new().pubkey(); + let mut stakes = vec![(pubkey0, 1), (pubkey1, 1)]; + sort_stakes(&mut stakes); + assert_eq!(stakes, vec![(pubkey0, 1), (pubkey1, 1)]); + } }