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::{backtrace_enabled, Config, MaxSteps};
@@ -90,15 +91,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,
         }
     }
@@ -137,7 +144,11 @@ impl Execution {
     where
         F: FnOnce() + Send + 'static,
     {
-        let state = RefCell::new(ExecutionState::new(config.clone(), Rc::clone(&self.scheduler)));
+        let state = RefCell::new(ExecutionState::new(
+            config.clone(),
+            Rc::clone(&self.scheduler),
+            Rc::clone(&self.time_model),
+        ));
 
         init_panic_hook(config.clone());
         CurrentSchedule::init(self.initial_schedule.clone());
@@ -243,6 +254,12 @@ impl Execution {
     #[inline]
     fn run_to_competion(&mut self, immediately_return_on_panic: bool) -> Result<(), StepError> {
         loop {
+            // 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: Option<Rc<RefCell<PooledContinuation>>> = ExecutionState::with(|state| {
                 state.schedule()?;
                 state.advance_to_next_task();
@@ -338,6 +355,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)]
@@ -362,7 +383,11 @@ impl ScheduledTask {
 }
 
 impl ExecutionState {
-    fn new(config: Config, scheduler: Rc<RefCell<dyn Scheduler>>) -> Self {
+    fn new(
+        config: Config,
+        scheduler: Rc<RefCell<dyn Scheduler>>,
+        time_model: Rc<RefCell<dyn TimeModel>>,
+    ) -> Self {
         Self {
             config,
             tasks: SmallVec::new(),
@@ -378,17 +403,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?")
     }
 
@@ -604,6 +633,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);
 
@@ -619,6 +650,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| {
             if std::thread::panicking() {
                 return true;
@@ -723,11 +761,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};

shuttle/src/sync/time/constant_stepped.rs

@@ -0,0 +1,146 @@
+use std::{
+    cmp::{max, Reverse},
+    collections::{BinaryHeap, HashMap},
+    task::Waker,
+};
+
+use tracing::{debug, warn};
+
+use crate::{current::TaskId, runtime::execution::ExecutionState};
+
+use super::{Duration, Instant, TimeDistribution, TimeModel};
+
+/// A time model where time advances by a constant amount for each scheduling step
+#[derive(Clone, Debug)]
+pub struct ConstantSteppedTimeModel {
+    distribution: ConstantTimeDistribution,
+    current_step_size: std::time::Duration,
+    current_time_elapsed: std::time::Duration,
+    waiters: BinaryHeap<Reverse<(std::time::Duration, TaskId, u64)>>,
+    wakers: HashMap<u64, Waker>,
+}
+
+unsafe impl Send for ConstantSteppedTimeModel {}
+
+impl ConstantSteppedTimeModel {
+    /// Create a ConstantSteppedTimeModel
+    pub fn new(step_size: std::time::Duration) -> Self {
+        let distribution = ConstantTimeDistribution::new(step_size);
+        Self {
+            distribution,
+            current_step_size: distribution.sample(),
+            current_time_elapsed: std::time::Duration::from_secs(0),
+            waiters: BinaryHeap::new(),
+            wakers: HashMap::new(),
+        }
+    }
+
+    fn unblock_expired(&mut self) {
+        while let Some(waker_key) = self.waiters.peek().and_then(|Reverse((t, _, sleep_id))| {
+            if *t <= self.current_time_elapsed {
+                Some(*sleep_id)
+            } else {
+                None
+            }
+        }) {
+            _ = self.waiters.pop();
+            if let Some(waker) = self.wakers.remove(&waker_key) {
+                waker.wake();
+            }
+        }
+    }
+
+    /// Get the currently sleeping tasks and deadlines. May contain duplicates
+    pub fn get_waiters(&self) -> &[Reverse<(std::time::Duration, TaskId, u64)>] {
+        self.waiters.as_slice()
+    }
+
+    /// Manually wake a task without affecting the global clock
+    pub fn wake_frozen(&mut self, sleep_id: u64) {
+        if let Some(waker) = self.wakers.remove(&sleep_id) {
+            waker.wake();
+        }
+    }
+}
+
+impl TimeModel for ConstantSteppedTimeModel {
+    fn pause(&mut self) {
+        warn!("Pausing stepped model has no effect")
+    }
+
+    fn resume(&mut self) {
+        warn!("Resuming stepped model has no effect")
+    }
+
+    fn step(&mut self) {
+        debug!("step");
+        self.current_time_elapsed += self.current_step_size;
+        self.unblock_expired();
+    }
+
+    fn reset(&mut self) {
+        self.current_step_size = self.distribution.sample();
+        self.current_time_elapsed = std::time::Duration::from_secs(0);
+        self.waiters.clear();
+        self.wakers.clear();
+    }
+
+    fn instant(&self) -> Instant {
+        Instant::Simulated(self.current_time_elapsed)
+    }
+
+    fn wake_next(&mut self) -> bool {
+        if self.waiters.is_empty() {
+            return false;
+        }
+        if let Some(Reverse((time, _, _))) = self.waiters.peek() {
+            self.current_time_elapsed = max(self.current_time_elapsed, *time);
+        }
+        self.unblock_expired();
+        true
+    }
+
+    #[allow(clippy::useless_conversion)]
+    fn advance(&mut self, dur: Duration) {
+        self.current_time_elapsed += dur.into();
+    }
+
+    fn register_sleep(&mut self, deadline: Instant, sleep_id: u64, waker: Option<Waker>) -> bool {
+        let deadline = deadline.unwrap_simulated();
+        if deadline <= self.current_time_elapsed {
+            return true;
+        }
+
+        if let Some(waker) = waker {
+            let task_id = ExecutionState::with(|s| s.current().id());
+            let item = (deadline, task_id, sleep_id);
+            self.waiters.push(Reverse(item));
+            self.wakers.insert(sleep_id, waker);
+        }
+        false
+    }
+
+    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
+        self
+    }
+}
+
+/// A constant distribution; each sample returns the same time
+#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
+pub struct ConstantTimeDistribution {
+    /// The time that will be returned on sampling
+    pub time: std::time::Duration,
+}
+
+impl ConstantTimeDistribution {
+    /// Create a new constant time distribution
+    pub fn new(time: std::time::Duration) -> Self {
+        Self { time }
+    }
+}
+
+impl TimeDistribution<std::time::Duration> for ConstantTimeDistribution {
+    fn sample(&self) -> std::time::Duration {
+        self.time
+    }
+}