2019-02-20 18:21:08 -08:00
|
|
|
use rand::distributions::{Distribution, WeightedIndex};
|
|
|
|
use rand::SeedableRng;
|
|
|
|
use rand_chacha::ChaChaRng;
|
2019-02-20 18:49:29 -08:00
|
|
|
use solana_runtime::bank::Bank;
|
2019-02-19 19:16:27 -08:00
|
|
|
use solana_sdk::pubkey::Pubkey;
|
|
|
|
use std::ops::Index;
|
|
|
|
|
|
|
|
/// Round-robin leader schedule.
|
2019-02-20 18:21:08 -08:00
|
|
|
#[derive(Debug, PartialEq)]
|
2019-02-19 19:16:27 -08:00
|
|
|
pub struct LeaderSchedule {
|
|
|
|
slot_leaders: Vec<Pubkey>,
|
|
|
|
}
|
|
|
|
|
|
|
|
impl LeaderSchedule {
|
2019-02-20 18:49:29 -08:00
|
|
|
pub fn new(ids_and_stakes: &[(Pubkey, u64)], seed: &[u8; 32], len: u64) -> Self {
|
2019-02-20 18:21:08 -08:00
|
|
|
let (pubkeys, stakes): (Vec<Pubkey>, Vec<u64>) = 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();
|
2019-02-20 18:49:29 -08:00
|
|
|
let slot_leaders = (0..len)
|
2019-02-20 18:21:08 -08:00
|
|
|
.map(|_| pubkeys[weighted_index.sample(&mut rng)])
|
|
|
|
.collect();
|
|
|
|
|
2019-02-19 19:16:27 -08:00
|
|
|
Self { slot_leaders }
|
|
|
|
}
|
2019-02-20 18:49:29 -08:00
|
|
|
|
2019-02-21 10:59:03 -08:00
|
|
|
pub fn new_with_bank(bank: &Bank) -> Self {
|
2019-02-20 18:49:29 -08:00
|
|
|
let mut seed = [0u8; 32];
|
|
|
|
seed.copy_from_slice(bank.last_id().as_ref());
|
2019-02-23 21:56:41 -08:00
|
|
|
let stakes: Vec<_> = bank.staked_nodes().into_iter().collect();
|
|
|
|
Self::new(&stakes, &seed, bank.slots_per_epoch())
|
2019-02-20 18:49:29 -08:00
|
|
|
}
|
2019-02-19 19:16:27 -08:00
|
|
|
}
|
|
|
|
|
|
|
|
impl Index<usize> for LeaderSchedule {
|
|
|
|
type Output = Pubkey;
|
|
|
|
fn index(&self, index: usize) -> &Pubkey {
|
|
|
|
&self.slot_leaders[index % self.slot_leaders.len()]
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2019-02-21 12:32:25 -08:00
|
|
|
pub trait LeaderScheduleUtil {
|
2019-02-21 10:59:03 -08:00
|
|
|
/// Return the leader schedule for the current epoch.
|
|
|
|
fn leader_schedule(&self) -> LeaderSchedule;
|
2019-02-23 13:43:34 -08:00
|
|
|
|
|
|
|
/// Return the leader id for the current slot.
|
|
|
|
fn slot_leader(&self) -> Pubkey;
|
2019-02-21 10:59:03 -08:00
|
|
|
}
|
|
|
|
|
|
|
|
impl LeaderScheduleUtil for Bank {
|
|
|
|
fn leader_schedule(&self) -> LeaderSchedule {
|
|
|
|
match self.leader_schedule_bank() {
|
|
|
|
None => LeaderSchedule::new_with_bank(self),
|
|
|
|
Some(bank) => LeaderSchedule::new_with_bank(&bank),
|
|
|
|
}
|
|
|
|
}
|
2019-02-23 13:43:34 -08:00
|
|
|
|
|
|
|
fn slot_leader(&self) -> Pubkey {
|
|
|
|
self.leader_schedule()[self.slot_index() as usize]
|
|
|
|
}
|
2019-02-21 10:59:03 -08:00
|
|
|
}
|
|
|
|
|
2019-02-19 19:16:27 -08:00
|
|
|
#[cfg(test)]
|
|
|
|
mod tests {
|
|
|
|
use super::*;
|
2019-02-20 18:49:29 -08:00
|
|
|
use solana_sdk::genesis_block::GenesisBlock;
|
2019-02-19 19:16:27 -08:00
|
|
|
use solana_sdk::signature::{Keypair, KeypairUtil};
|
2019-02-21 10:59:03 -08:00
|
|
|
use std::iter;
|
2019-02-20 18:21:08 -08:00
|
|
|
|
2019-02-19 19:16:27 -08:00
|
|
|
#[test]
|
|
|
|
fn test_leader_schedule_index() {
|
|
|
|
let pubkey0 = Keypair::new().pubkey();
|
|
|
|
let pubkey1 = Keypair::new().pubkey();
|
2019-02-20 18:21:08 -08:00
|
|
|
let leader_schedule = LeaderSchedule {
|
|
|
|
slot_leaders: vec![pubkey0, pubkey1],
|
|
|
|
};
|
2019-02-19 19:16:27 -08:00
|
|
|
assert_eq!(leader_schedule[0], pubkey0);
|
|
|
|
assert_eq!(leader_schedule[1], pubkey1);
|
|
|
|
assert_eq!(leader_schedule[2], pubkey0);
|
|
|
|
}
|
2019-02-20 18:21:08 -08:00
|
|
|
|
|
|
|
#[test]
|
2019-02-20 18:49:29 -08:00
|
|
|
fn test_leader_schedule_basic() {
|
2019-02-20 18:21:08 -08:00
|
|
|
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());
|
2019-02-20 18:49:29 -08:00
|
|
|
let len = num_keys * 10;
|
|
|
|
let leader_schedule = LeaderSchedule::new(&stakes, &seed_bytes, len);
|
|
|
|
let leader_schedule2 = LeaderSchedule::new(&stakes, &seed_bytes, len);
|
|
|
|
assert_eq!(leader_schedule.slot_leaders.len() as u64, len);
|
2019-02-20 18:21:08 -08:00
|
|
|
// Check that the same schedule is reproducibly generated
|
|
|
|
assert_eq!(leader_schedule, leader_schedule2);
|
|
|
|
}
|
2019-02-20 18:49:29 -08:00
|
|
|
|
|
|
|
#[test]
|
|
|
|
fn test_leader_schedule_via_bank() {
|
|
|
|
let pubkey = Keypair::new().pubkey();
|
|
|
|
let (genesis_block, _mint_keypair) = GenesisBlock::new_with_leader(2, pubkey, 2);
|
|
|
|
let bank = Bank::new(&genesis_block);
|
2019-02-21 10:59:03 -08:00
|
|
|
let leader_schedule = LeaderSchedule::new_with_bank(&bank);
|
|
|
|
let len = bank.slots_per_epoch() as usize;
|
|
|
|
let expected: Vec<_> = iter::repeat(pubkey).take(len).collect();
|
|
|
|
assert_eq!(leader_schedule.slot_leaders, expected);
|
|
|
|
assert_eq!(bank.leader_schedule().slot_leaders, expected); // Same thing, but with the trait
|
2019-02-20 18:49:29 -08:00
|
|
|
}
|
2019-02-23 13:43:34 -08:00
|
|
|
|
|
|
|
#[test]
|
|
|
|
fn test_leader_schedule_slot_leader_basic() {
|
|
|
|
let pubkey = Keypair::new().pubkey();
|
|
|
|
let bank = Bank::new(&GenesisBlock::new_with_leader(2, pubkey, 2).0);
|
|
|
|
assert_eq!(bank.slot_leader(), pubkey);
|
|
|
|
}
|
2019-02-19 19:16:27 -08:00
|
|
|
}
|