diff --git a/src/active_stakers.rs b/src/active_stakers.rs index 25ffe3bdba..87719f9614 100644 --- a/src/active_stakers.rs +++ b/src/active_stakers.rs @@ -1,4 +1,3 @@ -use crate::leader_schedule::LeaderSchedule; use solana_runtime::bank::Bank; use solana_sdk::pubkey::Pubkey; use solana_sdk::timing::{DEFAULT_SLOTS_PER_EPOCH, DEFAULT_TICKS_PER_SLOT}; @@ -60,7 +59,7 @@ impl ActiveStakers { Self::new_with_bounds(bank, DEFAULT_ACTIVE_WINDOW_TICK_LENGTH, bank.tick_height()) } - pub fn ranked_stakes(&self) -> Vec<(Pubkey, u64)> { + pub fn sorted_stakes(&self) -> Vec<(Pubkey, u64)> { self.stakes.clone() } @@ -75,10 +74,6 @@ impl ActiveStakers { pubkeys.sort_unstable(); pubkeys } - - pub fn leader_schedule(&self) -> LeaderSchedule { - LeaderSchedule::new(self.pubkeys()) - } } #[cfg(test)] diff --git a/src/leader_schedule.rs b/src/leader_schedule.rs index a0cb6ebb5e..f26344d745 100644 --- a/src/leader_schedule.rs +++ b/src/leader_schedule.rs @@ -1,13 +1,29 @@ +use rand::distributions::{Distribution, WeightedIndex}; +use rand::SeedableRng; +use rand_chacha::ChaChaRng; use solana_sdk::pubkey::Pubkey; use std::ops::Index; /// Round-robin leader schedule. +#[derive(Debug, PartialEq)] pub struct LeaderSchedule { slot_leaders: Vec, } impl LeaderSchedule { - pub fn new(slot_leaders: Vec) -> Self { + pub fn new(ids_and_stakes: &[(Pubkey, u64)], seed: &[u8; 32], slots_per_epoch: u64) -> Self { + let (pubkeys, stakes): (Vec, Vec) = ids_and_stakes + .iter() + .map(|&(ref id, ref stake)| (id, stake)) + .unzip(); + + // Should have no zero weighted stakes + let mut rng = ChaChaRng::from_seed(*seed); + let weighted_index = WeightedIndex::new(stakes).unwrap(); + let slot_leaders = (0..slots_per_epoch) + .map(|_| pubkeys[weighted_index.sample(&mut rng)]) + .collect(); + Self { slot_leaders } } } @@ -23,13 +39,34 @@ impl Index for LeaderSchedule { mod tests { use super::*; use solana_sdk::signature::{Keypair, KeypairUtil}; + #[test] fn test_leader_schedule_index() { let pubkey0 = Keypair::new().pubkey(); let pubkey1 = Keypair::new().pubkey(); - let leader_schedule = LeaderSchedule::new(vec![pubkey0, pubkey1]); + let leader_schedule = LeaderSchedule { + slot_leaders: vec![pubkey0, pubkey1], + }; assert_eq!(leader_schedule[0], pubkey0); assert_eq!(leader_schedule[1], pubkey1); assert_eq!(leader_schedule[2], pubkey0); } + + #[test] + fn test_new_leader_schedule() { + let num_keys = 10; + let stakes: Vec<_> = (0..num_keys) + .map(|i| (Keypair::new().pubkey(), i)) + .collect(); + + let seed = Keypair::new().pubkey(); + let mut seed_bytes = [0u8; 32]; + seed_bytes.copy_from_slice(seed.as_ref()); + let slots_per_epoch = num_keys * 10; + let leader_schedule = LeaderSchedule::new(&stakes, &seed_bytes, slots_per_epoch); + let leader_schedule2 = LeaderSchedule::new(&stakes, &seed_bytes, slots_per_epoch); + assert_eq!(leader_schedule.slot_leaders.len() as u64, slots_per_epoch); + // Check that the same schedule is reproducibly generated + assert_eq!(leader_schedule, leader_schedule2); + } } diff --git a/src/leader_scheduler.rs b/src/leader_scheduler.rs index 113086bcc1..c1bf6dc35a 100644 --- a/src/leader_scheduler.rs +++ b/src/leader_scheduler.rs @@ -230,7 +230,7 @@ impl LeaderScheduler { self.seed = Self::calculate_seed(tick_height); let ranked_active_set = ActiveStakers::new_with_bounds(&bank, self.active_window_tick_length, tick_height) - .ranked_stakes(); + .sorted_stakes(); if ranked_active_set.is_empty() { info!(