diff --git a/src/active_stakers.rs b/src/active_stakers.rs new file mode 100644 index 0000000000..b45f1168b1 --- /dev/null +++ b/src/active_stakers.rs @@ -0,0 +1,58 @@ +use crate::leader_schedule::LeaderSchedule; +use hashbrown::{HashMap, HashSet}; +use solana_runtime::bank::Bank; +use solana_sdk::pubkey::Pubkey; +use solana_sdk::vote_program::VoteState; + +// 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 + .votes + .back() + .filter(|vote| vote.tick_height >= lower_bound && vote.tick_height <= upper_bound) + .is_some() +} + +/// The set of stakers that have voted near the time of construction +pub struct ActiveStakers { + stakes: HashMap, +} + +impl ActiveStakers { + pub fn new_with_upper_bound(bank: &Bank, lower_bound: u64, upper_bound: u64) -> Self { + let stakes = bank + .vote_states(|vote_state| is_active_staker(vote_state, lower_bound, upper_bound)) + .iter() + .filter_map(|vote_state| { + let pubkey = vote_state.staker_id; + let stake = bank.get_balance(&pubkey); + if stake > 0 { + Some((pubkey, stake)) + } else { + None + } + }) + .collect(); + Self { stakes } + } + + pub fn new(bank: &Bank, lower_bound: u64) -> Self { + Self::new_with_upper_bound(bank, lower_bound, bank.tick_height()) + } + + /// Return a map from staker pubkeys to their respective stakes. + pub fn stakes(&self) -> HashMap { + self.stakes.clone() + } + + /// Return the pubkeys of each staker. + pub fn stakers(&self) -> HashSet { + self.stakes.keys().cloned().collect() + } + + pub fn leader_schedule(&self) -> LeaderSchedule { + let mut stakers: Vec<_> = self.stakes.keys().cloned().collect(); + stakers.sort(); + LeaderSchedule::new(stakers) + } +} diff --git a/src/leader_schedule.rs b/src/leader_schedule.rs new file mode 100644 index 0000000000..a0cb6ebb5e --- /dev/null +++ b/src/leader_schedule.rs @@ -0,0 +1,35 @@ +use solana_sdk::pubkey::Pubkey; +use std::ops::Index; + +/// Round-robin leader schedule. +pub struct LeaderSchedule { + slot_leaders: Vec, +} + +impl LeaderSchedule { + pub fn new(slot_leaders: Vec) -> Self { + Self { slot_leaders } + } +} + +impl Index for LeaderSchedule { + type Output = Pubkey; + fn index(&self, index: usize) -> &Pubkey { + &self.slot_leaders[index % self.slot_leaders.len()] + } +} + +#[cfg(test)] +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]); + assert_eq!(leader_schedule[0], pubkey0); + assert_eq!(leader_schedule[1], pubkey1); + assert_eq!(leader_schedule[2], pubkey0); + } +} diff --git a/src/leader_scheduler.rs b/src/leader_scheduler.rs index 0d7b702b9f..7519576e52 100644 --- a/src/leader_scheduler.rs +++ b/src/leader_scheduler.rs @@ -1,6 +1,7 @@ //! The `leader_scheduler` module implements a structure and functions for tracking and //! managing the schedule for leader rotation +use crate::active_stakers::ActiveStakers; use crate::entry::{create_ticks, next_entry_mut, Entry}; use crate::voting_keypair::VotingKeypair; use bincode::serialize; @@ -12,7 +13,6 @@ use solana_sdk::pubkey::Pubkey; use solana_sdk::signature::{Keypair, KeypairUtil}; use solana_sdk::system_transaction::SystemTransaction; use solana_sdk::timing::{DEFAULT_SLOTS_PER_EPOCH, DEFAULT_TICKS_PER_SLOT}; -use solana_sdk::vote_program::VoteState; use solana_sdk::vote_transaction::VoteTransaction; use std::io::Cursor; use std::sync::Arc; @@ -202,15 +202,6 @@ impl LeaderScheduler { } } - // 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 - .votes - .back() - .filter(|vote| vote.tick_height >= lower_bound && vote.tick_height <= upper_bound) - .is_some() - } - // TODO: We use a HashSet for now because a single validator could potentially register // multiple vote account. Once that is no longer possible (see the TODO in vote_program.rs, // process_transaction(), case VoteInstruction::RegisterAccount), we can use a vector. @@ -223,10 +214,7 @@ impl LeaderScheduler { upper_bound ); - bank.vote_states(|vote_state| Self::is_active_staker(vote_state, lower_bound, upper_bound)) - .iter() - .map(|vote_state| vote_state.staker_id) - .collect() + ActiveStakers::new_with_upper_bound(&bank, lower_bound, upper_bound).stakers() } // Updates the leader schedule to include ticks from tick_height to the first tick of the next epoch diff --git a/src/lib.rs b/src/lib.rs index 34eb4b324f..02c0772b4d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,6 +7,7 @@ //! #![cfg_attr(feature = "unstable", feature(test))] +pub mod active_stakers; pub mod bank_forks; pub mod banking_stage; pub mod blob_fetch_stage; @@ -39,6 +40,7 @@ pub mod fullnode; pub mod gen_keys; pub mod gossip_service; pub mod leader_confirmation_service; +pub mod leader_schedule; pub mod leader_scheduler; pub mod local_vote_signer_service; pub mod packet;