Seed tokio with the world rng (#245)

* Introduce `rt::Config` type

In preparation for making tokio's rng seed configurable, introduce a new
`rt::Config` struct that holds configuration for creating tokio
runtimes and pass that around instead of the bare `enable_io` bool.

* Seed tokio runtimes using the world rng

If the `tokio_unstable` flag is set, seed host and client runtimes using
the world rng. This makes the simulation more deterministic. In
particular, the order in which tokio polls `select!` branches is
determined by this seed.

Author: teskje

src/rt.rs

@@ -3,6 +3,7 @@ use std::mem;
 use std::pin::Pin;
 use std::sync::Arc;
 
+use rand::rngs::SmallRng;
 use tokio::runtime::Runtime;
 use tokio::task::JoinHandle;
 use tokio::task::LocalSet;
@@ -28,6 +29,18 @@ pub enum Kind<'a> {
     NoSoftware,
 }
 
+/// Runtime configuration.
+#[derive(Clone, Default)]
+pub struct Config {
+    /// Whether io is enabled on the runtime.
+    pub enable_io: bool,
+
+    /// Rng used to seed the tokio runtime.
+    ///
+    /// Only used if the `tokio_unstable` cfg flag is set.
+    pub rng: Option<SmallRng>,
+}
+
 /// Per host simulated runtime.
 ///
 /// The tokio runtime is paused (see [`Builder::start_paused`]), which gives us
@@ -50,16 +63,16 @@ pub struct Rt<'a> {
     /// consumed to check for error, which is propagated up to fail the simulation.
     handle: Option<JoinHandle<Result>>,
 
-    /// Whether io is enabled on this runtime.
-    enable_io: bool,
+    /// Runtime configuration.
+    config: Config,
 }
 
 impl<'a> Rt<'a> {
-    pub(crate) fn client<F>(nodename: Arc<str>, client: F, enable_io: bool) -> Self
+    pub(crate) fn client<F>(nodename: Arc<str>, client: F, mut config: Config) -> Self
     where
         F: Future<Output = Result> + 'static,
     {
-        let (tokio, local) = init(enable_io);
+        let (tokio, local) = init(&mut config);
 
         let handle = with(&tokio, &local, || tokio::task::spawn_local(client));
 
@@ -69,16 +82,16 @@ impl<'a> Rt<'a> {
             local,
             nodename,
             handle: Some(handle),
-            enable_io,
+            config,
         }
     }
 
-    pub(crate) fn host<F, Fut>(nodename: Arc<str>, software: F, enable_io: bool) -> Self
+    pub(crate) fn host<F, Fut>(nodename: Arc<str>, software: F, mut config: Config) -> Self
     where
         F: Fn() -> Fut + 'a,
         Fut: Future<Output = Result> + 'static,
     {
-        let (tokio, local) = init(enable_io);
+        let (tokio, local) = init(&mut config);
 
         let software: Software = Box::new(move || Box::pin(software()));
         let handle = with(&tokio, &local, || tokio::task::spawn_local(software()));
@@ -89,20 +102,21 @@ impl<'a> Rt<'a> {
             local,
             nodename,
             handle: Some(handle),
-            enable_io,
+            config,
         }
     }
 
     pub(crate) fn no_software() -> Self {
-        let (tokio, local) = init(false);
+        let mut config = Config::default();
+        let (tokio, local) = init(&mut config);
 
         Self {
             kind: Kind::NoSoftware,
             tokio,
             local,
             nodename: String::new().into(),
             handle: None,
-            enable_io: false,
+            config,
         }
     }
 
@@ -207,23 +221,30 @@ impl<'a> Rt<'a> {
     ///
     /// Both the [`Runtime`] and [`LocalSet`] are replaced with new instances.
     fn cancel_tasks(&mut self) {
-        let (tokio, local) = init(self.enable_io);
+        let (tokio, local) = init(&mut self.config);
 
         _ = mem::replace(&mut self.tokio, tokio);
         drop(mem::replace(&mut self.local, local));
     }
 }
 
-fn init(enable_io: bool) -> (Runtime, LocalSet) {
+fn init(config: &mut Config) -> (Runtime, LocalSet) {
     let mut tokio_builder = tokio::runtime::Builder::new_current_thread();
 
     #[cfg(tokio_unstable)]
     tokio_builder.unhandled_panic(tokio::runtime::UnhandledPanic::ShutdownRuntime);
 
-    if enable_io {
+    if config.enable_io {
         tokio_builder.enable_io();
     }
 
+    #[cfg(tokio_unstable)]
+    if let Some(rng) = &mut config.rng {
+        let bytes: [u8; 32] = rand::Rng::gen(rng);
+        let seed = tokio::runtime::RngSeed::from_bytes(&bytes);
+        tokio_builder.rng_seed(seed);
+    }
+
     let tokio = tokio_builder
         .enable_time()
         .start_paused(true)

src/sim.rs

@@ -1,4 +1,4 @@
-use rand::seq::SliceRandom;
+use rand::{rngs::SmallRng, seq::SliceRandom, Rng, SeedableRng};
 use std::cell::RefCell;
 use std::future::Future;
 use std::net::IpAddr;
@@ -91,9 +91,14 @@ impl<'a> Sim<'a> {
             world.register(addr, &nodename, HostTimer::new(self.elapsed), &self.config);
         }
 
-        let rt = World::enter(&self.world, || {
-            Rt::client(nodename, client, self.config.enable_tokio_io)
-        });
+        let seed = self.world.borrow_mut().rng.gen();
+        let rng = SmallRng::from_seed(seed);
+        let config = rt::Config {
+            enable_io: self.config.enable_tokio_io,
+            rng: Some(rng),
+        };
+
+        let rt = World::enter(&self.world, || Rt::client(nodename, client, config));
 
         self.rts.insert(addr, rt);
     }
@@ -126,9 +131,14 @@ impl<'a> Sim<'a> {
             world.register(addr, &nodename, HostTimer::new(self.elapsed), &self.config);
         }
 
-        let rt = World::enter(&self.world, || {
-            Rt::host(nodename, host, self.config.enable_tokio_io)
-        });
+        let seed = self.world.borrow_mut().rng.gen();
+        let rng = SmallRng::from_seed(seed);
+        let config = rt::Config {
+            enable_io: self.config.enable_tokio_io,
+            rng: Some(rng),
+        };
+
+        let rt = World::enter(&self.world, || Rt::host(nodename, host, config));
 
         self.rts.insert(addr, rt);
     }