simulation runtime task panic handling

Author: Shubham8287

crates/runtime/src/lib.rs

@@ -49,13 +49,32 @@ pub mod sync {
     pub use tokio::sync::watch;
 }
 
-#[derive(Clone)]
 pub enum Handle {
     Tokio(TokioHandle),
     #[cfg(feature = "simulation")]
     Simulation(sim::Handle),
 }
 
+impl Clone for Handle {
+    fn clone(&self) -> Self {
+        match self {
+            Self::Tokio(handle) => Self::Tokio(handle.clone()),
+            #[cfg(feature = "simulation")]
+            Self::Simulation(handle) => Self::Simulation(handle.clone()),
+        }
+    }
+}
+
+impl fmt::Debug for Handle {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::Tokio(_) => f.write_str("Handle::Tokio"),
+            #[cfg(feature = "simulation")]
+            Self::Simulation(_) => f.write_str("Handle::Simulation"),
+        }
+    }
+}
+
 pub struct JoinHandle<T> {
     inner: JoinHandleInner<T>,
 }
@@ -135,6 +154,16 @@ impl fmt::Display for JoinError {
     }
 }
 
+impl JoinError {
+    pub fn is_panic(&self) -> bool {
+        match &self.inner {
+            JoinErrorInner::Tokio(err) => err.is_panic(),
+            #[cfg(feature = "simulation")]
+            JoinErrorInner::Simulation(err) => err.is_panic(),
+        }
+    }
+}
+
 impl std::error::Error for JoinError {}
 
 impl<T> JoinHandleInner<T> {
@@ -234,6 +263,27 @@ impl Handle {
     pub fn simulation(handle: sim::Handle) -> Self {
         Self::Simulation(handle)
     }
+
+    pub fn on_simulation_node(&self, node: sim::NodeId) -> Self {
+        match self {
+            Self::Tokio(_) => panic!("Handle::on_simulation_node requires a simulation runtime"),
+            Self::Simulation(handle) => Self::Simulation(handle.on_node(node)),
+        }
+    }
+
+    pub fn create_simulation_node(&self) -> Option<sim::NodeBuilder> {
+        match self {
+            Self::Tokio(_) => panic!("Handle::create_simulation_node requires a simulation runtime"),
+            Self::Simulation(handle) => Some(handle.create_node()),
+        }
+    }
+
+    pub fn drain_simulation_task_panics(&self) -> Vec<sim::TaskPanic> {
+        match self {
+            Self::Tokio(_) => panic!("Handle::drain_simulation_task_panics requires a simulation runtime"),
+            Self::Simulation(handle) => handle.drain_task_panics(),
+        }
+    }
 }
 
 impl Handle {
@@ -244,7 +294,7 @@ impl Handle {
             },
             #[cfg(feature = "simulation")]
             Self::Simulation(handle) => JoinHandle {
-                inner: JoinHandleInner::Simulation(handle.spawn_on(sim::NodeId::MAIN, future)),
+                inner: JoinHandleInner::Simulation(handle.spawn(future)),
             },
         }
     }
@@ -268,7 +318,7 @@ impl Handle {
             // the simulation backend.
             #[cfg(feature = "simulation")]
             Self::Simulation(handle) => handle
-                .spawn_on(sim::NodeId::MAIN, async move { f() })
+                .spawn(async move { f() })
                 .await
                 .expect("simulation spawn_blocking task should not be cancelled"),
         }
@@ -295,6 +345,14 @@ impl Handle {
             Self::Simulation(handle) => handle.sleep(duration).await,
         }
     }
+
+    pub fn block_on<F: Future>(&self, future: F) -> F::Output {
+        match self {
+            Self::Tokio(handle) => tokio::task::block_in_place(|| handle.block_on(future)),
+            #[cfg(feature = "simulation")]
+            Self::Simulation(handle) => handle.block_on(future),
+        }
+    }
 }
 
 #[cfg(test)]

crates/runtime/src/sim/executor/mod.rs

