solana_leader_schedule/
lib.rs

1//! Solana leader schedule.
2
3#![cfg(feature = "agave-unstable-api")]
4#![allow(clippy::arithmetic_side_effects)]
5
6use {
7    agave_random::weighted::WeightedU64Index,
8    rand_chacha::{ChaChaRng, rand_core::SeedableRng},
9    solana_clock::Epoch,
10    solana_pubkey::Pubkey,
11    std::{iter, num::NonZeroUsize, sync::Arc},
12};
13
14mod vote_keyed;
15/// Stake-weighted leader schedule for one epoch.
16pub use vote_keyed::LeaderSchedule;
17
18/// Number of consecutive slots assigned to the same leader before the schedule
19/// advances to the next validator.
20pub const NUM_CONSECUTIVE_LEADER_SLOTS: NonZeroUsize = NonZeroUsize::new(4).unwrap();
21
22#[derive(Copy, Clone, Default, Debug, PartialEq, Eq)]
23pub struct SlotLeader {
24    pub id: Pubkey,
25    pub vote_address: Pubkey,
26}
27
28impl SlotLeader {
29    pub fn new_unique() -> Self {
30        SlotLeader {
31            id: Pubkey::new_unique(),
32            vote_address: Pubkey::new_unique(),
33        }
34    }
35}
36
37// Used for testing
38#[derive(Clone, Debug)]
39pub struct FixedSchedule {
40    pub leader_schedule: Arc<LeaderSchedule>,
41}
42
43// Note: passing in zero keyed stakes will cause a panic.
44fn stake_weighted_slot_leaders(
45    mut slot_leader_stakes: Vec<(SlotLeader, u64)>,
46    epoch: Epoch,
47    len: usize,
48    repeat: NonZeroUsize,
49) -> Vec<SlotLeader> {
50    let repeat = repeat.get();
51    debug_assert!(
52        len.is_multiple_of(repeat),
53        "expected `len` {len} to be divisible by `repeat` {repeat}"
54    );
55    sort_stakes(&mut slot_leader_stakes);
56    let (slot_leaders, stakes): (Vec<_>, Vec<_>) = slot_leader_stakes.into_iter().unzip();
57    let weighted_index = WeightedU64Index::new(stakes).unwrap();
58    let mut seed = [0u8; 32];
59    seed[0..8].copy_from_slice(&epoch.to_le_bytes());
60    let rng = &mut ChaChaRng::from_seed(seed);
61    iter::repeat_with(|| slot_leaders[weighted_index.sample(rng)])
62        .take(len / repeat)
63        .collect()
64}
65
66fn sort_stakes(stakes: &mut Vec<(SlotLeader, u64)>) {
67    // Sort first by stake. If stakes are the same, sort by vote address pubkey
68    // to ensure a deterministic result.
69    // Note: Use unstable sort, because we dedup right after to remove the equal
70    // elements.
71    stakes.sort_unstable_by(|(l_leader, l_stake), (r_leader, r_stake)| {
72        if r_stake == l_stake {
73            r_leader.vote_address.cmp(&l_leader.vote_address)
74        } else {
75            r_stake.cmp(l_stake)
76        }
77    });
78
79    // Now that it's sorted, we can do an O(n) dedup.
80    stakes.dedup_by(|(l_leader, l_stake), (r_leader, r_stake)| {
81        r_stake == l_stake && r_leader.vote_address == l_leader.vote_address
82    });
83}
84
85#[cfg(test)]
86mod tests {
87    use {super::*, itertools::Itertools, rand::Rng, std::iter::repeat_with, test_case::test_case};
88
89    const NZ_1: NonZeroUsize = NonZeroUsize::new(1).unwrap();
90    const NZ_2: NonZeroUsize = NonZeroUsize::new(2).unwrap();
91
92    #[test]
93    fn test_get_leader_upcoming_slots() {
94        const NUM_SLOTS: usize = 97;
95        let mut rng = rand::rng();
96        let unique_leaders: Vec<_> = repeat_with(SlotLeader::new_unique).take(4).collect();
97        let schedule: Vec<_> = repeat_with(|| unique_leaders[rng.random_range(0..3)])
98            .take(19)
99            .collect();
100        let schedule = LeaderSchedule::new_from_schedule(schedule, NZ_1);
101        let leaders = (0..NUM_SLOTS)
102            .map(|i| (schedule[i as u64].id, i))
103            .into_group_map();
104        for leader in &unique_leaders {
105            let index = leaders.get(&leader.id).cloned().unwrap_or_default();
106            for offset in 0..NUM_SLOTS {
107                let upcoming_slots: Vec<_> = schedule
108                    .get_leader_upcoming_slots(&leader.id, offset)
109                    .take_while(|s| *s < NUM_SLOTS)
110                    .collect();
111                let expected: Vec<_> = index.iter().copied().skip_while(|s| *s < offset).collect();
112                assert_eq!(upcoming_slots, expected);
113            }
114        }
115    }
116
117    #[test]
118    fn test_sort_stakes_basic() {
119        let leader0 = SlotLeader::new_rand();
120        let leader1 = SlotLeader::new_rand();
121        let mut stakes = vec![(leader0, 1), (leader1, 2)];
122        sort_stakes(&mut stakes);
123        assert_eq!(stakes, vec![(leader1, 2), (leader0, 1)]);
124    }
125
126    #[test]
127    fn test_sort_stakes_with_dup() {
128        let leader0 = SlotLeader::new_rand();
129        let leader1 = SlotLeader::new_rand();
130        let mut stakes = vec![(leader0, 1), (leader1, 2), (leader0, 1)];
131        sort_stakes(&mut stakes);
132        assert_eq!(stakes, vec![(leader1, 2), (leader0, 1)]);
133    }
134
135    #[test]
136    fn test_sort_stakes_with_equal_stakes() {
137        let leader0 = SlotLeader {
138            id: solana_pubkey::new_rand(),
139            vote_address: Pubkey::default(),
140        };
141        let leader1 = SlotLeader {
142            id: solana_pubkey::new_rand(),
143            vote_address: solana_pubkey::new_rand(),
144        };
145        let mut stakes = vec![(leader0, 1), (leader1, 1)];
146        sort_stakes(&mut stakes);
147        assert_eq!(stakes, vec![(leader1, 1), (leader0, 1)]);
148    }
149
150    fn pubkey_from_u16(n: u16) -> Pubkey {
151        let mut bytes = [0; 32];
152        bytes[0..2].copy_from_slice(&n.to_le_bytes());
153        Pubkey::new_from_array(bytes)
154    }
155
156    #[test_case(1, &[10, 20, 30], 12, NZ_1, &[1, 1, 2, 1, 1, 0, 0, 1, 2, 1, 0, 1])]
157    #[test_case(1, &[10, 20, 30], 12, NZ_2, &[1, 1, 2, 1, 1, 0])]
158    #[test_case(1, &[30, 10, 20], 12, NZ_1, &[2, 2, 0, 2, 2, 1, 1, 2, 0, 2, 1, 2])]
159    #[test_case(1, &[30, 10, 20], 12, NZ_2, &[2, 2, 0, 2, 2, 1])]
160    #[test_case(1, &[10, 20, 25, 30], 12, NZ_1, &[2, 2, 3, 1, 2, 0, 1, 1, 3, 2, 1, 2])]
161    #[test_case(1, &[10, 20, 25, 30, 35, 40, 100], 15, NZ_1,
162                &[4, 5, 6, 3, 4, 1, 2, 3, 6, 4, 2, 4, 5, 6, 6])]
163    #[test_case(1, &[10, 20, 25, 30, 35, 40, 100, 1000], 15, NZ_1,
164                &[7, 7, 7, 7, 7, 4, 6, 7, 7, 7, 6, 7, 7, 7, 7])]
165    #[test_case(1, &[10, 20, 25, 30, 35, 40, 100, 1000, 10_000], 20, NZ_1,
166                &[8, 8, 8, 8, 8, 7, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 7])]
167    #[test_case(1, &[10, 20, 25, 30, 35, 40, 100, 1000, 10_000], 25, NZ_1,
168                &[8, 8, 8, 8, 8, 7, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 7, 8, 8, 8, 8, 8])]
169    #[test_case(457468, &[10, 20, 30], 12, NZ_1, &[2, 2, 0, 1, 0, 2, 1, 2, 1, 2, 2, 2])]
170    #[test_case(457468, &[10, 20, 30], 12, NZ_2, &[2, 2, 0, 1, 0, 2])]
171    #[test_case(457469, &[10, 20, 30], 12, NZ_1, &[1, 2, 2, 2, 2, 2, 2, 1, 0, 2, 2, 0])]
172    #[test_case(457470, &[10, 20, 30], 12, NZ_1, &[2, 1, 1, 1, 1, 1, 1, 1, 1, 2, 0, 2])]
173    #[test_case(3466545, &[10, 20, 30], 12, NZ_1, &[2, 2, 0, 0, 2, 1, 1, 1, 0, 0, 2, 2])]
174    #[test_case(3466545, &[10, 20, 30], 13, NZ_1, &[2, 2, 0, 0, 2, 1, 1, 1, 0, 0, 2, 2, 1])]
175    #[test_case(3466545, &[10, 20, 30], 14, NZ_1, &[2, 2, 0, 0, 2, 1, 1, 1, 0, 0, 2, 2, 1, 2])]
176    #[test_case(3466545, &[10, 20, 30], 14, NZ_2, &[2, 2, 0, 0, 2, 1, 1])]
177    fn test_stake_leader_schedule_exact_order(
178        epoch: u64,
179        stakes: &[u64],
180        len: usize,
181        repeat: NonZeroUsize,
182        expected_order: &[usize],
183    ) {
184        let slot_leaders: Vec<_> = (0..stakes.len() as u16)
185            .map(|seed| SlotLeader {
186                id: pubkey_from_u16(seed),
187                vote_address: Pubkey::new_unique(),
188            })
189            .collect();
190        let stakes = slot_leaders
191            .iter()
192            .copied()
193            .zip(stakes.iter().copied())
194            .collect();
195        let order: Vec<_> = stake_weighted_slot_leaders(stakes, epoch, len, repeat)
196            .into_iter()
197            .map(|slot_leader| {
198                slot_leaders
199                    .iter()
200                    .find_position(|item| *item == &slot_leader)
201                    .unwrap()
202                    .0
203            })
204            .collect();
205        assert_eq!(order, expected_order);
206    }
207
208    #[test_case(42, 1_000, 0, "4XU6LEarBUmBkAvXRsjeyLu3N8CcgrvbRFrNiJi2jECk")]
209    #[test_case(42, 10_000, 0, "G2MGFXgdLATXWr1336i8PTcaUMc4GbJRMJdbxiarCttr")]
210    #[test_case(42, 10_000, 1, "9xLLKyyqF5YrdwPSDbqh5oVamSF7cqPqQLxEyHTexEiP")]
211    #[test_case(42, 10_000, 2, "AJ6NQi2p5SnRz9mqESqkW2PwVoT2vYy1fmKdaHxNFUAf")]
212    #[test_case(42, 10_000, 3, "2oLjZggMwDTQhzdB4KN5VQisyeRw6MZbBBdjosNZK5xR")]
213    #[test_case(346436, 1_000, 0, "59SnXMS4NzTSib8TNykiJgFQBeAVxUqsAvQm7JtkodPQ")]
214    #[test_case(346436, 1_000, 1, "BEB2nC9MBALPbgwGKGfHu6V88QG7doScx65cAd6VjnRk")]
215    #[test_case(346436, 1_000, 2, "3aLE5S6xLEU9yg5EZQH27qrC86aC2dG8KLh4NbcapXpy")]
216    #[test_case(346436, 1_000, 3, "H2bw3Y2AjxJyK7smy1ZBB4LJ7MY3i9bPQM3YdAChAww2")]
217    #[test_case(454357, 10_000, 0, "4BLanrC5t7vzNXx62javKtjCmCkd8yZfZpVrjT4eUpNQ")]
218    #[test_case(454357, 10_000, 1, "FyvbdxpVchendERMnzH2KDceqydpXtJarrfFXoLQEXgQ")]
219    #[test_case(454357, 10_000, 2, "7KwK44Y7V3GzJLN8aGZtM8EEfAYmRvaiDyKYV6jg4MQn")]
220    #[test_case(454357, 10_000, 3, "E9XL5BLhCJ4Emyfs8jTUsQetfA8QZj78LcnN63dPp7jJ")]
221    fn test_long_leader_schedule_hashed(
222        epoch: Epoch,
223        len: usize,
224        stake_pow: u32,
225        expected_hash: &str,
226    ) {
227        fn hash_slot_leader_vote_addresses(v: &[SlotLeader]) -> String {
228            use solana_sha256_hasher::Hasher;
229
230            let mut hasher = Hasher::default();
231
232            for slot_leader in v {
233                hasher.hash(&slot_leader.vote_address.to_bytes());
234            }
235
236            hasher.result().to_string()
237        }
238        let slot_leaders: Vec<_> = (0..=u16::MAX)
239            .map(|seed| SlotLeader {
240                vote_address: pubkey_from_u16(seed),
241                id: Pubkey::new_unique(),
242            })
243            .collect();
244        let stakes = slot_leaders
245            .iter()
246            .copied()
247            .enumerate()
248            .map(|(i, slot_leader)| (slot_leader, i.pow(stake_pow) as u64))
249            .collect();
250        let schedule = stake_weighted_slot_leaders(stakes, epoch, len, NZ_1);
251        assert_eq!(hash_slot_leader_vote_addresses(&schedule), expected_hash);
252    }
253
254    impl SlotLeader {
255        pub fn new_rand() -> Self {
256            SlotLeader {
257                id: solana_pubkey::new_rand(),
258                vote_address: solana_pubkey::new_rand(),
259            }
260        }
261    }
262
263    #[test]
264    #[should_panic]
265    fn test_zero_stake_panics() {
266        let _ = stake_weighted_slot_leaders(
267            vec![(SlotLeader::new_unique(), 0), (SlotLeader::new_unique(), 0)],
268            0,
269            5,
270            NZ_1,
271        );
272    }
273}
solana_leader_schedule/lib.rs

solana_leader_schedule/
lib.rs
