Adding TimeModels to Shuttle

Author: Dylan Wolff

shuttle/Cargo.toml

@@ -15,6 +15,7 @@ cfg-if = "1.0"
 generator = "0.8.1"
 hex = "0.4.2"
 owo-colors = "3.5.0"
+pin-project = "1.1.3"
 rand_core = "0.6.4"
 rand = "0.8.5"
 rand_pcg = "0.3.1"
@@ -38,7 +39,6 @@ tempfile = "3.2.0"
 test-log = { version = "0.2.8", default-features = false, features = ["trace"] }
 tracing-subscriber = { version = "0.3.9", features = ["env-filter"] }
 trybuild = "1.0"
-pin-project = "1.1.3"
 # The following line is necessary to ensure vector-clocks are used for integration tests
 # To run performance tests without vector clocks using `cargo bench`, this line must be commented out
 shuttle = { path = ".", features = ["vector-clocks"] }
@@ -55,6 +55,7 @@ annotation = ["dep:serde", "dep:serde_json", "dep:regex"]
 # are otherwise always enabled via a dev-dependency to ensure all *test* assertions utilizing vector
 # clocks behave correctly during testing
 bench-no-vector-clocks = []
+advanced-time-models = []
 
 [[bench]]
 name = "lock"

shuttle/src/runtime/execution.rs

@@ -5,6 +5,7 @@ use crate::runtime::task::labels::Labels;
 use crate::runtime::task::{ChildLabelFn, Task, TaskId, TaskName, TaskSignature, DEFAULT_INLINE_TASKS};
 use crate::runtime::thread::continuation::PooledContinuation;
 use crate::scheduler::{Schedule, Scheduler};
+use crate::sync::time::TimeModel;
 use crate::sync::{ResourceSignature, ResourceType};
 use crate::thread::thread_fn;
 use crate::{Config, MaxSteps};
