Better error messages for usages of ExecutionState::with (#219)

Author: sarsko

shuttle/src/current.rs

@@ -87,12 +87,15 @@ pub fn set_name_for_task(task_id: TaskId, task_name: impl Into<TaskName>) -> Opt
     // when running with something like the `tracing_subscriber::fmt` subscriber.
     // This either has to be lived with, or the task name should be set via the `ChildLabelFn` mechanism, or a different subscriber should be used
     // (if this is done, then `record_steps_in_span` should be set to true as well), or Shuttle will have to be chanegd to recreate the Span
-    ExecutionState::try_with(|state| {
+    let res = ExecutionState::try_with(|state| {
         state
             .get_mut(task_id)
             .step_span
             .record("task", format!("{task_name:?}"));
     });
+    if let Err(e) = res {
+        tracing::error!("`set_name_for_task` failed with error: {e:?}");
+    }
     set_label_for_task::<TaskName>(task_id, task_name)
 }
 

shuttle/src/future/mod.rs

@@ -61,11 +61,14 @@ pub struct AbortHandle {
 impl AbortHandle {
     /// Abort the task associated with the handle.
     pub fn abort(&self) {
-        ExecutionState::try_with(|state| {
+        let res = ExecutionState::try_with(|state| {
             if !state.is_finished() {
                 state.get_mut(self.task_id).abort();
             }
         });
+        if let Err(e) = res {
+            tracing::error!("`AbortHandle::abort` failed with error: {e:?}");
+        }
     }
 
     /// Returns `true` if this task is finished, otherwise returns `false`.
@@ -108,11 +111,14 @@ impl<T> Default for JoinHandleInner<T> {
 impl<T> JoinHandle<T> {
     /// Abort the task associated with the handle.
     pub fn abort(&self) {
-        ExecutionState::try_with(|state| {
+        let res = ExecutionState::try_with(|state| {
             if !state.is_finished() {
                 state.get_mut(self.task_id).abort();
             }
         });
+        if let Err(e) = res {
+            tracing::error!("`JoinHandle::abort` failed with error: {e:?}");
+        }
     }
 
     /// Returns `true` if this task is finished, otherwise returns `false`.

shuttle/src/runtime/execution.rs

@@ -363,6 +363,15 @@ impl ScheduledTask {
     }
 }
 
+/// Error type for when an `ExecutionState::with` fails
+#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
+pub enum ExecutionStateBorrowError {
+    /// `ExecutionState` is currently not set
+    NotSet,
+    /// We are trying to borrow `ExecutionState` while it is already borrowed
+    AlreadyBorrowed,
+}
+
 impl ExecutionState {
     fn new(config: Config, scheduler: Rc<RefCell<dyn Scheduler>>) -> Self {
         Self {
@@ -387,30 +396,46 @@ impl ExecutionState {
     /// access to the state of the execution to influence scheduling (e.g. to register a task as
     /// blocked).
     #[inline]
+    #[track_caller]
     pub(crate) fn with<F, T>(f: F) -> T
     where
         F: FnOnce(&mut ExecutionState) -> T,
     {
-        Self::try_with(f).expect("Shuttle internal error: cannot access ExecutionState. are you trying to access a Shuttle primitive from outside a Shuttle test?")
+        Self::try_with(f).unwrap_or_else(|e| {
+            eprintln!("`ExecutionState::try_with` failed with error: {e:?}");
+            eprintln!(
+                "Backtrace for `with`: {:#?}",
+                std::backtrace::Backtrace::force_capture()
+            );
+            match e {
+                ExecutionStateBorrowError::AlreadyBorrowed => panic!("`ExecutionState::with` panicked because `ExecutionState` is already borrowed."),
+                ExecutionStateBorrowError::NotSet => panic!("`ExecutionState::with` panicked because `ExecutionState` is not set. Are you accessing a Shuttle primitive outside of a Shuttle test?"),
+            }
+        })
     }
 
     /// Like `with`, but returns None instead of panicking if there is no current ExecutionState or
     /// if the current ExecutionState is already borrowed.
     #[inline]
-    pub(crate) fn try_with<F, T>(f: F) -> Option<T>
+    #[track_caller]
+    pub(crate) fn try_with<F, T>(f: F) -> Result<T, ExecutionStateBorrowError>
     where
         F: FnOnce(&mut ExecutionState) -> T,
     {
+        trace!(
+            "ExecutionState::try_with called from {:?}",
+            std::panic::Location::caller()
+        );
         if EXECUTION_STATE.is_set() {
             EXECUTION_STATE.with(|cell| {
                 if let Ok(mut state) = cell.try_borrow_mut() {
-                    Some(f(&mut state))
+                    Ok(f(&mut state))
                 } else {
-                    None
+                    Err(ExecutionStateBorrowError::AlreadyBorrowed)
                 }
             })
         } else {
-            None
+            Err(ExecutionStateBorrowError::NotSet)
         }
     }
 
@@ -717,7 +742,7 @@ impl ExecutionState {
                 "<unknown>".into()
             }
         })
-        .unwrap_or_else(|| "ExecutionState is either not set or currently borrowed. This should not happen".to_string())
+        .unwrap_or_else(|e| format!("Tried to get ExecutionState, but got the following error: {e:?}"))
     }
 
     /// Generate a random u64 from the current scheduler and return it.

shuttle/tests/basic/execution.rs

@@ -240,7 +240,9 @@ fn context_switches_mutex() {
 /// Check that we get a good failure message if accessing a Shuttle primitive from outside an
 /// execution.
 #[test]
-#[should_panic(expected = "are you trying to access a Shuttle primitive from outside a Shuttle test?")]
+#[should_panic(
+    expected = "`ExecutionState::with` panicked because `ExecutionState` is not set. Are you accessing a Shuttle primitive outside of a Shuttle test?"
+)]
 fn failure_outside_execution() {
     let lock = shuttle::sync::Mutex::new(0u64);
     let _ = lock.lock().unwrap();