shuttle_engine/annotations/
mod.rs

1//! Annotated schedules. When an execution is scheduled using an
2//! `AnnotationScheduler`, Shuttle will produce a file that contains
3//! additional information about the execution, such as the kind of step that
4//! was taken (was a task created, were permits acquired from a semaphore, etc)
5//! as well as the task's vector clocks and thus any causal dependence between
6//! the tasks. The resulting file can be visualized using the Shuttle Explorer
7//! IDE extension.
8
9// TODO: the types defined here with `derive(Serialize)` are all parsed from
10//       JSON output by Shuttle Explorer; if any changes are made, they should
11//       also be reflected in the parsing
12// TODO: introduce version numbers to make sure breaking changes are noticed
13
14// `annotation` turns on `vector-clocks`, but `bench-no-vector-clocks` overrides it and would
15// leave every clock in the annotated schedule empty, hiding all causal dependence from the
16// Shuttle Explorer.
17#[cfg(all(feature = "annotation", feature = "bench-no-vector-clocks"))]
18compile_error!(
19    "the `annotation` feature requires vector clocks, so it cannot be combined with `bench-no-vector-clocks`"
20);
21
22cfg_if::cfg_if! {
23    if #[cfg(feature = "annotation")] {
24        use crate::annotation_file;
25        use crate::runtime::{
26            execution::ExecutionState,
27            task::{clock::VectorClock, Task, TaskId},
28        };
29        use serde::Serialize;
30        use std::cell::RefCell;
31        use std::collections::HashMap;
32        use std::thread_local;
33
34        thread_local! {
35            static ANNOTATION_STATE: RefCell<Option<AnnotationState>> = const { RefCell::new(None) };
36        }
37
38        #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
39        pub struct ObjectId(usize);
40
41        pub const DUMMY_OBJECT_ID: ObjectId = ObjectId(usize::MAX);
42
43        pub const ANNOTATION_VERSION: usize = 0;
44
45        /// Information about a file path found in one or more backtraces in the
46        /// annotated schedule. The path is stored in this type; instances of this
47        /// type are stored in the `files` vector in `AnnotationState`, and backtrace
48        /// frames then refer to paths using the index into the vector.
49        #[derive(Serialize)]
50        struct FileInfo {
51            path: String,
52        }
53
54        /// Information about a function name found in one or more backtraces in the
55        /// annotated schedule. The name is stored in this type; instances of this
56        /// type are stored in the `functions` vector in `AnnotationState`, and
57        /// backtrace frames then refer to functions using the index into the vector.
58        #[derive(Serialize)]
59        struct FunctionInfo {
60            name: String,
61        }
62
63        /// A backtrace frame.
64        #[derive(Serialize)]
65        struct Frame(
66            // file (index into `state.files`)
67            usize,
68            // function (index into `state.functions`)
69            usize,
70            // line
71            usize,
72            // column
73            usize,
74        );
75
76        /// Information about a shared object, i.e., a synchronization primitive
77        /// based on a batch semaphore.
78        #[derive(Serialize)]
79        struct ObjectInfo {
80            created_by: TaskId,
81            created_at: usize,
82            name: Option<String>,
83            kind: Option<String>,
84        }
85
86        /// Information about a task.
87        #[derive(Serialize)]
88        struct TaskInfo {
89            created_by: TaskId,
90            first_step: usize,
91            last_step: usize,
92            name: Option<String>,
93        }
94
95        #[derive(Debug, Serialize)]
96        enum AnnotationEvent {
97            SemaphoreCreated(ObjectId),
98            SemaphoreClosed(ObjectId),
99            SemaphoreAcquireFast(ObjectId, usize),
100            SemaphoreAcquireBlocked(ObjectId, usize),
101            SemaphoreAcquireUnblocked(ObjectId, TaskId, usize),
102            SemaphoreTryAcquire(ObjectId, usize, bool),
103            SemaphoreRelease(ObjectId, usize),
104
105            TaskCreated(TaskId, bool),
106            TaskTerminated,
107
108            Random,
109            Tick,
110        }
111
112        #[derive(Serialize)]
113        struct EventInfo(
114            // which task did something/yielded?
115            TaskId,
116            // backtrace
117            Option<Vec<Frame>>,
118            // event kind
119            AnnotationEvent,
120            // (if available,) clock of the task
121            // TODO: should always be available?
122            Option<VectorClock>,
123            // which other tasks were available to schedule, if this was a scheduled tick
124            Option<Vec<TaskId>>,
125        );
126
127        #[derive(Default, Serialize)]
128        struct AnnotationState {
129            version: usize,
130            files: Vec<FileInfo>,
131            #[serde(skip)]
132            path_to_file: HashMap<String, usize>,
133            functions: Vec<FunctionInfo>,
134            #[serde(skip)]
135            name_to_function: HashMap<String, usize>,
136            objects: Vec<ObjectInfo>,
137            tasks: Vec<TaskInfo>,
138            events: Vec<EventInfo>,
139
140            #[serde(skip)]
141            last_runnable_ids: Option<Vec<TaskId>>,
142            #[serde(skip)]
143            last_task_id: Option<TaskId>,
144            #[serde(skip)]
145            max_task_id: Option<TaskId>,
146        }
147
148        impl Serialize for VectorClock {
149            fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
150            where
151                S: serde::ser::Serializer,
152            {
153                use serde::ser::SerializeSeq;
154                // `VectorClock` derefs to `[u32]` in both its real and its stubbed
155                // form, so go through the slice rather than the `time` field, which
156                // only exists when the `vector-clocks` feature is enabled.
157                let mut seq = serializer.serialize_seq(Some(self.len()))?;
158                for e in self.iter() {
159                    seq.serialize_element(e)?;
160                }
161                seq.end()
162            }
163        }
164
165        impl Serialize for TaskId {
166            fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
167            where
168                S: serde::ser::Serializer,
169            {
170                usize::from(*self).serialize(serializer)
171            }
172        }
173
174        fn record_event(event: AnnotationEvent) {
175            with_state(move |state| {
176                let task_id = state.last_task_id.expect("no last task ID");
177
178                let task_id_num = usize::from(task_id);
179                assert!(task_id_num < state.tasks.len());
180                state.tasks[task_id_num].first_step = state.tasks[task_id_num].first_step.min(state.events.len());
181                state.tasks[task_id_num].last_step = state.tasks[task_id_num].last_step.max(state.events.len());
182
183                use std::backtrace::{Backtrace, BacktraceStatus};
184                use std::sync::OnceLock;
185                use regex::Regex;
186
187                // Here is a fragment of a backtrace for reference:
188                // ```
189                // 2: core::panicking::assert_failed_inner
190                // 3: core::panicking::assert_failed
191                //           at /rustc/129f3b9964af4d4a709d1383930ade12dfe7c081/library/core/src/panicking.rs:364:5
192                // 4: shuttle_clients::tests::example_impl
193                //           at ./src/example.rs:15:5
194                // 5: core::ops::function::Fn::call
195                //           at /rustc/129f3b9964af4d4a709d1383930ade12dfe7c081/library/core/src/ops/function.rs:79:5
196                // ```
197                // We want to capture the frames (numbered lines above) which
198                // refer to files local to the project being run, as well as
199                // the function name, and line/column info.
200                // TODO: for now, "local to the project" is detected based on
201                //       the path starting with `./src/`. Find an alternative
202                //       way to do this?
203                // The following regex matches frames with local paths. We rely
204                // on the string format of the backtrace because there is no
205                // better API. At the time of writing, even the unstable feature
206                // `backtrace_frames` does not provide a better interface.
207                // https://doc.rust-lang.org/std/backtrace/struct.BacktraceFrame.html
208                static RE: OnceLock<Regex> = OnceLock::new();
209                //                                         _num      function_name  path           line     col
210                let regex = RE.get_or_init(|| Regex::new(r"([0-9]+): ([^\n]+)\n +at (\./src/[^:]+):([0-9]+):([0-9]+)\b").unwrap());
211
212                // Whether or not the following call actually captures a backtrace
213                // depends on the environment variables `RUST_BACKTRACE` and
214                // `RUST_LIB_BACKTRACE`. See:
215                // https://doc.rust-lang.org/std/backtrace/index.html#environment-variables
216                // TODO: alternatively, we could use `Backtrace::force_capture`
217                //       and use our own environment flag.
218                let bt = Backtrace::capture();
219                let info = if bt.status() == BacktraceStatus::Captured {
220                    Some(regex
221                        // apply regex to debug-formatted backtrace
222                        .captures_iter(&format!("{bt}"))
223                        // for each match, extract the captured groups
224                        .map(|group| group.extract().1)
225                        // then store the extracted data into a `Frame`
226                        .map(|[_num, function_name, path, line_str, col_str]| {
227                            // intern file path in `state.files`
228                            let path_idx = *state
229                                .path_to_file
230                                .entry(path.to_string())
231                                .or_insert_with(|| {
232                                    let idx = state.files.len();
233                                    state.files.push(FileInfo {
234                                        path: path.to_string(),
235                                    });
236                                    idx
237                                });
238
239                            // intern function name in `state.functions`
240                            let function_idx = *state
241                                .name_to_function
242                                .entry(function_name.to_string())
243                                .or_insert_with(|| {
244                                    let idx = state.functions.len();
245                                    state.functions.push(FunctionInfo {
246                                        name: function_name.to_string(),
247                                    });
248                                    idx
249                                });
250
251                            Frame(
252                                path_idx,                           // file
253                                function_idx,                       // function
254                                line_str.parse::<usize>().unwrap(), // line
255                                col_str.parse::<usize>().unwrap(),  // col
256                            )
257                        })
258                        .collect::<Vec<_>>())
259                } else {
260                    None
261                };
262
263                state.events.push(EventInfo(
264                    task_id,
265                    info,
266                    event,
267                    ExecutionState::try_with(|state| state.get_clock(task_id).clone()).ok(),
268                    state.last_runnable_ids.take(),
269                ))
270            });
271        }
272
273        fn with_state<R, F: FnOnce(&mut AnnotationState) -> R>(f: F) -> Option<R> {
274            ANNOTATION_STATE.with(|cell| {
275                let mut bw = cell.borrow_mut();
276                let state = bw.as_mut()?;
277                Some(f(state))
278            })
279        }
280
281        fn record_object() -> ObjectId {
282            with_state(|state| {
283                let id = ObjectId(state.objects.len());
284                state.objects.push(ObjectInfo {
285                    created_by: state.last_task_id.unwrap(),
286                    created_at: state.events.len(),
287                    name: None,
288                    kind: None,
289                });
290                id
291            })
292            .unwrap_or(DUMMY_OBJECT_ID)
293        }
294
295        pub fn start_annotations() {
296            ANNOTATION_STATE.with(|cell| {
297                let mut bw = cell.borrow_mut();
298                assert!(bw.is_none(), "annotations already started");
299                let state = AnnotationState {
300                    version: ANNOTATION_VERSION,
301                    last_task_id: Some(0.into()),
302                    ..Default::default()
303                };
304                *bw = Some(state);
305            });
306        }
307
308        pub fn stop_annotations() {
309            ANNOTATION_STATE.with(|cell| {
310                let mut bw = cell.borrow_mut();
311                let state = bw.take().expect("annotations not started");
312                if state.max_task_id.is_none() {
313                    // nothing to output
314                    return;
315                };
316                let json = serde_json::to_string(&state).unwrap();
317                std::fs::write(
318                    annotation_file(),
319                    json,
320                )
321                .unwrap();
322            });
323        }
324
325        pub fn record_semaphore_created() -> ObjectId {
326            let object_id = record_object();
327            record_event(AnnotationEvent::SemaphoreCreated(object_id));
328            object_id
329        }
330
331        pub fn record_semaphore_closed(object_id: ObjectId) {
332            record_event(AnnotationEvent::SemaphoreClosed(object_id));
333        }
334
335        pub fn record_semaphore_acquire_fast(object_id: ObjectId, num_permits: usize) {
336            record_event(AnnotationEvent::SemaphoreAcquireFast(object_id, num_permits));
337        }
338
339        pub fn record_semaphore_acquire_blocked(object_id: ObjectId, num_permits: usize) {
340            record_event(AnnotationEvent::SemaphoreAcquireBlocked(object_id, num_permits));
341        }
342
343        pub fn record_semaphore_acquire_unblocked(object_id: ObjectId, unblocked_task_id: TaskId, num_permits: usize) {
344            record_event(AnnotationEvent::SemaphoreAcquireUnblocked(
345                object_id,
346                unblocked_task_id,
347                num_permits,
348            ));
349        }
350
351        pub fn record_semaphore_try_acquire(object_id: ObjectId, num_permits: usize, successful: bool) {
352            record_event(AnnotationEvent::SemaphoreTryAcquire(object_id, num_permits, successful));
353        }
354
355        pub fn record_semaphore_release(object_id: ObjectId, num_permits: usize) {
356            record_event(AnnotationEvent::SemaphoreRelease(object_id, num_permits));
357        }
358
359        pub fn record_task_created(task_id: TaskId, is_future: bool) {
360            with_state(move |state| {
361                assert_eq!(state.tasks.len(), usize::from(task_id));
362                state.tasks.push(TaskInfo {
363                    created_by: state.last_task_id.unwrap(),
364                    first_step: usize::MAX,
365                    last_step: 0,
366                    name: None,
367                });
368            });
369            record_event(AnnotationEvent::TaskCreated(task_id, is_future));
370        }
371
372        pub fn record_task_terminated() {
373            record_event(AnnotationEvent::TaskTerminated);
374        }
375
376        pub fn record_name_for_object(object_id: ObjectId, name: Option<&str>, kind: Option<&str>) {
377            with_state(move |state| {
378                if let Some(object_info) = state.objects.get_mut(object_id.0) {
379                    if name.is_some() {
380                        object_info.name = name.map(|name| name.to_string());
381                    }
382                    if kind.is_some() {
383                        object_info.kind = kind.map(|kind| kind.to_string());
384                    }
385                } // TODO: else panic? warn?
386            });
387        }
388
389        pub fn record_name_for_task(task_id: TaskId, name: &crate::current::TaskName) {
390            with_state(|state| {
391                if let Some(task_info) = state.tasks.get_mut(usize::from(task_id)) {
392                    let name: &String = name.into();
393                    task_info.name = Some(name.to_string());
394                } // TODO: else panic? warn?
395            });
396        }
397
398        pub fn record_random() {
399            record_event(AnnotationEvent::Random);
400        }
401
402        pub fn record_schedule(choice: TaskId, runnable_tasks: &[&Task]) {
403            with_state(|state| {
404                let choice_id_num = usize::from(choice);
405                state.tasks[choice_id_num].first_step = state.tasks[choice_id_num].first_step.min(state.events.len());
406                state.tasks[choice_id_num].last_step = state.tasks[choice_id_num].last_step.max(state.events.len());
407                assert!(
408                    state.last_runnable_ids.is_none(),
409                    "multiple schedule calls without a Tick"
410                );
411                state.last_runnable_ids = Some(runnable_tasks.iter().map(|task| task.id()).collect::<Vec<_>>());
412                state.last_task_id = Some(choice);
413                state.max_task_id = state.max_task_id.max(Some(choice));
414            });
415        }
416
417        pub fn record_tick() {
418            record_event(AnnotationEvent::Tick);
419        }
420    } else {
421        use crate::runtime::task::{Task, TaskId};
422
423        #[derive(Clone, Copy, Debug, PartialEq, Eq)]
424        pub struct ObjectId;
425
426        pub const DUMMY_OBJECT_ID: ObjectId = ObjectId;
427
428        #[inline(always)]
429        pub fn start_annotations() {}
430
431        #[inline(always)]
432        pub fn stop_annotations() {}
433
434        #[inline(always)]
435        pub fn record_semaphore_created() -> ObjectId {
436            DUMMY_OBJECT_ID
437        }
438
439        #[inline(always)]
440        pub fn record_semaphore_closed(_object_id: ObjectId) {}
441
442        #[inline(always)]
443        pub fn record_semaphore_acquire_fast(_object_id: ObjectId, _num_permits: usize) {}
444
445        #[inline(always)]
446        pub fn record_semaphore_acquire_blocked(_object_id: ObjectId, _num_permits: usize) {}
447
448        #[inline(always)]
449        pub fn record_semaphore_acquire_unblocked(_object_id: ObjectId, _unblocked_task_id: TaskId, _num_permits: usize) {}
450
451        #[inline(always)]
452        pub fn record_semaphore_try_acquire(_object_id: ObjectId, _num_permits: usize, _successful: bool) {}
453
454        #[inline(always)]
455        pub fn record_semaphore_release(_object_id: ObjectId, _num_permits: usize) {}
456
457        #[inline(always)]
458        pub fn record_task_created(_task_id: TaskId, _future: bool) {}
459
460        #[inline(always)]
461        pub fn record_task_terminated() {}
462
463        #[inline(always)]
464        pub fn record_name_for_object(_object_id: ObjectId, _name: Option<&str>, _kind: Option<&str>) {}
465
466        #[inline(always)]
467        pub fn record_name_for_task(_task_id: TaskId, _name: &crate::current::TaskName) {}
468
469        #[inline(always)]
470        pub fn record_random() {}
471
472        #[inline(always)]
473        pub fn record_schedule(_choice: TaskId, _runnable_tasks: &[&Task]) {}
474
475        #[inline(always)]
476        pub fn record_tick() {}
477    }
478}
479
480/// Trait to record information about shared objects, such as their name and
481/// type. See implementation in [`crate::future::batch_semaphore::BatchSemaphore`], which actually records the
482/// name into the schedule, other types should forward calls into their
483/// underlying primitive, as in `shuttle::sync::Mutex`.
484pub trait WithName {
485    /// Set the name and kind (full type path) of this object.
486    fn with_name_and_kind(self, name: Option<&str>, kind: Option<&str>) -> Self;
487
488    /// Set the name of this object.
489    fn with_name(self, name: &str) -> Self
490    where
491        Self: Sized,
492    {
493        self.with_name_and_kind(Some(name), None)
494    }
495
496    /// Set the kind (full type path) of this object.
497    fn with_kind(self, kind: &str) -> Self
498    where
499        Self: Sized,
500    {
501        self.with_name_and_kind(None, Some(kind))
502    }
503}
shuttle_engine/annotations/mod.rs

shuttle_engine/annotations/
mod.rs
