shuttle_engine/runtime/
runner.rs

1use crate::config::Config;
2use crate::runtime::execution::Execution;
3use crate::runtime::task::{Task, TaskId};
4use crate::runtime::thread::continuation::{ContinuationPool, CONTINUATION_POOL};
5use crate::scheduler::metrics::MetricsScheduler;
6use crate::scheduler::{Schedule, Scheduler};
7use std::cell::RefCell;
8use std::fmt;
9use std::panic::{self, Location};
10use std::rc::Rc;
11use std::sync::atomic::{AtomicBool, Ordering};
12use std::sync::mpsc;
13use std::sync::Arc;
14use std::thread;
15use std::time::Instant;
16use tracing::{span, Level};
17
18// A helper struct which on `drop` exits all current spans, then enters the span which was entered when it was constructed.
19// The reason this exists is to solve the "span-stacking" issue which occurs when there is a panic inside `run` which is
20// then caught by `panic::catch_unwind()` (such as when Shuttle is run inside proptest).
21// In other words: it enables correct spans when doing proptest minimization.
22#[must_use]
23struct ResetSpanOnDrop {
24    span: tracing::Span,
25}
26
27impl ResetSpanOnDrop {
28    fn new() -> Self {
29        Self {
30            span: tracing::Span::current().clone(),
31        }
32    }
33}
34
35impl Drop for ResetSpanOnDrop {
36    // Exits all current spans, then enters the span which was entered when `self` was constructed.
37    fn drop(&mut self) {
38        tracing::dispatcher::get_default(|subscriber| {
39            while let Some(span_id) = tracing::Span::current().id().as_ref() {
40                subscriber.exit(span_id);
41            }
42            if let Some(span_id) = self.span.id().as_ref() {
43                subscriber.enter(span_id);
44            }
45        });
46    }
47}
48
49/// A `Runner` is the entry-point for testing concurrent code.
50///
51/// It takes as input a function to test and a `Scheduler` to run it under. It then executes that
52/// function as many times as dictated by the scheduler; each execution has its scheduling decisions
53/// resolved by the scheduler, which can make different choices for each execution.
54#[derive(Debug)]
55pub struct Runner<S: ?Sized + Scheduler> {
56    scheduler: Rc<RefCell<MetricsScheduler<S>>>,
57    config: Config,
58}
59
60impl<S: Scheduler + 'static> Runner<S> {
61    /// Construct a new `Runner` that will use the given `Scheduler` to control the test.
62    pub fn new(scheduler: S, config: Config) -> Self {
63        let metrics_scheduler = MetricsScheduler::new(scheduler);
64
65        Self {
66            scheduler: Rc::new(RefCell::new(metrics_scheduler)),
67            config,
68        }
69    }
70
71    /// Test the given function and return the number of times the function was invoked during the
72    /// test (i.e., the number of iterations run).
73    #[track_caller]
74    pub fn run<F>(self, f: F) -> usize
75    where
76        F: Fn() + Send + Sync + 'static,
77    {
78        let _span_drop_guard = ResetSpanOnDrop::new();
79        // Share continuations across executions to avoid reallocating them
80        // TODO it would be a lot nicer if this were a more generic "context" thing that we passed
81        // TODO around explicitly rather than being a thread local
82        CONTINUATION_POOL.set(&ContinuationPool::new(), || {
83            let f = Arc::new(f);
84
85            let start = Instant::now();
86
87            let mut i = 0;
88
89            loop {
90                if self.config.max_time.map(|t| start.elapsed() > t).unwrap_or(false) {
91                    break;
92                }
93
94                let schedule = match self.scheduler.borrow_mut().new_execution() {
95                    None => break,
96                    Some(s) => s,
97                };
98
99                let execution = Execution::new(self.scheduler.clone(), schedule);
100                let f = Arc::clone(&f);
101
102                // This is a slightly lazy way to ensure that everything outside of the "execution" span gets
103                // established correctly between executions. Fully `exit`ing and fully `enter`ing (explicitly
104                // `enter`/`exit` all `Span`s) would most likely obviate the need for this.
105                let _span_drop_guard2 = ResetSpanOnDrop::new();
106
107                span!(Level::ERROR, "execution", i)
108                    .in_scope(|| execution.run(&self.config, move || f(), Location::caller()));
109
110                i += 1;
111            }
112            i
113        })
114    }
115}
116
117/// A `PortfolioRunner` is the same as a `Runner`, except that it can run multiple different
118/// schedulers (a "portfolio" of schedulers) in parallel. If any of the schedulers finds a failing
119/// execution of the test, the entire run fails.
120pub struct PortfolioRunner {
121    schedulers: Vec<Box<dyn Scheduler + Send + 'static>>,
122    stop_on_first_failure: bool,
123    config: Config,
124}
125
126impl PortfolioRunner {
127    /// Construct a new `PortfolioRunner` with no schedulers. If `stop_on_first_failure` is true,
128    /// all schedulers will be terminated as soon as any fails; if false, they will keep running
129    /// and potentially find multiple bugs.
130    pub fn new(stop_on_first_failure: bool, config: Config) -> Self {
131        Self {
132            schedulers: Vec::new(),
133            stop_on_first_failure,
134            config,
135        }
136    }
137
138    /// Add the given scheduler to the portfolio of schedulers to run the test with.
139    pub fn add(&mut self, scheduler: impl Scheduler + Send + 'static) {
140        self.schedulers.push(Box::new(scheduler));
141    }
142
143    /// Test the given function against all schedulers in parallel. If any of the schedulers finds
144    /// a failing execution, this function panics.
145    pub fn run<F>(self, f: F)
146    where
147        F: Fn() + Send + Sync + 'static,
148    {
149        #[derive(Debug, Clone, Copy, PartialEq, Eq)]
150        enum ThreadResult {
151            Passed,
152            Failed,
153        }
154
155        let (tx, rx) = mpsc::sync_channel::<ThreadResult>(0);
156        let stop_signal = Arc::new(AtomicBool::new(false));
157        let config = self.config;
158        let f = Arc::new(f);
159
160        let threads = self
161            .schedulers
162            .into_iter()
163            .enumerate()
164            .map(|(i, scheduler)| {
165                let f = Arc::clone(&f);
166                let tx = tx.clone();
167                let stop_signal = stop_signal.clone();
168                let config = config.clone();
169
170                thread::spawn(move || {
171                    let scheduler = PortfolioStoppableScheduler { scheduler, stop_signal };
172
173                    let runner = Runner::new(scheduler, config);
174
175                    span!(Level::ERROR, "job", i).in_scope(|| {
176                        let ret = panic::catch_unwind(panic::AssertUnwindSafe(|| runner.run(move || f())));
177
178                        match ret {
179                            Ok(_) => tx.send(ThreadResult::Passed),
180                            Err(e) => {
181                                tx.send(ThreadResult::Failed).unwrap();
182                                panic::resume_unwind(e);
183                            }
184                        }
185                    })
186                })
187            })
188            .collect::<Vec<_>>();
189
190        // Wait for each thread to pass or fail, and if any fails, tell all threads to stop early
191        for _ in 0..threads.len() {
192            if rx.recv().unwrap() == ThreadResult::Failed && self.stop_on_first_failure {
193                stop_signal.store(true, Ordering::SeqCst);
194            }
195        }
196
197        // Join all threads and propagate the first panic we see (note that this might not be the
198        // same panic that caused us to stop, if multiple threads panic around the same time, but
199        // that's probably OK).
200        let mut panic = None;
201        for thread in threads {
202            if let Err(e) = thread.join() {
203                panic = Some(e);
204            }
205        }
206        assert!(stop_signal.load(Ordering::SeqCst) == panic.is_some());
207        if let Some(e) = panic {
208            std::panic::resume_unwind(e);
209        }
210    }
211}
212
213impl fmt::Debug for PortfolioRunner {
214    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
215        f.debug_struct("PortfolioRunner")
216            .field("schedulers", &self.schedulers.len())
217            .field("stop_on_first_failure", &self.stop_on_first_failure)
218            .field("config", &self.config)
219            .finish()
220    }
221}
222
223/// A wrapper around a `Scheduler` that can be told to stop early by setting a flag. We use this to
224/// abort all jobs in a `PortfolioRunner` as soon as any job fails.
225#[derive(Debug)]
226struct PortfolioStoppableScheduler<S> {
227    scheduler: S,
228    stop_signal: Arc<AtomicBool>,
229}
230
231impl<S: Scheduler> Scheduler for PortfolioStoppableScheduler<S> {
232    fn new_execution(&mut self) -> Option<Schedule> {
233        if self.stop_signal.load(Ordering::SeqCst) {
234            None
235        } else {
236            self.scheduler.new_execution()
237        }
238    }
239
240    fn next_task(
241        &mut self,
242        runnable_tasks: &[&Task],
243        current_task: Option<TaskId>,
244        is_yielding: bool,
245    ) -> Option<TaskId> {
246        if self.stop_signal.load(Ordering::SeqCst) {
247            None
248        } else {
249            self.scheduler.next_task(runnable_tasks, current_task, is_yielding)
250        }
251    }
252
253    fn next_u64(&mut self) -> u64 {
254        self.scheduler.next_u64()
255    }
256}
shuttle_engine/runtime/runner.rs

shuttle_engine/runtime/
runner.rs
