shuttle_engine/scheduler/
metrics.rs

1use crate::runtime::task::{Task, TaskId};
2use crate::scheduler::{Schedule, Scheduler};
3use tracing::info;
4
5/// A `MetricsScheduler` wraps an inner `Scheduler` and collects metrics about the schedules it's
6/// generating.
7#[derive(Debug)]
8pub struct MetricsScheduler<S: ?Sized + Scheduler> {
9    inner: Box<S>,
10
11    iterations: usize,
12    iteration_divisor: usize,
13
14    steps: usize,
15    steps_metric: CountSummaryMetric,
16
17    last_task: TaskId,
18    // Number of times the scheduled task changed
19    context_switches: usize,
20    context_switches_metric: CountSummaryMetric,
21    // Number of times the scheduled task changed when the previous task was still runnable
22    preemptions: usize,
23    preemptions_metric: CountSummaryMetric,
24
25    random_choices: usize,
26    random_choices_metric: CountSummaryMetric,
27}
28
29impl<S: Scheduler> MetricsScheduler<S> {
30    /// Create a new `MetricsScheduler` by wrapping the given `Scheduler` implementation.
31    pub fn new(inner: S) -> Self {
32        Self {
33            inner: Box::new(inner),
34
35            iterations: 0,
36            iteration_divisor: 10,
37
38            steps: 0,
39            steps_metric: CountSummaryMetric::new(),
40
41            last_task: TaskId::from(0),
42            context_switches: 0,
43            context_switches_metric: CountSummaryMetric::new(),
44            preemptions: 0,
45            preemptions_metric: CountSummaryMetric::new(),
46
47            random_choices: 0,
48            random_choices_metric: CountSummaryMetric::new(),
49        }
50    }
51}
52
53impl<S: ?Sized + Scheduler> MetricsScheduler<S> {
54    fn record_and_reset_metrics(&mut self) {
55        self.steps_metric.record(self.steps);
56        self.steps = 0;
57
58        self.context_switches_metric.record(self.context_switches);
59        self.context_switches = 0;
60        self.preemptions_metric.record(self.preemptions);
61        self.preemptions = 0;
62
63        self.random_choices_metric.record(self.random_choices);
64        self.random_choices = 0;
65    }
66}
67
68impl<S: Scheduler> Scheduler for MetricsScheduler<S> {
69    fn new_execution(&mut self) -> Option<Schedule> {
70        if self.iterations > 0 {
71            self.record_and_reset_metrics();
72
73            if self.iterations.is_multiple_of(self.iteration_divisor) {
74                info!(iterations = self.iterations);
75
76                if self.iterations == self.iteration_divisor * 10 {
77                    self.iteration_divisor *= 10;
78                }
79            }
80        }
81        self.iterations += 1;
82
83        self.inner.new_execution()
84    }
85
86    fn next_task(
87        &mut self,
88        runnable_tasks: &[&Task],
89        current_task: Option<TaskId>,
90        is_yielding: bool,
91    ) -> Option<TaskId> {
92        let choice = self.inner.next_task(runnable_tasks, current_task, is_yielding)?;
93
94        self.steps += 1;
95        if choice != self.last_task {
96            self.context_switches += 1;
97            if runnable_tasks.iter().any(|t| t.id() == self.last_task) {
98                self.preemptions += 1;
99            }
100        }
101        self.last_task = choice;
102
103        Some(choice)
104    }
105
106    fn next_u64(&mut self) -> u64 {
107        self.steps += 1;
108        self.random_choices += 1;
109        self.inner.next_u64()
110    }
111}
112
113impl<S: ?Sized + Scheduler> Drop for MetricsScheduler<S> {
114    fn drop(&mut self) {
115        // If steps > 0 then we didn't get a chance to record the metrics for the current execution
116        // (it's probably panicking), so record them now
117        if self.steps > 0 {
118            self.record_and_reset_metrics();
119            self.iterations += 1;
120        }
121
122        info!(
123            iterations = self.iterations.saturating_sub(1),
124            steps = %self.steps_metric,
125            context_switches = %self.context_switches_metric,
126            preemptions = %self.preemptions_metric,
127            random_choices = %self.random_choices_metric,
128            "run finished"
129        );
130    }
131}
132
133/// A simple thing that can record a stream of `usize` values, and then report their min, max,
134/// and average.
135#[derive(Debug)]
136struct CountSummaryMetric {
137    min: usize,
138    sum: usize,
139    max: usize,
140    n: usize,
141}
142
143impl CountSummaryMetric {
144    fn new() -> Self {
145        Self {
146            min: usize::MAX,
147            sum: 0,
148            max: 0,
149            n: 0,
150        }
151    }
152
153    fn record(&mut self, value: usize) {
154        if value < self.min {
155            self.min = value;
156        }
157        if value > self.max {
158            self.max = value;
159        }
160        self.sum += value;
161        self.n += 1;
162    }
163}
164
165impl std::fmt::Display for CountSummaryMetric {
166    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
167        write!(f, "[min={}, max={},", self.min, self.max)?;
168        if self.n > 0 {
169            let avg = self.sum as f64 / self.n as f64;
170            write!(f, " avg={avg:.1}")?;
171        }
172        write!(f, "]")
173    }
174}
shuttle_engine/scheduler/metrics.rs

shuttle_engine/scheduler/
metrics.rs
