Add primitive ActiveStakers and LeaderSchedule objects

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<Pubkey, u64>,
+}
+
+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<Pubkey, u64> {
+        self.stakes.clone()
+    }
+
+    /// Return the pubkeys of each staker.
+    pub fn stakers(&self) -> HashSet<Pubkey> {
+        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)
+    }
+}

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<Pubkey>,
+}
+
+impl LeaderSchedule {
+    pub fn new(slot_leaders: Vec<Pubkey>) -> Self {
+        Self { slot_leaders }
+    }
+}
+
+impl Index<usize> 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);
+    }
+}

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))
+        ActiveStakers::new_with_upper_bound(&bank, lower_bound, upper_bound).stakers()
-            .iter()
-            .map(|vote_state| vote_state.staker_id)
-            .collect()
     }
 
     // Updates the leader schedule to include ticks from tick_height to the first tick of the next epoch

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;