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

shuttle_engine/runtime/
runner.rs
