diff --git a/tokio/src/runtime/context.rs b/tokio/src/runtime/context.rs index 880b67205..c0d3be325 100644 --- a/tokio/src/runtime/context.rs +++ b/tokio/src/runtime/context.rs @@ -4,7 +4,7 @@ use crate::runtime::coop; use std::cell::Cell; #[cfg(any(feature = "rt", feature = "macros"))] -use crate::util::rand::{FastRand, RngSeed}; +use crate::util::rand::FastRand; cfg_rt! { mod scoped; @@ -47,7 +47,7 @@ struct Context { runtime: Cell, #[cfg(any(feature = "rt", feature = "macros"))] - rng: FastRand, + rng: Cell>, /// Tracks the amount of "work" a task may still do before yielding back to /// the sheduler @@ -64,7 +64,7 @@ struct Context { } tokio_thread_local! { - static CONTEXT: Context = { + static CONTEXT: Context = const { Context { #[cfg(feature = "rt")] thread_id: Cell::new(None), @@ -90,7 +90,7 @@ tokio_thread_local! { runtime: Cell::new(EnterRuntime::NotEntered), #[cfg(any(feature = "rt", feature = "macros"))] - rng: FastRand::new(RngSeed::new()), + rng: Cell::new(None), budget: Cell::new(coop::Budget::unconstrained()), @@ -112,7 +112,12 @@ tokio_thread_local! { #[cfg(any(feature = "macros", all(feature = "sync", feature = "rt")))] pub(crate) fn thread_rng_n(n: u32) -> u32 { - CONTEXT.with(|ctx| ctx.rng.fastrand_n(n)) + CONTEXT.with(|ctx| { + let mut rng = ctx.rng.get().unwrap_or_else(FastRand::new); + let ret = rng.fastrand_n(n); + ctx.rng.set(Some(rng)); + ret + }) } pub(super) fn budget(f: impl FnOnce(&Cell) -> R) -> Result { @@ -121,6 +126,7 @@ pub(super) fn budget(f: impl FnOnce(&Cell) -> R) -> Result { unsafe impl Sync for Scoped {} impl Scoped { - pub(super) fn new() -> Scoped { + pub(super) const fn new() -> Scoped { Scoped { inner: Cell::new(ptr::null()), } diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs index 76f8d0be4..947b6fb7f 100644 --- a/tokio/src/runtime/scheduler/multi_thread/worker.rs +++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs @@ -252,7 +252,7 @@ pub(super) fn create( park: Some(park), global_queue_interval: stats.tuned_global_queue_interval(&config), stats, - rand: FastRand::new(config.seed_generator.next_seed()), + rand: FastRand::from_seed(config.seed_generator.next_seed()), })); remotes.push(Remote { steal, unpark }); diff --git a/tokio/src/util/rand.rs b/tokio/src/util/rand.rs index 7018c1ee0..d96c8d37e 100644 --- a/tokio/src/util/rand.rs +++ b/tokio/src/util/rand.rs @@ -1,48 +1,9 @@ -use std::cell::Cell; - cfg_rt! { - use std::sync::Mutex; + mod rt; + pub(crate) use rt::RngSeedGenerator; - /// A deterministic generator for seeds (and other generators). - /// - /// Given the same initial seed, the generator will output the same sequence of seeds. - /// - /// Since the seed generator will be kept in a runtime handle, we need to wrap `FastRand` - /// in a Mutex to make it thread safe. Different to the `FastRand` that we keep in a - /// thread local store, the expectation is that seed generation will not need to happen - /// very frequently, so the cost of the mutex should be minimal. - #[derive(Debug)] - pub(crate) struct RngSeedGenerator { - /// Internal state for the seed generator. We keep it in a Mutex so that we can safely - /// use it across multiple threads. - state: Mutex, - } - - impl RngSeedGenerator { - /// Returns a new generator from the provided seed. - pub(crate) fn new(seed: RngSeed) -> Self { - Self { - state: Mutex::new(FastRand::new(seed)), - } - } - - /// Returns the next seed in the sequence. - pub(crate) fn next_seed(&self) -> RngSeed { - let rng = self - .state - .lock() - .expect("RNG seed generator is internally corrupt"); - - let s = rng.fastrand(); - let r = rng.fastrand(); - - RngSeed::from_pair(s, r) - } - - /// Directly creates a generator using the next seed. - pub(crate) fn next_generator(&self) -> Self { - RngSeedGenerator::new(self.next_seed()) - } + cfg_unstable! { + mod rt_unstable; } } @@ -57,31 +18,25 @@ pub struct RngSeed { r: u32, } +/// Fast random number generate. +/// +/// Implement xorshift64+: 2 32-bit xorshift sequences added together. +/// Shift triplet `[17,7,16]` was calculated as indicated in Marsaglia's +/// Xorshift paper: +/// This generator passes the SmallCrush suite, part of TestU01 framework: +/// +#[derive(Clone, Copy, Debug)] +pub(crate) struct FastRand { + one: u32, + two: u32, +} + impl RngSeed { /// Creates a random seed using loom internally. pub(crate) fn new() -> Self { Self::from_u64(crate::loom::rand::seed()) } - cfg_unstable! { - /// Generates a seed from the provided byte slice. - /// - /// # Example - /// - /// ``` - /// # use tokio::runtime::RngSeed; - /// let seed = RngSeed::from_bytes(b"make me a seed"); - /// ``` - #[cfg(feature = "rt")] - pub fn from_bytes(bytes: &[u8]) -> Self { - use std::{collections::hash_map::DefaultHasher, hash::Hasher}; - - let mut hasher = DefaultHasher::default(); - hasher.write(bytes); - Self::from_u64(hasher.finish()) - } - } - fn from_u64(seed: u64) -> Self { let one = (seed >> 32) as u32; let mut two = seed as u32; @@ -98,41 +53,19 @@ impl RngSeed { Self { s, r } } } -/// Fast random number generate. -/// -/// Implement xorshift64+: 2 32-bit xorshift sequences added together. -/// Shift triplet `[17,7,16]` was calculated as indicated in Marsaglia's -/// Xorshift paper: -/// This generator passes the SmallCrush suite, part of TestU01 framework: -/// -#[derive(Debug)] -pub(crate) struct FastRand { - one: Cell, - two: Cell, -} impl FastRand { - /// Initializes a new, thread-local, fast random number generator. - pub(crate) fn new(seed: RngSeed) -> FastRand { - FastRand { - one: Cell::new(seed.s), - two: Cell::new(seed.r), - } + /// Initialize a new fast random number generator using the default source of entropy. + pub(crate) fn new() -> FastRand { + FastRand::from_seed(RngSeed::new()) } - /// Replaces the state of the random number generator with the provided seed, returning - /// the seed that represents the previous state of the random number generator. - /// - /// The random number generator will become equivalent to one created with - /// the same seed. - #[cfg(feature = "rt")] - pub(crate) fn replace_seed(&self, seed: RngSeed) -> RngSeed { - let old_seed = RngSeed::from_pair(self.one.get(), self.two.get()); - - self.one.replace(seed.s); - self.two.replace(seed.r); - - old_seed + /// Initializes a new, thread-local, fast random number generator. + pub(crate) fn from_seed(seed: RngSeed) -> FastRand { + FastRand { + one: seed.s, + two: seed.r, + } } #[cfg(any( @@ -140,22 +73,22 @@ impl FastRand { feature = "rt-multi-thread", all(feature = "sync", feature = "rt") ))] - pub(crate) fn fastrand_n(&self, n: u32) -> u32 { + pub(crate) fn fastrand_n(&mut self, n: u32) -> u32 { // This is similar to fastrand() % n, but faster. // See https://lemire.me/blog/2016/06/27/a-fast-alternative-to-the-modulo-reduction/ let mul = (self.fastrand() as u64).wrapping_mul(n as u64); (mul >> 32) as u32 } - fn fastrand(&self) -> u32 { - let mut s1 = self.one.get(); - let s0 = self.two.get(); + fn fastrand(&mut self) -> u32 { + let mut s1 = self.one; + let s0 = self.two; s1 ^= s1 << 17; s1 = s1 ^ s0 ^ s1 >> 7 ^ s0 >> 16; - self.one.set(s0); - self.two.set(s1); + self.one = s0; + self.two = s1; s0.wrapping_add(s1) } diff --git a/tokio/src/util/rand/rt.rs b/tokio/src/util/rand/rt.rs new file mode 100644 index 000000000..6584ba4f2 --- /dev/null +++ b/tokio/src/util/rand/rt.rs @@ -0,0 +1,61 @@ +use super::{FastRand, RngSeed}; + +use std::sync::Mutex; + +/// A deterministic generator for seeds (and other generators). +/// +/// Given the same initial seed, the generator will output the same sequence of seeds. +/// +/// Since the seed generator will be kept in a runtime handle, we need to wrap `FastRand` +/// in a Mutex to make it thread safe. Different to the `FastRand` that we keep in a +/// thread local store, the expectation is that seed generation will not need to happen +/// very frequently, so the cost of the mutex should be minimal. +#[derive(Debug)] +pub(crate) struct RngSeedGenerator { + /// Internal state for the seed generator. We keep it in a Mutex so that we can safely + /// use it across multiple threads. + state: Mutex, +} + +impl RngSeedGenerator { + /// Returns a new generator from the provided seed. + pub(crate) fn new(seed: RngSeed) -> Self { + Self { + state: Mutex::new(FastRand::from_seed(seed)), + } + } + + /// Returns the next seed in the sequence. + pub(crate) fn next_seed(&self) -> RngSeed { + let mut rng = self + .state + .lock() + .expect("RNG seed generator is internally corrupt"); + + let s = rng.fastrand(); + let r = rng.fastrand(); + + RngSeed::from_pair(s, r) + } + + /// Directly creates a generator using the next seed. + pub(crate) fn next_generator(&self) -> Self { + RngSeedGenerator::new(self.next_seed()) + } +} + +impl FastRand { + /// Replaces the state of the random number generator with the provided seed, returning + /// the seed that represents the previous state of the random number generator. + /// + /// The random number generator will become equivalent to one created with + /// the same seed. + pub(crate) fn replace_seed(&mut self, seed: RngSeed) -> RngSeed { + let old_seed = RngSeed::from_pair(self.one, self.two); + + self.one = seed.s; + self.two = seed.r; + + old_seed + } +} diff --git a/tokio/src/util/rand/rt_unstable.rs b/tokio/src/util/rand/rt_unstable.rs new file mode 100644 index 000000000..3a8613eb0 --- /dev/null +++ b/tokio/src/util/rand/rt_unstable.rs @@ -0,0 +1,20 @@ +use super::RngSeed; + +use std::collections::hash_map::DefaultHasher; +use std::hash::Hasher; + +impl RngSeed { + /// Generates a seed from the provided byte slice. + /// + /// # Example + /// + /// ``` + /// # use tokio::runtime::RngSeed; + /// let seed = RngSeed::from_bytes(b"make me a seed"); + /// ``` + pub fn from_bytes(bytes: &[u8]) -> Self { + let mut hasher = DefaultHasher::default(); + hasher.write(bytes); + Self::from_u64(hasher.finish()) + } +}