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

shuttle_engine/annotations/
mod.rs
