Adding semantic next_event information to Task struct

Author: Dylan Wolff

shuttle/src/future/batch_semaphore.rs

@@ -1,7 +1,7 @@
 //! A counting semaphore supporting both async and sync operations.
 use crate::current;
 use crate::runtime::execution::ExecutionState;
-use crate::runtime::task::{clock::VectorClock, TaskId};
+use crate::runtime::task::{clock::VectorClock, Event, TaskId};
 use crate::runtime::thread;
 use crate::sync::{ResourceSignature, ResourceType};
 use std::cell::RefCell;
@@ -289,7 +289,6 @@ impl BatchSemaphoreState {
 pub struct BatchSemaphore {
     state: RefCell<BatchSemaphoreState>,
     fairness: Fairness,
-    #[allow(unused)]
     signature: ResourceSignature,
 }
 
@@ -392,8 +391,9 @@ impl BatchSemaphore {
 
     /// Closes the semaphore. This prevents the semaphore from issuing new
     /// permits and notifies all pending waiters.
+    #[track_caller]
     pub fn close(&self) {
-        thread::switch_task();
+        thread::switch(Event::batch_semaphore_rel(&self.signature));
 
         self.init_object_id();
         let mut state = self.state.borrow_mut();
@@ -436,8 +436,9 @@ impl BatchSemaphore {
     /// If the permits are available, returns Ok(())
     /// If the semaphore is closed, returns `Err(TryAcquireError::Closed)`
     /// If there aren't enough permits, returns `Err(TryAcquireError::NoPermits)`
+    #[track_caller]
     pub fn try_acquire(&self, num_permits: usize) -> Result<(), TryAcquireError> {
-        thread::switch_task();
+        thread::switch(Event::batch_semaphore_acq(&self.signature));
 
         self.init_object_id();
         let mut state = self.state.borrow_mut();
@@ -538,22 +539,25 @@ impl BatchSemaphore {
     }
 
     /// Acquire the specified number of permits (async API)
+    #[track_caller]
     pub fn acquire(&self, num_permits: usize) -> Acquire<'_> {
         // No switch here; switch should be triggered on polling future
         self.init_object_id();
         Acquire::new(self, num_permits)
     }
 
     /// Acquire the specified number of permits (blocking API)
+    #[track_caller]
     pub fn acquire_blocking(&self, num_permits: usize) -> Result<(), AcquireError> {
         // No switch here; switch should be triggered on polling future
         self.init_object_id();
         crate::future::block_on(self.acquire(num_permits))
     }
 
     /// Release `num_permits` back to the Semaphore
+    #[track_caller]
     pub fn release(&self, num_permits: usize) {
-        thread::switch_task();
+        thread::switch(Event::batch_semaphore_rel(&self.signature));
 
         self.init_object_id();
         if num_permits == 0 {
@@ -645,6 +649,7 @@ pub struct Acquire<'a> {
 }
 
 impl<'a> Acquire<'a> {
+    #[track_caller]
     fn new(semaphore: &'a BatchSemaphore, num_permits: usize) -> Self {
         let waiter = Arc::new(Waiter::new(num_permits));
         Self {
@@ -675,7 +680,7 @@ impl Future for Acquire<'_> {
         // available, the extra waiter doesn't affect other tasks -- all waiters and active tasks will race
         // to acquire permits when the current holder releases.
         if self.never_polled && (will_succeed || blocking_changes_state) {
-            thread::switch_task();
+            thread::switch(Event::batch_semaphore_acq(&self.semaphore.signature));
             self.never_polled = false;
         }
 
@@ -764,12 +769,25 @@ impl Future for Acquire<'_> {
                             self.waiter.is_queued.store(true, Ordering::SeqCst);
                         }
                         trace!("Acquire::poll for waiter {:?} that is enqueued", self.waiter);
+
+                        let event = Event::batch_semaphore_acq(&self.semaphore.signature);
+                        ExecutionState::with(|s| unsafe { s.current_mut().set_next_event(event) });
+                        // SAFETY: This is safe because the current task immediately suspends after this future
+                        // returns Poll::Pending (src/future/mod.rs). Whenever a task resumes, the `next_event`
+                        // is unset, so there is no opportunity to corrupt the reference to our signature while
+                        // it is set as the `next_task`.
                         Poll::Pending
                     }
                     Err(TryAcquireError::Closed) => unreachable!(),
                 }
             } else {
                 // No progress made, future is still pending.
+                let event = Event::batch_semaphore_acq(&self.semaphore.signature);
+                // SAFETY: This is safe because the current task immediately suspends after this future
+                // returns Poll::Pending (src/future/mod.rs). Whenever a task resumes, the `next_event`
+                // is unset, so there is no opportunity to corrupt the reference to our signature while
+                // it is set as the `next_task`.
+                ExecutionState::with(|s| unsafe { s.current_mut().set_next_event(event) });
                 Poll::Pending
             }
         }

shuttle/src/future/mod.rs

@@ -251,7 +251,7 @@ pub fn block_on<F: Future>(future: F) -> F::Output {
             Poll::Ready(result) => break result,
             Poll::Pending => {
                 ExecutionState::with(|state| state.current_mut().sleep_unless_woken());
-                thread::switch_task();
+                thread::switch_keep_event();
             }
         }
     }

shuttle/src/runtime/execution.rs

@@ -2,7 +2,7 @@ use crate::runtime::failure::{init_panic_hook, persist_failure, persist_task_fai
 use crate::runtime::storage::{StorageKey, StorageMap};
 use crate::runtime::task::clock::VectorClock;
 use crate::runtime::task::labels::Labels;
-use crate::runtime::task::{ChildLabelFn, Task, TaskId, TaskName, TaskSignature, DEFAULT_INLINE_TASKS};
+use crate::runtime::task::{ChildLabelFn, Event, Task, TaskId, TaskName, TaskSignature, DEFAULT_INLINE_TASKS};
 use crate::runtime::thread;
 use crate::runtime::thread::continuation::PooledContinuation;
 use crate::scheduler::{Schedule, Scheduler};
@@ -472,7 +472,8 @@ impl ExecutionState {
     where
         F: Future<Output = ()> + 'static,
     {
-        thread::switch_task();
+        let signature = ExecutionState::with(|state| state.current_mut().signature.new_child(caller));
+        thread::switch(Event::Spawn(&signature));
         let task_id = Self::with(|state| {
             let schedule_len = state.current_schedule.len();
             let parent_span_id = state.top_level_span.id();
@@ -495,7 +496,7 @@ impl ExecutionState {
                 schedule_len,
                 tag,
                 Some(state.current().id()),
-                state.current_mut().signature.new_child(caller),
+                signature,
             );
 
             state.tasks.push(task);
@@ -515,7 +516,8 @@ impl ExecutionState {
         mut initial_clock: Option<VectorClock>,
         caller: &'static Location<'static>,
     ) -> TaskId {
-        thread::switch_task();
+        let signature = ExecutionState::with(|state| state.current_mut().signature.new_child(caller));
+        thread::switch(Event::Spawn(&signature));
         let task_id = Self::with(|state| {
             let parent_span_id = state.top_level_span.id();
             let task_id = TaskId(state.tasks.len());
@@ -544,7 +546,7 @@ impl ExecutionState {
                 schedule_len,
                 tag,
                 Some(state.current().id()),
-                state.current_mut().signature.new_child(caller),
+                signature,
             );
             state.tasks.push(task);
 
@@ -819,7 +821,7 @@ impl ExecutionState {
             .scheduler
             .borrow_mut()
             .next_task(task_refs, self.current_task.id(), is_yielding)
-            .map(ScheduledTask::Some)
+            .map(|task| ScheduledTask::Some(task.id()))
             .unwrap_or(ScheduledTask::Stopped);
 
         // Tracing this `in_scope` is purely a matter of taste. We do it because

shuttle/src/runtime/runner.rs

@@ -237,12 +237,12 @@ impl<S: Scheduler> Scheduler for PortfolioStoppableScheduler<S> {
         }
     }
 
-    fn next_task(
+    fn next_task<'a>(
         &mut self,
-        runnable_tasks: &[&Task],
+        runnable_tasks: &'a [&'a Task],
         current_task: Option<TaskId>,
         is_yielding: bool,
-    ) -> Option<TaskId> {
+    ) -> Option<&'a Task> {
         if self.stop_signal.load(Ordering::SeqCst) {
             None
         } else {

shuttle/src/runtime/task/mod.rs

@@ -251,6 +251,107 @@ impl PartialEq for TaskSignature {
 
 impl Eq for TaskSignature {}
 
+pub(crate) type Loc = &'static Location<'static>;
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub(crate) enum Event<'a> {
+    AtomicRead(&'a ResourceSignature, Loc),
+    AtomicWrite(&'a ResourceSignature, Loc),
+    AtomicReadWrite(&'a ResourceSignature, Loc),
+    BatchSemaphoreAcq(&'a ResourceSignature, Loc),
+    BatchSemaphoreRel(&'a ResourceSignature, Loc),
+    BarrierWait(&'a ResourceSignature, Loc),
+    CondvarWait(&'a ResourceSignature, Loc),
+    CondvarNotify(Loc),
+    Park(Loc),
+    Unpark(&'a TaskSignature, Loc),
+    ChannelSend(&'a ResourceSignature, Loc),
+    ChannelRecv(&'a ResourceSignature, Loc),
+    Spawn(&'a TaskSignature),
+    Yield(Loc),
+    Sleep(Loc),
+    Exit,
+    Join(&'a TaskSignature, Loc),
+    Unknown,
+}
+
+impl<'a> Event<'a> {
+    #[track_caller]
+    pub(crate) fn atomic_read(sig: &'a ResourceSignature) -> Self {
+        Self::AtomicRead(sig, Location::caller())
+    }
+
+    #[track_caller]
+    pub(crate) fn atomic_write(sig: &'a ResourceSignature) -> Self {
+        Self::AtomicWrite(sig, Location::caller())
+    }
+
+    #[track_caller]
+    pub(crate) fn atomic_read_write(sig: &'a ResourceSignature) -> Self {
+        Self::AtomicReadWrite(sig, Location::caller())
+    }
+
+    #[track_caller]
+    pub(crate) fn batch_semaphore_acq(sig: &'a ResourceSignature) -> Self {
+        Self::BatchSemaphoreAcq(sig, Location::caller())
+    }
+
+    #[track_caller]
+    pub(crate) fn batch_semaphore_rel(sig: &'a ResourceSignature) -> Self {
+        Self::BatchSemaphoreRel(sig, Location::caller())
+    }
+
+    #[track_caller]
+    pub(crate) fn barrier_wait(sig: &'a ResourceSignature) -> Self {
+        Self::BarrierWait(sig, Location::caller())
+    }
+
+    #[track_caller]
+    pub(crate) fn condvar_wait(sig: &'a ResourceSignature) -> Self {
+        Self::CondvarWait(sig, Location::caller())
+    }
+
+    #[track_caller]
+    pub(crate) fn condvar_notify() -> Self {
+        Self::CondvarNotify(Location::caller())
+    }
+
+    #[track_caller]
+    pub(crate) fn park() -> Self {
+        Self::Park(Location::caller())
+    }
+
+    #[track_caller]
+    pub(crate) fn unpark(sig: &'a TaskSignature) -> Self {
+        Self::Unpark(sig, Location::caller())
+    }
+
+    #[track_caller]
+    pub(crate) fn channel_send(sig: &'a ResourceSignature) -> Self {
+        Self::ChannelSend(sig, Location::caller())
+    }
+
+    #[track_caller]
+    pub(crate) fn channel_recv(sig: &'a ResourceSignature) -> Self {
+        Self::ChannelRecv(sig, Location::caller())
+    }
+
+    #[track_caller]
+    pub(crate) fn yield_now() -> Self {
+        Self::Yield(Location::caller())
+    }
+
+    #[track_caller]
+    pub(crate) fn sleep() -> Self {
+        Self::Sleep(Location::caller())
+    }
+
+    #[track_caller]
+    pub(crate) fn join(sig: &'a TaskSignature) -> Self {
+        Self::Join(sig, Location::caller())
+    }
+}
+
 /// A `Task` represents a user-level unit of concurrency. Each task has an `id` that is unique within
 /// the execution, and a `state` reflecting whether the task is runnable (enabled) or not.
 #[derive(Debug)]
@@ -271,6 +372,8 @@ pub struct Task {
     // Remember whether the waker was invoked while we were running
     woken: bool,
 
+    next_event: Event<'static>,
+
     name: Option<String>,
 
     local_storage: StorageMap,
@@ -342,6 +445,7 @@ impl Task {
             waiter: None,
             waker,
             woken: false,
+            next_event: Event::Unknown,
             detached: false,
             park_state: ParkState::default(),
             name,
@@ -416,7 +520,7 @@ impl Task {
                 let cx = &mut Context::from_waker(&waker);
                 while future.as_mut().poll(cx).is_pending() {
                     ExecutionState::with(|state| state.current_mut().sleep_unless_woken());
-                    thread::switch_task();
+                    thread::switch_keep_event();
                 }
             }),
             stack_size,
@@ -651,6 +755,19 @@ impl Task {
         TASK_ID_TO_TAGS.with(|cell| cell.borrow_mut().insert(self.id(), tag.clone()));
         self.tag.replace(tag)
     }
+
+    /// Get the next_event with a downcast lifetime tied to self.
+    pub(crate) fn next_event(&self) -> &Event<'_> {
+        unsafe { std::mem::transmute(&self.next_event) }
+    }
+
+    pub(crate) unsafe fn set_next_event(&mut self, event: Event<'_>) {
+        self.next_event = unsafe { std::mem::transmute::<Event<'_>, Event<'static>>(event) };
+    }
+
+    pub(crate) fn unset_next_event(&mut self) {
+        self.next_event = Event::Unknown;
+    }
 }
 
 #[derive(PartialEq, Eq, Clone, Copy, Debug)]

shuttle/src/runtime/thread/continuation.rs

@@ -5,6 +5,7 @@
 #![allow(deprecated)]
 
 use crate::runtime::execution::ExecutionState;
+use crate::runtime::task::Event;
 use generator::{Generator, Gn};
 use scoped_tls::scoped_thread_local;
 use std::cell::{Cell, RefCell};
@@ -277,13 +278,24 @@ unsafe impl Send for PooledContinuation {}
 /// is a visible operation, meaning that both scheduling points are necessary for complete
 /// exploration of all possible behaviors.
 #[track_caller]
-pub(crate) fn switch_task() {
+pub(crate) fn switch(event: Event<'_>) {
+    // SAFETY we cast the lifetime of the Event to 'static when embedding it into the current Task
+    // This is safe because (1) we have a valid reference to the Event's data for the scope of this function
+    // and (2) the static reference is dropped at the end of the scope of this function when the next_event
+    // is set back to Unknown
+    ExecutionState::with(|s| unsafe { s.current_mut().set_next_event(event) });
+    switch_keep_event()
+}
+
+#[track_caller]
+pub(crate) fn switch_keep_event() {
     crate::annotations::record_tick();
     trace!("switch from {}", Location::caller());
     if ExecutionState::maybe_yield() {
         let r = generator::yield_(ContinuationOutput::Yielded).unwrap();
         assert!(matches!(r, ContinuationInput::Resume));
     }
+    ExecutionState::with(|s| s.current_mut().unset_next_event());
 }
 
 #[cfg(test)]

shuttle/src/runtime/thread/mod.rs

@@ -1,3 +1,3 @@
 pub(crate) mod continuation;
-
-pub(crate) use continuation::switch_task;
+pub(crate) use continuation::switch;
+pub(crate) use continuation::switch_keep_event;