rt: add rng_seed option to runtime::Builder

tokio/src/runtime/blocking/pool.rs

@@ -8,6 +8,7 @@ use crate::runtime::builder::ThreadNameFn;
 use crate::runtime::context;
 use crate::runtime::task::{self, JoinHandle};
 use crate::runtime::{Builder, Callback, Handle};
+use crate::util::{replace_thread_rng, RngSeedGenerator};
 
 use std::collections::{HashMap, VecDeque};
 use std::fmt;
@@ -48,6 +49,9 @@ struct Inner {
 
     // Customizable wait timeout.
     keep_alive: Duration,
+
+    // Random number seed
+    seed_generator: RngSeedGenerator,
 }
 
 struct Shared {
@@ -182,6 +186,7 @@ impl BlockingPool {
                     before_stop: builder.before_stop.clone(),
                     thread_cap,
                     keep_alive,
+                    seed_generator: builder.seed_generator.next_generator(),
                 }),
             },
             shutdown_rx,
@@ -431,6 +436,8 @@ impl Inner {
         if let Some(f) = &self.after_start {
             f()
         }
+        // We own this thread so thee is no need to replace the RngSeed once we're done.
+        let _ = replace_thread_rng(self.seed_generator.next_seed());
 
         let mut shared = self.shared.lock();
         let mut join_on_thread = None;

tokio/src/runtime/builder.rs

@@ -1,5 +1,6 @@
 use crate::runtime::handle::Handle;
 use crate::runtime::{blocking, driver, Callback, Runtime, Spawner};
+use crate::util::{RngSeed, RngSeedGenerator};
 
 use std::fmt;
 use std::io;
@@ -90,6 +91,9 @@ pub struct Builder {
     /// This option should only be exposed as unstable.
     pub(super) disable_lifo_slot: bool,
 
+    /// Specify a random number generator seed to provide deterministic results
+    pub(super) seed_generator: RngSeedGenerator,
+
     #[cfg(tokio_unstable)]
     pub(super) unhandled_panic: UnhandledPanic,
 }
@@ -255,6 +259,8 @@ impl Builder {
             global_queue_interval,
             event_interval,
 
+            seed_generator: RngSeedGenerator::new(RngSeed::new()),
+
             #[cfg(tokio_unstable)]
             unhandled_panic: UnhandledPanic::Ignore,
 
@@ -828,6 +834,36 @@ impl Builder {
         pub fn disable_lifo_slot(&mut self) -> &mut Self {
             self.disable_lifo_slot = true;
             self
+        }
+
+                /// Specifies the random number generation seed to use within all threads associated
+        /// with the runtime being built.
+        ///
+        /// This option is intended to make certain parts of the runtime deterministic.
+        /// Specifically, it affects the [`tokio::select!`] macro and the work stealing
+        /// algorithm. In the case of [`tokio::select!`] it will ensure that the order that
+        /// branches are polled is deterministic.
+        ///
+        /// In the case of work stealing, it's a little more complicated. Each worker will
+        /// be given a deterministic seed so that the starting peer for each work stealing
+        /// search will be deterministic.
+        ///
+        /// # Examples
+        ///
+        /// ```
+        /// # use tokio::runtime::{self, RngSeed};
+        /// # pub fn main() {
+        /// let seed = RngSeed::from_bytes(b"place your seed here");
+        /// let rt = runtime::Builder::new_current_thread()
+        ///     .rng_seed(seed)
+        ///     .build();
+        /// # }
+        /// ```
+        ///
+        /// [`tokio::select!`]: crate::select
+        pub fn rng_seed(&mut self, seed: RngSeed) -> &mut Self {
+            self.seed_generator = RngSeedGenerator::new(seed);
+            self
         }
     }
 
@@ -855,6 +891,7 @@ impl Builder {
                 #[cfg(tokio_unstable)]
                 unhandled_panic: self.unhandled_panic.clone(),
                 disable_lifo_slot: self.disable_lifo_slot,
+                seed_generator: self.seed_generator.next_generator(),
             },
         );
         let spawner = Spawner::CurrentThread(scheduler.spawner().clone());
@@ -866,6 +903,7 @@ impl Builder {
             signal_handle: resources.signal_handle,
             clock: resources.clock,
             blocking_spawner,
+            seed_generator: self.seed_generator.next_generator(),
         });
 
         Ok(Runtime {
@@ -975,6 +1013,7 @@ cfg_rt_multi_thread! {
                     #[cfg(tokio_unstable)]
                     unhandled_panic: self.unhandled_panic.clone(),
                     disable_lifo_slot: self.disable_lifo_slot,
+                    seed_generator: self.seed_generator.next_generator(),
                 },
             );
             let spawner = Spawner::MultiThread(scheduler.spawner().clone());
@@ -986,6 +1025,7 @@ cfg_rt_multi_thread! {
                 signal_handle: resources.signal_handle,
                 clock: resources.clock,
                 blocking_spawner,
+                seed_generator: self.seed_generator.next_generator(),
             });
 
             // Create the runtime handle

tokio/src/runtime/config.rs

@@ -1,5 +1,6 @@
 #![cfg_attr(any(not(feature = "full"), tokio_wasm), allow(dead_code))]
 use crate::runtime::Callback;
