rt: make CONTEXT const TLS (#5757)

This makes initializing `Context` const, which lets us use const
thread-locals. The next step will be to ensure `Context` does not have a
drop impl.

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<EnterRuntime>,
 
     #[cfg(any(feature = "rt", feature = "macros"))]
-    rng: FastRand,
+    rng: Cell<Option<FastRand>>,
 
     /// 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<R>(f: impl FnOnce(&Cell<coop::Budget>) -> R) -> Result<R, AccessError> {
@@ -121,6 +126,7 @@ pub(super) fn budget<R>(f: impl FnOnce(&Cell<coop::Budget>) -> R) -> Result<R, A
 
 cfg_rt! {
     use crate::runtime::{ThreadId, TryCurrentError};
+    use crate::util::rand::RngSeed;
 
     use std::fmt;
 
@@ -295,7 +301,9 @@ cfg_rt! {
             let rng_seed = handle.seed_generator().next_seed();
 
             let old_handle = self.handle.borrow_mut().replace(handle.clone());
-            let old_seed = self.rng.replace_seed(rng_seed);
+            let mut rng = self.rng.get().unwrap_or_else(FastRand::new);
+            let old_seed = rng.replace_seed(rng_seed);
+            self.rng.set(Some(rng));
 
             SetCurrentGuard {
                 old_handle,
@@ -308,7 +316,10 @@ cfg_rt! {
         fn drop(&mut self) {
             CONTEXT.with(|ctx| {
                 *ctx.handle.borrow_mut() = self.old_handle.take();
-                ctx.rng.replace_seed(self.old_seed.clone());
+
+                let mut rng = ctx.rng.get().unwrap_or_else(FastRand::new);
+                rng.replace_seed(self.old_seed.clone());
+                ctx.rng.set(Some(rng));
             });
         }
     }

tokio/src/runtime/context/scoped.rs

@@ -9,7 +9,7 @@ pub(super) struct Scoped<T> {
 unsafe impl<T> Sync for Scoped<T> {}
 
 impl<T> Scoped<T> {
-    pub(super) fn new() -> Scoped<T> {
+    pub(super) const fn new() -> Scoped<T> {
         Scoped {
             inner: Cell::new(ptr::null()),
         }

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 });

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<FastRand>,
-    }
-
-    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: <https://www.jstatsoft.org/article/view/v008i14/xorshift.pdf>
+/// This generator passes the SmallCrush suite, part of TestU01 framework:
+/// <http://simul.iro.umontreal.ca/testu01/tu01.html>
+#[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: <https://www.jstatsoft.org/article/view/v008i14/xorshift.pdf>
-/// This generator passes the SmallCrush suite, part of TestU01 framework:
-/// <http://simul.iro.umontreal.ca/testu01/tu01.html>
-#[derive(Debug)]
-pub(crate) struct FastRand {
-    one: Cell<u32>,
-    two: Cell<u32>,
-}
 
 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)
     }

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<FastRand>,
+}
+
+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
+    }
+}

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())
+    }
+}