@@ -83,6 +83,11 @@ pub struct Node {
     config: Arc<NodeConfig>,
 }
 
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct TaskPanic {
+    pub node: NodeId,
+}
+
 impl Node {
     /// Return the stable identifier for this simulated node.
     pub fn id(&self) -> NodeId {
@@ -158,10 +163,15 @@ impl Runtime {
         self.executor.elapsed()
     }
 
+    pub fn drain_task_panics(&self) -> Vec<TaskPanic> {
+        self.executor.drain_task_panics()
+    }
+
     /// Get a cloneable handle for spawning tasks and accessing runtime services.
     pub fn handle(&self) -> Handle {
         Handle {
             executor: Arc::clone(&self.executor),
+            node: NodeId::MAIN,
         }
     }
 
@@ -244,9 +254,32 @@ impl Runtime {
 #[derive(Clone)]
 pub struct Handle {
     executor: Arc<Executor>,
+    node: NodeId,
 }
 
 impl Handle {
+    /// Return a handle that spawns default work on `node`.
+    pub fn on_node(&self, node: NodeId) -> Self {
+        Self {
+            executor: Arc::clone(&self.executor),
+            node,
+        }
+    }
+
+    /// Return the node this handle uses for default spawned work.
+    pub fn node(&self) -> NodeId {
+        self.node
+    }
+
+    /// Spawn a `Send` future onto this handle's current simulated node.
+    pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
+    where
+        F: Future + Send + 'static,
+        F::Output: Send + 'static,
+    {
+        self.spawn_on(self.node, future)
+    }
+
     /// Create a new simulated node owned by this runtime.
     pub fn create_node(&self) -> NodeBuilder {
         NodeBuilder {
@@ -260,7 +293,7 @@ impl Handle {
         let config = self.executor.node_config(id);
         Node {
             id,
-            handle: self.clone(),
+            handle: self.on_node(id),
             config,
         }
     }
@@ -319,6 +352,10 @@ impl Handle {
         self.executor.time.timeout(duration, future).await
     }
 
+    pub fn block_on<F: Future>(&self, future: F) -> F::Output {
+        self.executor.block_on(future)
+    }
+
     pub fn enable_buggify(&self) {
         self.executor.enable_buggify();
     }
@@ -342,6 +379,10 @@ impl Handle {
     pub fn buggify_with_prob(&self, probability: f64) -> bool {
         self.executor.buggify_with_prob(probability)
     }
+
+    pub fn drain_task_panics(&self) -> Vec<TaskPanic> {
+        self.executor.drain_task_panics()
+    }
 }
 
 /// Core single-threaded scheduler backing a simulation [`Runtime`].
@@ -356,6 +397,7 @@ struct Executor {
     next_node: AtomicU64,
     rng: Rng,
     time: TimeHandle,
+    task_panics: Arc<Mutex<Vec<TaskPanic>>>,
 }
 
 impl Executor {
@@ -371,13 +413,18 @@ impl Executor {
             next_node: AtomicU64::new(1),
             rng: Rng::new(config.seed),
             time: TimeHandle::new(),
+            task_panics: Arc::new(Mutex::new(Vec::new())),
         }
     }
 
     fn elapsed(&self) -> Duration {
         self.time.now()
     }
 
+    fn drain_task_panics(&self) -> Vec<TaskPanic> {
+        core::mem::take(&mut *self.task_panics.lock())
+    }
+
     fn enable_buggify(&self) {
         self.rng.enable_buggify();
     }
@@ -439,7 +486,7 @@ impl Executor {
         self.assert_known_node(node);
 
         let abort = AbortHandle::new();
-        let abortable = Abortable::new(future, abort.clone());
+        let abortable = Abortable::new(future, abort.clone(), node, Arc::clone(&self.task_panics));
         let sender = self.sender.clone();
         let (runnable, task) = async_task::Builder::new()
             .metadata(node)
@@ -458,7 +505,7 @@ impl Executor {
         self.assert_known_node(node);
 
         let abort = AbortHandle::new();
-        let abortable = Abortable::new(future, abort.clone());
+        let abortable = Abortable::new(future, abort.clone(), node, Arc::clone(&self.task_panics));
         let sender = self.sender.clone();
         let (runnable, task) = unsafe {
             async_task::Builder::new()
@@ -731,7 +778,30 @@ mod tests {
         let err = runtime
             .block_on(task)
             .expect_err("aborted task should surface JoinError instead of panicking");
-        assert_eq!(err, JoinError);
+        assert_eq!(err, JoinError::Cancelled);
+    }
+
+    #[test]
+    fn spawned_task_panic_is_captured_without_stopping_runtime() {
+        let mut runtime = Runtime::new(10);
+        let node = runtime.create_node().name("panic").build();
+        let node_id = node.id();
+        let task = node.spawn(async move {
+            yield_now().await;
+            panic!("node boom");
+        });
+
+        runtime.block_on(async {
+            yield_now().await;
+            yield_now().await;
+        });
+
+        let panics = runtime.drain_task_panics();
+        assert_eq!(panics, vec![TaskPanic { node: node_id }]);
+        let err = runtime
+            .block_on(task)
+            .expect_err("panicked task should surface JoinError instead of unwinding");
+        assert!(err.is_panic());
     }
 
     #[cfg(feature = "simulation")]

crates/runtime/src/sim/executor/task.rs

@@ -1,4 +1,4 @@
-use alloc::sync::Arc;
+use alloc::{sync::Arc, vec::Vec};
 use core::{
     fmt,
     future::Future,
@@ -9,7 +9,7 @@ use core::{
 
 use spin::Mutex;
 
-use super::NodeId;
+use super::{NodeId, TaskPanic};
 
 /// A spawned simulated task.
 ///
@@ -87,12 +87,24 @@ impl AbortHandle {
     }
 }
 
-#[derive(Clone, Copy, Debug, Eq, PartialEq)]
-pub struct JoinError;
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub enum JoinError {
+    Cancelled,
+    Panicked,
+}
 
 impl fmt::Display for JoinError {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.write_str("task was cancelled")
+        match self {
+            Self::Cancelled => f.write_str("task was cancelled"),
+            Self::Panicked => f.write_str("task panicked"),
+        }
+    }
+}
+
+impl JoinError {
+    pub fn is_panic(&self) -> bool {
+        matches!(self, Self::Panicked)
     }
 }
 
@@ -124,11 +136,18 @@ impl AbortState {
 pub(crate) struct Abortable<F> {
     future: F,
     abort: AbortHandle,
+    node: NodeId,
+    task_panics: Arc<Mutex<Vec<TaskPanic>>>,
 }
 
 impl<F> Abortable<F> {
-    pub(crate) fn new(future: F, abort: AbortHandle) -> Self {
-        Self { future, abort }
+    pub(crate) fn new(future: F, abort: AbortHandle, node: NodeId, task_panics: Arc<Mutex<Vec<TaskPanic>>>) -> Self {
+        Self {
+            future,
+            abort,
+            node,
+            task_panics,
+        }
     }
 }
 
@@ -138,7 +157,7 @@ impl<F: Future> Future for Abortable<F> {
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         // Check cancellation before doing any work.
         if self.abort.state.aborted.load(Ordering::Relaxed) {
-            return Poll::Ready(Err(JoinError));
+            return Poll::Ready(Err(JoinError::Cancelled));
         }
 
         // Register the waker so abort() can wake this task.
@@ -153,7 +172,15 @@ impl<F: Future> Future for Abortable<F> {
         //      address of `future` relative to `self`.
         //   3. The caller guarantees `self` stays pinned for the lifetime of the
         //      future.
+        let node = self.node;
+        let task_panics = Arc::clone(&self.task_panics);
         let mut future = unsafe { self.map_unchecked_mut(|this| &mut this.future) };
-        future.as_mut().poll(cx).map(Ok)
+        match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| future.as_mut().poll(cx))) {
+            Ok(poll) => poll.map(Ok),
+            Err(_) => {
+                task_panics.lock().push(TaskPanic { node });
+                Poll::Ready(Err(JoinError::Panicked))
+            }
+        }
     }
 }

crates/runtime/src/sim/mod.rs

@@ -4,7 +4,7 @@ mod rng;
 pub mod time;
 
 pub use executor::{
-    yield_now, AbortHandle, Handle, JoinError, JoinHandle, Node, NodeBuilder, NodeId, Runtime, RuntimeConfig,
+    yield_now, AbortHandle, Handle, JoinError, JoinHandle, Node, NodeBuilder, NodeId, Runtime, RuntimeConfig, TaskPanic,
 };
 pub(crate) use rng::DeterminismLog;
 pub use rng::{GlobalRng, Rng};