Add tag (awslabs#98)

* Add tag

* Tag: i64 -> u64

* Address comments from @jorajeev and @kraglb

* Add spawn_and_join test

* Add NOTE to get_current_task

* Address comments from @jorajeev

* Address Clippy too_many_arguments lint

* Address comments from @jorajeev

Author: sarsko

src/current.rs

@@ -7,7 +7,7 @@
 
 use crate::runtime::execution::ExecutionState;
 use crate::runtime::task::clock::VectorClock;
-use crate::runtime::task::TaskId;
+pub use crate::runtime::task::{Tag, TaskId};
 
 /// The number of context switches that happened so far in the current Shuttle execution.
 ///
@@ -22,7 +22,7 @@ pub fn context_switches() -> usize {
 
 /// Get the current thread's vector clock
 pub fn clock() -> VectorClock {
-    crate::runtime::execution::ExecutionState::with(|state| {
+    ExecutionState::with(|state| {
         let me = state.current();
         state.get_clock(me.id()).clone()
     })
@@ -32,3 +32,19 @@ pub fn clock() -> VectorClock {
 pub fn clock_for(task_id: TaskId) -> VectorClock {
     ExecutionState::with(|state| state.get_clock(task_id).clone())
 }
+
+/// Sets the `tag` field of the current task.
+/// Returns the `tag` which was there previously.
+pub fn set_tag_for_current_task(tag: Tag) -> Tag {
+    ExecutionState::set_tag_for_current_task(tag)
+}
+
+/// Gets the `tag` field of the current task.
+pub fn get_tag_for_current_task() -> Tag {
+    ExecutionState::get_tag_for_current_task()
+}
+
+/// Gets the `TaskId` of the current task, or `None` if there is no current task.
+pub fn get_current_task() -> Option<TaskId> {
+    ExecutionState::with(|s| Some(s.try_current()?.id()))
+}

src/runtime/execution.rs

@@ -1,6 +1,7 @@
 use crate::runtime::failure::{init_panic_hook, persist_failure, persist_task_failure};
 use crate::runtime::storage::{StorageKey, StorageMap};
 use crate::runtime::task::clock::VectorClock;
+use crate::runtime::task::Tag;
 use crate::runtime::task::{Task, TaskId, DEFAULT_INLINE_TASKS};
 use crate::runtime::thread::continuation::PooledContinuation;
 use crate::scheduler::{Schedule, Scheduler};
@@ -283,6 +284,7 @@ impl ExecutionState {
                 .map(|t| t.span.clone())
                 .unwrap_or_else(tracing::Span::current);
             let task_id = TaskId(state.tasks.len());
+            let tag = state.get_tag_or_default_for_current_task();
             let clock = state.increment_clock_mut(); // Increment the parent's clock
             clock.extend(task_id); // and extend it with an entry for the new task
 
@@ -294,6 +296,7 @@ impl ExecutionState {
                 clock.clone(),
                 parent_span,
                 schedule_len,
+                tag,
             );
 
             state.tasks.push(task);
@@ -312,7 +315,7 @@ impl ExecutionState {
     {
         Self::with(|state| {
             let task_id = TaskId(state.tasks.len());
-
+            let tag = state.get_tag_or_default_for_current_task();
             let clock = if let Some(ref mut clock) = initial_clock {
                 clock
             } else {
@@ -328,7 +331,7 @@ impl ExecutionState {
                 .try_current()
                 .map(|t| t.span.clone())
                 .unwrap_or_else(tracing::Span::current);
-            let task = Task::from_closure(f, stack_size, task_id, name, clock, parent_span, schedule_len);
+            let task = Task::from_closure(f, stack_size, task_id, name, clock, parent_span, schedule_len, tag);
             state.tasks.push(task);
             task_id
         })
@@ -605,6 +608,27 @@ impl ExecutionState {
             }
         });
     }
+
+    // Sets the `tag` field of the current task.
+    // Returns the `tag` which was there previously.
+    fn set_tag_for_current_task_internal(&mut self, tag: Tag) -> Tag {
+        self.current_mut().set_tag(tag)
+    }
+
+    pub(crate) fn set_tag_for_current_task(tag: Tag) -> Tag {
+        ExecutionState::with(|s| s.set_tag_for_current_task_internal(tag))
+    }
+
+    fn get_tag_or_default_for_current_task(&self) -> Tag {
+        match self.try_current() {
+            Some(current) => current.get_tag(),
+            None => Tag::default(),
+        }
+    }
+
+    pub(crate) fn get_tag_for_current_task() -> Tag {
+        ExecutionState::with(|s| s.get_tag_or_default_for_current_task())
+    }
 }
 
 #[cfg(debug_assertions)]

src/runtime/task/mod.rs

@@ -59,10 +59,14 @@ pub(crate) struct Task {
 
     local_storage: StorageMap,
     pub(super) span: tracing::Span,
+
+    // Arbitrarily settable tag which is inherited from the parent.
+    tag: Tag,
 }
 
 impl Task {
     /// Create a task from a continuation
+    #[allow(clippy::too_many_arguments)]
     fn new<F>(
         f: F,
         stack_size: usize,
@@ -71,6 +75,7 @@ impl Task {
         clock: VectorClock,
         parent_span: tracing::Span,
         schedule_len: usize,
+        tag: Tag,
     ) -> Self
     where
         F: FnOnce() + Send + 'static,
@@ -100,9 +105,11 @@ impl Task {
             name,
             span,
             local_storage: StorageMap::new(),
+            tag,
         }
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub(crate) fn from_closure<F>(
         f: F,
         stack_size: usize,
@@ -111,13 +118,15 @@ impl Task {
         clock: VectorClock,
         parent_span: tracing::Span,
         schedule_len: usize,
+        tag: Tag,
     ) -> Self
     where
         F: FnOnce() + Send + 'static,
     {
-        Self::new(f, stack_size, id, name, clock, parent_span, schedule_len)
+        Self::new(f, stack_size, id, name, clock, parent_span, schedule_len, tag)
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub(crate) fn from_future<F>(
         future: F,
         stack_size: usize,
@@ -126,6 +135,7 @@ impl Task {
         clock: VectorClock,
         parent_span: tracing::Span,
         schedule_len: usize,
+        tag: Tag,
     ) -> Self
     where
         F: Future<Output = ()> + Send + 'static,
@@ -147,6 +157,7 @@ impl Task {
             clock,
             parent_span,
             schedule_len,
+            tag,
         )
     }
 
@@ -335,6 +346,32 @@ impl Task {
             self.park_state.token_available = true;
         }
     }
+
+    pub(crate) fn get_tag(&self) -> Tag {
+        self.tag
+    }
+
+    /// Sets the `tag` field of the current task.
+    /// Returns the `tag` which was there previously.
+    pub(crate) fn set_tag(&mut self, tag: Tag) -> Tag {
+        std::mem::replace(&mut self.tag, tag)
+    }
+}
+
+/// A `Tag` is an arbitrarily settable value for each task.
+#[derive(PartialEq, Eq, Clone, Copy, Debug, Default, Hash, PartialOrd, Ord)]
+pub struct Tag(u64);
+
+impl From<u64> for Tag {
+    fn from(tag: u64) -> Self {
+        Tag(tag)
+    }
+}
+
+impl From<Tag> for u64 {
+    fn from(tag: Tag) -> u64 {
+        tag.0
+    }
 }
 
 #[derive(PartialEq, Eq, Clone, Copy, Debug)]

tests/basic/mod.rs

@@ -16,6 +16,7 @@ mod portfolio;
 mod replay;
 mod rwlock;
 mod shrink;
+mod tag;
 mod thread;
 mod timeout;
 mod uncontrolled_nondeterminism;

tests/basic/tag.rs

@@ -0,0 +1,135 @@
+use futures::future::join_all;
+use shuttle::{
+    check_dfs, check_random,
+    current::{get_tag_for_current_task, set_tag_for_current_task, Tag},
+    future::block_on,
+    thread,
+    thread::JoinHandle,
+};
+use test_log::test;
+
+fn spawn_some_futures_and_set_tag<F: (Fn(Tag, u64) -> Tag) + Send + Sync>(
+    tag_on_entry: Tag,
+    f: &'static F,
+    num_threads: u64,
+) {
+    let threads: Vec<_> = (0..num_threads)
+        .map(|i| {
+            shuttle::future::spawn(async move {
+                assert_eq!(get_tag_for_current_task(), tag_on_entry);
+                let new_tag = f(tag_on_entry, i);
+                set_tag_for_current_task(new_tag);
+                assert_eq!(get_tag_for_current_task(), new_tag);
+            })
+        })
+        .collect();
+
+    block_on(join_all(threads));
+
+    assert_eq!(get_tag_for_current_task(), tag_on_entry);
+}
+
+fn spawn_some_threads_and_set_tag<F: (Fn(Tag, u64) -> Tag) + Send + Sync>(
+    tag_on_entry: Tag,
+    f: &'static F,
+    num_threads: u64,
+) {
+    let threads: Vec<_> = (0..num_threads)
+        .map(|i| {
+            thread::spawn(move || {
+                assert_eq!(get_tag_for_current_task(), tag_on_entry);
+                let new_tag = f(tag_on_entry, i);
+                set_tag_for_current_task(new_tag);
+                assert_eq!(get_tag_for_current_task(), new_tag);
+            })
+        })
+        .collect();
+
+    threads.into_iter().for_each(|t| t.join().expect("Failed"));
+
+    assert_eq!(get_tag_for_current_task(), tag_on_entry);
+}
+
+fn spawn_threads_which_spawn_more_threads(
+    tag_on_entry: Tag,
+    num_threads_first_block: u64,
+    num_threads_second_block: u64,
+) {
+    let mut threads: Vec<_> = (0..num_threads_first_block)
+        .map(|i| {
+            thread::spawn(move || {
+                assert_eq!(get_tag_for_current_task(), tag_on_entry);
+                let new_tag = i.into();
+                set_tag_for_current_task(new_tag);
+                assert_eq!(get_tag_for_current_task(), new_tag);
+                spawn_some_threads_and_set_tag(new_tag, &|_, _| 123.into(), 13);
+                assert_eq!(get_tag_for_current_task(), new_tag);
+                spawn_some_threads_and_set_tag(new_tag, &|_, x| (x * 13).into(), 7);
+                assert_eq!(get_tag_for_current_task(), new_tag);
+                spawn_some_threads_and_set_tag(new_tag, &|p, x| ((u64::from(p) << 4) + x).into(), 19);
+                assert_eq!(get_tag_for_current_task(), new_tag);
+                spawn_some_futures_and_set_tag(new_tag, &|p, x| ((u64::from(p) << 4) & x).into(), 17);
+                assert_eq!(get_tag_for_current_task(), new_tag);
+            })
+        })
+        .collect();
+
+    assert_eq!(get_tag_for_current_task(), tag_on_entry);
+
+    let new_tag_main_thread = 987654321.into();
+    set_tag_for_current_task(new_tag_main_thread);
+    assert_eq!(get_tag_for_current_task(), new_tag_main_thread);
+
+    threads.extend(
+        (0..num_threads_second_block)
+            .map(|i| {
+                thread::spawn(move || {
+                    assert_eq!(get_tag_for_current_task(), new_tag_main_thread);
+                    let new_tag = i.into();
+                    set_tag_for_current_task(new_tag);
+                    assert_eq!(get_tag_for_current_task(), new_tag);
+                    spawn_some_threads_and_set_tag(new_tag, &|_, _| 123.into(), 13);
+                    assert_eq!(get_tag_for_current_task(), new_tag);
+                    spawn_some_threads_and_set_tag(new_tag, &|_, x| (x * 13).into(), 7);
+                    assert_eq!(get_tag_for_current_task(), new_tag);
+                    spawn_some_threads_and_set_tag(new_tag, &|p, x| ((u64::from(p) << 4) + x).into(), 19);
+                    assert_eq!(get_tag_for_current_task(), new_tag);
+                    spawn_some_futures_and_set_tag(new_tag, &|p, x| ((u64::from(p) << 4) & x).into(), 17);
+                    assert_eq!(get_tag_for_current_task(), new_tag);
+                })
+            })
+            .collect::<Vec<_>>(),
+    );
+
+    threads.into_iter().for_each(|t| t.join().expect("Failed"));
+
+    assert_eq!(get_tag_for_current_task(), new_tag_main_thread);
+}
+
+#[test]
+fn threads_which_spawn_threads_which_spawn_threads() {
+    check_random(|| spawn_threads_which_spawn_more_threads(Tag::default(), 3, 2), 10)
+}
+
+fn spawn_thread_and_set_tag(tag_on_entry: Tag, new_tag: Tag) -> JoinHandle<u64> {
+    thread::spawn(move || {
+        assert_eq!(get_tag_for_current_task(), tag_on_entry);
+        assert_eq!(set_tag_for_current_task(new_tag), tag_on_entry); // NOTE: Assertion with side effect
+        assert_eq!(get_tag_for_current_task(), new_tag);
+        new_tag.into()
+    })
+}
+
+fn spawn_and_join() {
+    set_tag_for_current_task(42.into());
+    let h1 = spawn_thread_and_set_tag(42.into(), 84.into());
+    set_tag_for_current_task(50.into());
+    let h2 = spawn_thread_and_set_tag(50.into(), 100.into());
+    let results = [h1.join().unwrap(), h2.join().unwrap()];
+    assert_eq!(results, [84, 100]);
+}
+
+#[test]
+fn test_spawn_and_join() {
+    check_dfs(spawn_and_join, None);
+}