@@ -48,15 +49,21 @@ thread_local! {
 /// static variable, but clients get access to it by calling `ExecutionState::with`.
 pub(crate) struct Execution {
     scheduler: Rc<RefCell<dyn Scheduler>>,
+    time_model: Rc<RefCell<dyn TimeModel>>,
     initial_schedule: Schedule,
 }
 
 impl Execution {
     /// Construct a new execution that will use the given scheduler. The execution should then be
     /// invoked via its `run` method, which takes as input the closure for task 0.
-    pub(crate) fn new(scheduler: Rc<RefCell<dyn Scheduler>>, initial_schedule: Schedule) -> Self {
+    pub(crate) fn new(
+        scheduler: Rc<RefCell<dyn Scheduler>>,
+        initial_schedule: Schedule,
+        time_model: Rc<RefCell<dyn TimeModel>>,
+    ) -> Self {
         Self {
             scheduler,
+            time_model,
             initial_schedule,
         }
     }
@@ -77,6 +84,7 @@ impl Execution {
         let state = RefCell::new(ExecutionState::new(
             config.clone(),
             Rc::clone(&self.scheduler),
+            Rc::clone(&self.time_model),
             self.initial_schedule.clone(),
         ));
 
@@ -95,6 +103,7 @@ impl Execution {
 
             // Cleanup the state before it goes out of `EXECUTION_STATE` scope
             ExecutionState::cleanup();
+            self.time_model.borrow_mut().reset();
         });
     }
 
@@ -107,6 +116,13 @@ impl Execution {
             Finished,
         }
 
+        // While there are no runnable tasks and tasks are able to be woken by the time model, continue waking tasks
+        while ExecutionState::num_runnable() == 0
+            && ExecutionState::with(|s| Rc::clone(&s.time_model))
+                .borrow_mut()
+                .wake_next()
+        {}
+
         let next_step = ExecutionState::with(|state| {
             if let Err(msg) = state.schedule() {
                 return NextStep::Failure(msg, state.current_schedule.clone());
@@ -281,6 +297,10 @@ pub(crate) struct ExecutionState {
     // Persistent Vec used as a bump allocator for references to runnable tasks to avoid slow allocation
     // on each scheduling decision. Should not be used outside of the `schedule` function
     runnable_tasks: Vec<*const Task>,
+
+    // Counter for unique timing resource ids (Sleeps, Timeouts and Intervals)
+    pub(crate) timer_id_counter: u64,
+    pub(crate) time_model: Rc<RefCell<dyn TimeModel>>,
 }
 
 #[derive(Debug, PartialEq, Eq, Clone, Copy)]
@@ -305,7 +325,12 @@ impl ScheduledTask {
 }
 
 impl ExecutionState {
-    fn new(config: Config, scheduler: Rc<RefCell<dyn Scheduler>>, initial_schedule: Schedule) -> Self {
+    fn new(
+        config: Config,
+        scheduler: Rc<RefCell<dyn Scheduler>>,
+        time_model: Rc<RefCell<dyn TimeModel>>,
+        initial_schedule: Schedule,
+    ) -> Self {
         Self {
             config,
             tasks: SmallVec::new(),
@@ -322,17 +347,21 @@ impl ExecutionState {
             has_cleaned_up: false,
             top_level_span: tracing::Span::current(),
             runnable_tasks: Vec::with_capacity(DEFAULT_INLINE_TASKS),
+            time_model,
+            timer_id_counter: 0,
         }
     }
 
     /// Invoke a closure with access to the current execution state. Library code uses this to gain
     /// access to the state of the execution to influence scheduling (e.g. to register a task as
     /// blocked).
     #[inline]
+    #[track_caller]
     pub(crate) fn with<F, T>(f: F) -> T
     where
         F: FnOnce(&mut ExecutionState) -> T,
     {
+        trace!("ExecutionState::with from {}", Location::caller());
         Self::try_with(f).expect("Shuttle internal error: cannot access ExecutionState. are you trying to access a Shuttle primitive from outside a Shuttle test?")
     }
 
@@ -550,6 +579,8 @@ impl ExecutionState {
         TASK_ID_TO_TAGS.with(|cell| cell.borrow_mut().clear());
         LABELS.with(|cell| cell.borrow_mut().clear());
 
+        Self::with(|s| s.timer_id_counter = 0);
+
         #[cfg(debug_assertions)]
         Self::with(|state| state.has_cleaned_up = true);
 
@@ -565,6 +596,13 @@ impl ExecutionState {
     /// is different from the currently running task, indicating that the current task should yield
     /// its execution.
     pub(crate) fn maybe_yield() -> bool {
+        // While there are no runnable tasks and tasks are able to be woken by the time model, continue waking tasks
+        while ExecutionState::num_runnable() == 0
+            && ExecutionState::with(|s| Rc::clone(&s.time_model))
+                .borrow_mut()
+                .wake_next()
+        {}
+
         Self::with(|state| {
             debug_assert!(
                 matches!(state.current_task, ScheduledTask::Some(_)) && state.next_task == ScheduledTask::None,
@@ -666,11 +704,19 @@ impl ExecutionState {
         Self::with(|state| state.context_switches)
     }
 
+    pub(crate) fn num_tasks(&self) -> usize {
+        self.tasks.len()
+    }
+
     #[track_caller]
     pub(crate) fn new_resource_signature(resource_type: ResourceType) -> ResourceSignature {
         ExecutionState::with(|s| s.current_mut().signature.new_resource(resource_type))
     }
 
+    pub(crate) fn num_runnable() -> usize {
+        Self::with(|state| state.tasks.iter().filter(|t| t.runnable()).count())
+    }
+
     pub(crate) fn get_storage<K: Into<StorageKey>, T: 'static>(&self, key: K) -> Option<&T> {
         self.storage
             .get(key.into())

shuttle/src/runtime/runner.rs

@@ -3,6 +3,7 @@ use crate::runtime::task::{Task, TaskId};
 use crate::runtime::thread::continuation::{ContinuationPool, CONTINUATION_POOL};
 use crate::scheduler::metrics::MetricsScheduler;
 use crate::scheduler::{Schedule, Scheduler};
+use crate::sync::time::{frozen::FrozenTimeModel, TimeModel};
 use crate::Config;
 use std::cell::RefCell;
 use std::fmt;
@@ -52,18 +53,33 @@ impl Drop for ResetSpanOnDrop {
 /// function as many times as dictated by the scheduler; each execution has its scheduling decisions
 /// resolved by the scheduler, which can make different choices for each execution.
 #[derive(Debug)]
-pub struct Runner<S: ?Sized + Scheduler> {
+pub struct Runner<S: ?Sized + Scheduler, T: TimeModel> {
     scheduler: Rc<RefCell<MetricsScheduler<S>>>,
+    time_model: Rc<RefCell<T>>,
     config: Config,
 }
 
-impl<S: Scheduler + 'static> Runner<S> {
+impl<S: Scheduler + 'static> Runner<S, FrozenTimeModel> {
     /// Construct a new `Runner` that will use the given `Scheduler` to control the test.
     pub fn new(scheduler: S, config: Config) -> Self {
         let metrics_scheduler = MetricsScheduler::new(scheduler);
 
         Self {
             scheduler: Rc::new(RefCell::new(metrics_scheduler)),
+            time_model: Rc::new(RefCell::new(FrozenTimeModel::new())),
+            config,
+        }
+    }
+}
+
+impl<S: Scheduler + 'static, T: TimeModel + 'static> Runner<S, T> {
+    /// Construct a new `Runner` that will use the given `Scheduler` to control the test.
+    pub fn new_with_time_model(scheduler: S, time_model: T, config: Config) -> Self {
+        let metrics_scheduler = MetricsScheduler::new(scheduler);
+
+        Self {
+            scheduler: Rc::new(RefCell::new(metrics_scheduler)),
+            time_model: Rc::new(RefCell::new(time_model)),
             config,
         }
     }
@@ -96,7 +112,7 @@ impl<S: Scheduler + 'static> Runner<S> {
                     Some(s) => s,
                 };
 
-                let execution = Execution::new(self.scheduler.clone(), schedule);
+                let execution = Execution::new(self.scheduler.clone(), schedule, self.time_model.clone());
                 let f = Arc::clone(&f);
 
                 // This is a slightly lazy way to ensure that everything outside of the "execution" span gets
@@ -122,7 +138,6 @@ pub struct PortfolioRunner {
     stop_on_first_failure: bool,
     config: Config,
 }
-
 impl PortfolioRunner {
     /// Construct a new `PortfolioRunner` with no schedulers. If `stop_on_first_failure` is true,
     /// all schedulers will be terminated as soon as any fails; if false, they will keep running

shuttle/src/runtime/thread/continuation.rs

@@ -257,10 +257,12 @@ unsafe impl Send for PooledContinuation {}
 /// Possibly yield back to the executor to perform a context switch.
 pub(crate) fn switch() {
     crate::annotations::record_tick();
+
     if ExecutionState::maybe_yield() {
         let r = generator::yield_(ContinuationOutput::Yielded).unwrap();
         assert!(matches!(r, ContinuationInput::Resume));
     }
+    ExecutionState::with(|s| Rc::clone(&s.time_model)).borrow_mut().step();
 }
 
 #[cfg(test)]

shuttle/src/sync/mod.rs

@@ -7,6 +7,7 @@ pub mod mpsc;
 mod mutex;
 mod once;
 mod rwlock;
+pub mod time;
 
 pub use barrier::{Barrier, BarrierWaitResult};
 pub use condvar::{Condvar, WaitTimeoutResult};