nuclei/rt/
rng.rs

1//! The per-thread random numbers behind the runtime's choices: where a task
2//! spawned from outside the workers goes, and which branch of a fair
3//! `select!` is polled first.
4//!
5//! Each thread has its own xorshift generator, seeded from its address and
6//! the clock; nothing depends on its quality but balance. A runtime built
7//! with an [`RngSeed`] seeds the generator of each of its workers, and of
8//! each thread while it is in the runtime's `block_on`, from that seed
9//! instead, so each such thread draws the same numbers run after run.
10
11use core::cell::Cell;
12
13thread_local! {
14    static STATE: Cell<u64> = Cell::new(fresh());
15}
16
17/// A state distinct per thread and per call: the address of a local and
18/// the clock.
19fn fresh() -> u64 {
20    let local = 0u8;
21    let addr = &local as *const u8 as u64;
22    let time = std::time::SystemTime::now()
23        .duration_since(std::time::UNIX_EPOCH)
24        .map_or(0, |d| d.as_nanos() as u64);
25    (addr ^ time.rotate_left(32)) | 1
26}
27
28/// The splitmix64 finalizer: spreads every bit of `z` over the word.
29fn mix(mut z: u64) -> u64 {
30    z = z.wrapping_add(0x9E37_79B9_7F4A_7C15);
31    z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
32    z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
33    z ^ (z >> 31)
34}
35
36/// The next number of this thread's generator.
37fn next() -> u64 {
38    STATE
39        .try_with(|s| {
40            let mut x = s.get();
41            x ^= x << 13;
42            x ^= x >> 7;
43            x ^= x << 17;
44            s.set(x);
45            x
46        })
47        // A thread on its way out has no generator left: any number serves.
48        .unwrap_or_else(|_| fresh())
49}
50
51/// A random number below `n`.
52pub(crate) fn random_below(n: usize) -> usize {
53    // Multiply-shift: unbiased enough for placement, no division.
54    ((u128::from(next()) * n as u128) >> 64) as usize
55}
56
57/// A seed for a runtime's random choices: where a task spawned from outside
58/// its workers goes, and which branch of a fair
59/// [`select!`](crate::select) is polled first.
60///
61/// A runtime built with [`Builder::rng_seed`](crate::runtime::Builder::rng_seed)
62/// seeds each of its workers, and each thread in its `block_on`, from the
63/// seed: run after run, such a thread draws the same numbers, so its
64/// choices repeat as long as what it chooses among is the same. The threads
65/// still interleave as the system runs them.
66///
67/// ```
68/// use nuclei::runtime::{RngSeed, Runtime};
69///
70/// let seed = RngSeed::from_bytes(b"a run to repeat");
71/// assert_eq!(seed, RngSeed::from_bytes(b"a run to repeat"));
72/// let rt = Runtime::builder().rng_seed(seed).build().unwrap();
73/// # drop(rt);
74/// ```
75#[derive(Clone, Copy, Debug, PartialEq, Eq)]
76pub struct RngSeed(u64);
77
78impl RngSeed {
79    /// A seed different each time.
80    pub fn new() -> Self {
81        Self(mix(next()))
82    }
83
84    /// The seed of `bytes`: the same bytes give the same seed on every
85    /// system.
86    pub fn from_bytes(bytes: &[u8]) -> Self {
87        // FNV-1a: fixed, so a seed never depends on the toolchain.
88        let hash = bytes.iter().fold(0xCBF2_9CE4_8422_2325u64, |h, &b| {
89            (h ^ u64::from(b)).wrapping_mul(0x0000_0100_0000_01B3)
90        });
91        Self(mix(hash))
92    }
93}
94
95impl Default for RngSeed {
96    fn default() -> Self {
97        Self::new()
98    }
99}
100
101/// This thread's generator set from a seed until the guard drops, then
102/// back to where it was.
103pub(crate) struct Reseeded(Option<u64>);
104
105impl Reseeded {
106    /// Seeds this thread's generator from `seed`, for the runtime's thread
107    /// `thread` (0 for a thread in its `block_on`, a worker's index plus one).
108    pub(crate) fn new(seed: RngSeed, thread: u64) -> Self {
109        let state = mix(seed.0 ^ mix(thread)) | 1;
110        Self(STATE.try_with(|s| s.replace(state)).ok())
111    }
112}
113
114impl Drop for Reseeded {
115    fn drop(&mut self) {
116        if let Some(previous) = self.0 {
117            let _ = STATE.try_with(|s| s.set(previous));
118        }
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125
126    #[test]
127    fn a_seed_draws_the_same_numbers_each_time_and_is_undone() {
128        let before: Vec<u64> = {
129            let _seeded = Reseeded::new(RngSeed::from_bytes(b"x"), 3);
130            (0..8).map(|_| next()).collect()
131        };
132        let again: Vec<u64> = {
133            let _seeded = Reseeded::new(RngSeed::from_bytes(b"x"), 3);
134            (0..8).map(|_| next()).collect()
135        };
136        assert_eq!(before, again);
137        let other_thread: Vec<u64> = {
138            let _seeded = Reseeded::new(RngSeed::from_bytes(b"x"), 4);
139            (0..8).map(|_| next()).collect()
140        };
141        assert_ne!(before, other_thread);
142        let other_seed: Vec<u64> = {
143            let _seeded = Reseeded::new(RngSeed::from_bytes(b"y"), 3);
144            (0..8).map(|_| next()).collect()
145        };
146        assert_ne!(before, other_seed);
147        // Undone: the thread's own generator goes on, not the seed's.
148        let after: Vec<u64> = (0..8).map(|_| next()).collect();
149        assert_ne!(before, after);
150    }
151
152    #[test]
153    fn numbers_stay_below_the_bound() {
154        for n in [1, 2, 3, 7, 64, 1000] {
155            assert!((0..1000).all(|_| random_below(n) < n));
156        }
157        assert_ne!(RngSeed::new(), RngSeed::new());
158    }
159}
nuclei/rt/rng.rs

nuclei/rt/
rng.rs
