Add reset_step_count

src/current.rs

@@ -95,6 +95,16 @@ pub fn me() -> TaskId {
     get_current_task().unwrap()
 }
 
+/// Sets the number of scheduling steps used (wrt. the step bound) to 0.
+///
+/// The idea behind this is to run the test with some step bound, and then call this function whenever it is known that progress has been made.
+/// This allows tests to run with tighter step bounds, and to scale a test up without also changing the step bound.
+///
+/// NOTE: Be careful when using this, as if used wrongly it can be used to make a test execute forever.
+pub fn reset_step_count() {
+    ExecutionState::with(|s| s.steps_reset_at = s.current_schedule.len());
+}
+
 /// Sets the `tag` field of the current task.
 /// Returns the `tag` which was there previously.
 #[deprecated]

src/runtime/execution.rs

@@ -246,12 +246,14 @@ pub(crate) struct ExecutionState {
     has_yielded: bool,
     // the number of scheduling decisions made so far
     context_switches: usize,
+    // the schedule length last time `reset_stop_bound()` was called
+    pub(crate) steps_reset_at: usize,
 
     // static values for the current execution
     storage: StorageMap,
 
     scheduler: Rc<RefCell<dyn Scheduler>>,
-    pub(super) current_schedule: Schedule,
+    pub(crate) current_schedule: Schedule,
 
     in_cleanup: bool,
 
@@ -292,6 +294,7 @@ impl ExecutionState {
             next_task: ScheduledTask::None,
             has_yielded: false,
             context_switches: 0,
+            steps_reset_at: 0,
             storage: StorageMap::new(),
             scheduler,
             current_schedule: initial_schedule,
@@ -631,6 +634,11 @@ impl ExecutionState {
         &mut task.clock
     }
 
+    /// Returns `true` if the test has exceeded the step bound, and `false` otherwise.
+    fn is_step_bound_exceeded(&self, max_steps: usize) -> bool {
+        self.current_schedule.len() - self.steps_reset_at >= max_steps
+    }
+
     /// Run the scheduler to choose the next task to run. `has_yielded` should be false if the
     /// scheduler is being invoked from within a running task. If scheduling fails, returns an Err
     /// with a String describing the failure.
@@ -644,14 +652,14 @@ impl ExecutionState {
         self.context_switches += 1;
 
         match self.config.max_steps {
-            MaxSteps::FailAfter(n) if self.current_schedule.len() >= n => {
+            MaxSteps::FailAfter(max_steps) if self.is_step_bound_exceeded(max_steps) => {
                 let msg = format!(
                     "exceeded max_steps bound {}. this might be caused by an unfair schedule (e.g., a spin loop)?",
-                    n
+                    max_steps
                 );
                 return Err(msg);
             }
-            MaxSteps::ContinueAfter(n) if self.current_schedule.len() >= n => {
+            MaxSteps::ContinueAfter(max_steps) if self.is_step_bound_exceeded(max_steps) => {
                 self.next_task = ScheduledTask::Stopped;
                 return Ok(());
             }

tests/basic/execution.rs

@@ -1,11 +1,14 @@
-use shuttle::scheduler::RandomScheduler;
-use shuttle::{check, check_dfs, current, thread, Config, MaxSteps, Runner};
+use shuttle::{
+    check, check_dfs, current,
+    scheduler::{DfsScheduler, RandomScheduler},
+    thread, Config, MaxSteps, Runner,
+};
 use std::panic::{catch_unwind, AssertUnwindSafe};
-use test_log::test;
 // Not actually trying to explore interleavings involving AtomicUsize, just using to smuggle a
 // mutable counter across threads
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
+use test_log::test;
 
 #[test]
 fn basic_scheduler_test() {
@@ -244,3 +247,43 @@ fn failure_outside_execution() {
     let lock = shuttle::sync::Mutex::new(0u64);
     let _ = lock.lock().unwrap();
 }
+
+fn reset_step_count(do_reset: bool, step_bound: usize) {
+    let mut config = Config::new();
+    config.max_steps = MaxSteps::FailAfter(step_bound);
+
+    let scheduler = DfsScheduler::new(None, false);
+    let runner = Runner::new(scheduler, config);
+
+    runner.run(move || {
+        (0..4)
+            .map(move |_| {
+                thread::spawn(move || {
+                    for _ in 0..3 {
+                        thread::yield_now();
+                    }
+                    if do_reset {
+                        shuttle::current::reset_step_count();
+                    }
+                })
+            })
+            .for_each(|jh| jh.join().unwrap())
+    });
+}
+
+#[test]
+#[should_panic(expected = "exceeded max_steps bound")]
+fn dont_reset_step_count() {
+    reset_step_count(false, 20);
+}
+
+#[test]
+fn do_reset_step_count() {
+    reset_step_count(true, 7);
+}
+
+#[test]
+#[should_panic(expected = "exceeded max_steps bound")]
+fn do_reset_step_count_panics() {
+    reset_step_count(true, 6);
+}