Skip to main content

shuttle_engine/scheduler/
mod.rs

1//! Core scheduling types: the `Scheduler` trait, `Schedule`, and supporting data structures.
2use std::fmt::Debug;
3
4pub mod data;
5pub mod metrics;
6pub mod serialization;
7
8pub use crate::runtime::task::{Task, TaskId};
9pub use data::{DataSource, RandomDataSource};
10
11/// A `Schedule` determines the order in which tasks are to be executed
12// TODO would be nice to make this generic in the type of `seed`, but for now all our seeds are u64s
13#[derive(Clone, Debug, Default, PartialEq, Eq)]
14pub struct Schedule {
15    /// The random seed for this schedule
16    pub seed: u64,
17    /// The steps of this schedule
18    pub steps: Vec<ScheduleStep>,
19}
20
21#[derive(Clone, Debug, PartialEq, Eq)]
22pub enum ScheduleStep {
23    /// A step that schedules a specific task
24    Task(TaskId),
25    /// A step that produces a random value
26    Random,
27}
28
29impl Schedule {
30    /// Create a new empty `Schedule` that starts with the given random seed.
31    pub fn new(seed: u64) -> Self {
32        Self { seed, steps: vec![] }
33    }
34
35    /// Create a new `Schedule` that begins by scheduling the given tasks.
36    pub fn new_from_task_ids<T>(seed: u64, task_ids: impl IntoIterator<Item = T>) -> Self
37    where
38        T: Into<TaskId>,
39    {
40        let steps = task_ids
41            .into_iter()
42            .map(|t| ScheduleStep::Task(t.into()))
43            .collect::<Vec<_>>();
44        Self { seed, steps }
45    }
46
47    /// Add the given task ID as the next step of the schedule.
48    pub fn push_task(&mut self, task: TaskId) {
49        self.steps.push(ScheduleStep::Task(task));
50    }
51
52    /// Add a choice of a random u64 value as the next step of the schedule.
53    pub fn push_random(&mut self) {
54        self.steps.push(ScheduleStep::Random);
55    }
56
57    /// Return the number of steps in the schedule.
58    pub fn len(&self) -> usize {
59        self.steps.len()
60    }
61
62    /// Return true if the schedule is empty.
63    pub fn is_empty(&self) -> bool {
64        self.steps.is_empty()
65    }
66}
67
68/// A `Scheduler` is an oracle that decides the order in which to execute concurrent tasks and the
69/// data to return to calls for random values.
70///
71/// The `Scheduler` lives across multiple executions of the test case, allowing it to retain some
72/// state and strategically explore different schedules. At the start of each test execution, the
73/// executor calls `new_execution()` to inform the scheduler that a new execution is starting. Then,
74/// for each scheduling decision, the executor calls `next_task` to determine which task to run.
75pub trait Scheduler {
76    /// Inform the `Scheduler` that a new execution is about to begin. If this function returns
77    /// None, the test will end rather than performing another execution. If it returns
78    /// `Some(schedule)`, the returned `Schedule` can be used to initialize a `ReplayScheduler` for
79    /// deterministic replay.
80    fn new_execution(&mut self) -> Option<Schedule>;
81
82    /// Decide which task to run next, given a list of runnable tasks and the currently running
83    /// tasks. This method returns `Some(task)` where `task` is the runnable task to be executed
84    /// next; it may also return `None`, indicating that the execution engine should stop exploring
85    /// the current schedule.
86    ///
87    /// `is_yielding` is a hint to the scheduler that `current_task` has asked to yield (e.g.,
88    /// during a spin loop) and should be deprioritized.
89    ///
90    /// The list of runnable tasks is guaranteed to be non-empty. If `current_task` is `None`, the
91    /// execution has not yet begun.
92    fn next_task(
93        &mut self,
94        runnable_tasks: &[&Task],
95        current_task: Option<TaskId>,
96        is_yielding: bool,
97    ) -> Option<TaskId>;
98
99    /// Choose the next u64 value to return to the currently running task.
100    fn next_u64(&mut self) -> u64;
101}
102
103impl Scheduler for Box<dyn Scheduler + Send> {
104    fn new_execution(&mut self) -> Option<Schedule> {
105        self.as_mut().new_execution()
106    }
107
108    fn next_task(
109        &mut self,
110        runnable_tasks: &[&Task],
111        current_task: Option<TaskId>,
112        is_yielding: bool,
113    ) -> Option<TaskId> {
114        self.as_mut().next_task(runnable_tasks, current_task, is_yielding)
115    }
116
117    fn next_u64(&mut self) -> u64 {
118        self.as_mut().next_u64()
119    }
120}