+use crate::util::RngSeedGenerator;
 
 pub(crate) struct Config {
     /// How many ticks before pulling a task from the global/remote queue?
@@ -23,6 +24,10 @@ pub(crate) struct Config {
     /// stop-gap, this unstable option lets users disable the LIFO task.
     pub(crate) disable_lifo_slot: bool,
 
+    /// Random number generator seed to configure runtimes to act in a
+    /// deterministic way.
+    pub(crate) seed_generator: RngSeedGenerator,
+
     #[cfg(tokio_unstable)]
     /// How to respond to unhandled task panics.
     pub(crate) unhandled_panic: crate::runtime::UnhandledPanic,

tokio/src/runtime/context.rs

@@ -1,5 +1,8 @@
 //! Thread local runtime context
-use crate::runtime::{Handle, TryCurrentError};
+use crate::{
+    runtime::{Handle, TryCurrentError},
+    util::{replace_thread_rng, RngSeed},
+};
 
 use std::cell::RefCell;
 
@@ -83,21 +86,29 @@ pub(crate) fn enter(new: Handle) -> EnterGuard {
 ///
 /// [`Handle`]: Handle
 pub(crate) fn try_enter(new: Handle) -> Option<EnterGuard> {
-    CONTEXT
-        .try_with(|ctx| {
-            let old = ctx.borrow_mut().replace(new);
-            EnterGuard(old)
-        })
-        .ok()
+    let rng_seed = new.as_inner().seed_generator.next_seed();
+    let old_handle = CONTEXT.try_with(|ctx| ctx.borrow_mut().replace(new)).ok()?;
+
+    let old_seed = replace_thread_rng(rng_seed);
+
+    Some(EnterGuard {
+        old_handle,
+        old_seed,
+    })
 }
 
 #[derive(Debug)]
-pub(crate) struct EnterGuard(Option<Handle>);
+pub(crate) struct EnterGuard {
+    old_handle: Option<Handle>,
+    old_seed: RngSeed,
+}
 
 impl Drop for EnterGuard {
     fn drop(&mut self) {
         CONTEXT.with(|ctx| {
-            *ctx.borrow_mut() = self.0.take();
+            *ctx.borrow_mut() = self.old_handle.take();
         });
+        // We discard the RngSeed associated with this guard
+        let _ = replace_thread_rng(self.old_seed.clone());
     }
 }

tokio/src/runtime/handle.rs

@@ -2,6 +2,7 @@
 #![cfg_attr(not(feature = "rt"), allow(dead_code))]
 
 use crate::runtime::driver;
+use crate::util::RngSeedGenerator;
 
 use std::sync::Arc;
 
@@ -58,6 +59,9 @@ pub(crate) struct HandleInner {
     /// Blocking pool spawner
     #[cfg(feature = "rt")]
     pub(crate) blocking_spawner: blocking::Spawner,
+
+    /// Current random number generator seed (when no)
+    pub(super) seed_generator: RngSeedGenerator,
 }
 
 cfg_rt! {

tokio/src/runtime/mod.rs

@@ -216,6 +216,7 @@ cfg_rt! {
     pub use self::builder::Builder;
     cfg_unstable! {
         pub use self::builder::UnhandledPanic;
+        pub use crate::util::RngSeed;
     }
 
     pub(crate) mod context;

tokio/src/runtime/scheduler/multi_thread/worker.rs

@@ -58,7 +58,6 @@
 
 use crate::coop;
 use crate::future::Future;
-use crate::loom::rand::seed;
 use crate::loom::sync::{Arc, Mutex};
 use crate::runtime;
 use crate::runtime::enter::EnterContext;
@@ -206,7 +205,7 @@ pub(super) fn create(size: usize, park: Parker, config: Config) -> (Arc<Shared>,
             is_shutdown: false,
             park: Some(park),
             metrics: MetricsBatch::new(),
-            rand: FastRand::new(seed()),
+            rand: FastRand::new(config.seed_generator.next_seed()),
         }));
 
         remotes.push(Remote { steal, unpark });

tokio/src/util/mod.rs

@@ -40,7 +40,7 @@ pub(crate) use wake_list::WakeList;
 ))]
 pub(crate) mod linked_list;
 
-#[cfg(any(feature = "rt-multi-thread", feature = "macros"))]
+#[cfg(any(feature = "rt", feature = "macros", feature = "stream"))]
 mod rand;
 
 cfg_rt! {
@@ -61,6 +61,14 @@ cfg_rt! {
     pub(crate) use rc_cell::RcCell;
 }
 
+#[allow(unreachable_pub)]
+pub use self::rand::RngSeed;
+pub(crate) use self::rand::{replace_thread_rng, RngSeedGenerator};
+
+#[cfg(any(feature = "macros"))]
+#[cfg_attr(not(feature = "macros"), allow(unreachable_pub))]
+pub use self::rand::thread_rng_n;
+
 cfg_rt_multi_thread! {
     pub(crate) use self::rand::FastRand;
 
@@ -70,8 +78,4 @@ cfg_rt_multi_thread! {
 
 pub(crate) mod trace;
 
-#[cfg(any(feature = "macros"))]
-#[cfg_attr(not(feature = "macros"), allow(unreachable_pub))]
-pub use self::rand::thread_rng_n;
-
 pub(crate) mod error;