solana_ledger/
leader_schedule.rs

1use {
2    rand::distributions::{Distribution, WeightedIndex},
3    rand_chacha::{rand_core::SeedableRng, ChaChaRng},
4    solana_clock::Epoch,
5    solana_pubkey::Pubkey,
6    std::{collections::HashMap, convert::identity, ops::Index, sync::Arc},
7};
8
9mod identity_keyed;
10mod vote_keyed;
11pub use {
12    identity_keyed::LeaderSchedule as IdentityKeyedLeaderSchedule,
13    vote_keyed::LeaderSchedule as VoteKeyedLeaderSchedule,
14};
15
16// Used for testing
17#[derive(Clone, Debug)]
18pub struct FixedSchedule {
19    pub leader_schedule: Arc<LeaderSchedule>,
20}
21
22/// Stake-weighted leader schedule for one epoch.
23pub type LeaderSchedule = Box<dyn LeaderScheduleVariant>;
24
25pub trait LeaderScheduleVariant:
26    std::fmt::Debug + Send + Sync + Index<u64, Output = Pubkey>
27{
28    fn get_slot_leaders(&self) -> &[Pubkey];
29    fn get_leader_slots_map(&self) -> &HashMap<Pubkey, Arc<Vec<usize>>>;
30
31    /// Get the vote account address for the given epoch slot index. This is
32    /// guaranteed to be Some if the leader schedule is keyed by vote account
33    fn get_vote_key_at_slot_index(&self, _epoch_slot_index: usize) -> Option<&Pubkey> {
34        None
35    }
36
37    fn get_leader_upcoming_slots(
38        &self,
39        pubkey: &Pubkey,
40        offset: usize, // Starting index.
41    ) -> Box<dyn Iterator<Item = usize>> {
42        let index = self
43            .get_leader_slots_map()
44            .get(pubkey)
45            .cloned()
46            .unwrap_or_default();
47        let num_slots = self.num_slots();
48        let size = index.len();
49        #[allow(clippy::reversed_empty_ranges)]
50        let range = if index.is_empty() {
51            1..=0 // Intentionally empty range of type RangeInclusive.
52        } else {
53            let offset = index
54                .binary_search(&(offset % num_slots))
55                .unwrap_or_else(identity)
56                + offset / num_slots * size;
57            offset..=usize::MAX
58        };
59        // The modular arithmetic here and above replicate Index implementation
60        // for LeaderSchedule, where the schedule keeps repeating endlessly.
61        // The '%' returns where in a cycle we are and the '/' returns how many
62        // times the schedule is repeated.
63        Box::new(range.map(move |k| index[k % size] + k / size * num_slots))
64    }
65
66    fn num_slots(&self) -> usize {
67        self.get_slot_leaders().len()
68    }
69}
70
71// Note: passing in zero keyed stakes will cause a panic.
72fn stake_weighted_slot_leaders(
73    mut keyed_stakes: Vec<(&Pubkey, u64)>,
74    epoch: Epoch,
75    len: u64,
76    repeat: u64,
77) -> Vec<Pubkey> {
78    sort_stakes(&mut keyed_stakes);
79    let (keys, stakes): (Vec<_>, Vec<_>) = keyed_stakes.into_iter().unzip();
80    let weighted_index = WeightedIndex::new(stakes).unwrap();
81    let mut seed = [0u8; 32];
82    seed[0..8].copy_from_slice(&epoch.to_le_bytes());
83    let rng = &mut ChaChaRng::from_seed(seed);
84    let mut current_slot_leader = Pubkey::default();
85    (0..len)
86        .map(|i| {
87            if i % repeat == 0 {
88                current_slot_leader = keys[weighted_index.sample(rng)];
89            }
90            current_slot_leader
91        })
92        .collect()
93}
94
95fn sort_stakes(stakes: &mut Vec<(&Pubkey, u64)>) {
96    // Sort first by stake. If stakes are the same, sort by pubkey to ensure a
97    // deterministic result.
98    // Note: Use unstable sort, because we dedup right after to remove the equal elements.
99    stakes.sort_unstable_by(|(l_pubkey, l_stake), (r_pubkey, r_stake)| {
100        if r_stake == l_stake {
101            r_pubkey.cmp(l_pubkey)
102        } else {
103            r_stake.cmp(l_stake)
104        }
105    });
106
107    // Now that it's sorted, we can do an O(n) dedup.
108    stakes.dedup();
109}
110
111#[cfg(test)]
112mod tests {
113    use {super::*, itertools::Itertools, rand::Rng, std::iter::repeat_with};
114
115    #[test]
116    fn test_get_leader_upcoming_slots() {
117        const NUM_SLOTS: usize = 97;
118        let mut rng = rand::thread_rng();
119        let pubkeys: Vec<_> = repeat_with(Pubkey::new_unique).take(4).collect();
120        let schedule: Vec<_> = repeat_with(|| pubkeys[rng.gen_range(0..3)])
121            .take(19)
122            .collect();
123        let schedule = IdentityKeyedLeaderSchedule::new_from_schedule(schedule);
124        let leaders = (0..NUM_SLOTS)
125            .map(|i| (schedule[i as u64], i))
126            .into_group_map();
127        for pubkey in &pubkeys {
128            let index = leaders.get(pubkey).cloned().unwrap_or_default();
129            for offset in 0..NUM_SLOTS {
130                let schedule: Vec<_> = schedule
131                    .get_leader_upcoming_slots(pubkey, offset)
132                    .take_while(|s| *s < NUM_SLOTS)
133                    .collect();
134                let index: Vec<_> = index.iter().copied().skip_while(|s| *s < offset).collect();
135                assert_eq!(schedule, index);
136            }
137        }
138    }
139
140    #[test]
141    fn test_sort_stakes_basic() {
142        let pubkey0 = solana_pubkey::new_rand();
143        let pubkey1 = solana_pubkey::new_rand();
144        let mut stakes = vec![(&pubkey0, 1), (&pubkey1, 2)];
145        sort_stakes(&mut stakes);
146        assert_eq!(stakes, vec![(&pubkey1, 2), (&pubkey0, 1)]);
147    }
148
149    #[test]
150    fn test_sort_stakes_with_dup() {
151        let pubkey0 = solana_pubkey::new_rand();
152        let pubkey1 = solana_pubkey::new_rand();
153        let mut stakes = vec![(&pubkey0, 1), (&pubkey1, 2), (&pubkey0, 1)];
154        sort_stakes(&mut stakes);
155        assert_eq!(stakes, vec![(&pubkey1, 2), (&pubkey0, 1)]);
156    }
157
158    #[test]
159    fn test_sort_stakes_with_equal_stakes() {
160        let pubkey0 = Pubkey::default();
161        let pubkey1 = solana_pubkey::new_rand();
162        let mut stakes = vec![(&pubkey0, 1), (&pubkey1, 1)];
163        sort_stakes(&mut stakes);
164        assert_eq!(stakes, vec![(&pubkey1, 1), (&pubkey0, 1)]);
165    }
166}
solana_ledger/leader_schedule.rs

solana_ledger/
leader_schedule.rs
