From 16748778823c0540f409af153bc5ad6d4ee7dcac Mon Sep 17 00:00:00 2001
From: Steve Lau <stevelauc@outlook.com>
Date: Sat, 2 Nov 2024 17:41:01 +0800
Subject: [PATCH] feat: impl spawn_catch_unwind()

---
 monoio/Cargo.toml     |  2 +-
 monoio/src/lib.rs     |  2 +-
 monoio/src/runtime.rs | 37 ++++++++++++++++++++++++++++++++++---
 3 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/monoio/Cargo.toml b/monoio/Cargo.toml
index eb25f186..7fe8216c 100644
--- a/monoio/Cargo.toml
+++ b/monoio/Cargo.toml
@@ -20,6 +20,7 @@ libc = "0.2"
 pin-project-lite = "0.2"
 socket2 = { version = "0.5", features = ["all"] }
 memchr = "2.7"
+futures = "0.3"
 
 bytes = { version = "1", optional = true }
 flume = { version = "0.11", optional = true }
@@ -55,7 +56,6 @@ nix = { version = "0.26", optional = true }
 io-uring = { version = "0.6", optional = true }
 
 [dev-dependencies]
-futures = "0.3"
 local-sync = "0.0.5"
 tempfile = "3.2"
 
diff --git a/monoio/src/lib.rs b/monoio/src/lib.rs
index 20cbb72d..9b733906 100644
--- a/monoio/src/lib.rs
+++ b/monoio/src/lib.rs
@@ -45,7 +45,7 @@ pub use driver::IoUringDriver;
 pub use driver::LegacyDriver;
 #[cfg(feature = "macros")]
 pub use monoio_macros::{main, test, test_all};
-pub use runtime::{spawn, Runtime};
+pub use runtime::{spawn, spawn_catch_unwind, Runtime};
 #[cfg(any(all(target_os = "linux", feature = "iouring"), feature = "legacy"))]
 pub use {builder::FusionDriver, runtime::FusionRuntime};
 
diff --git a/monoio/src/runtime.rs b/monoio/src/runtime.rs
index e4dc90af..be601f83 100644
--- a/monoio/src/runtime.rs
+++ b/monoio/src/runtime.rs
@@ -351,9 +351,6 @@ impl From<Runtime<TimeDriver<IoUringDriver>>> for FusionRuntime<TimeDriver<IoUri
 ///
 /// # Examples
 ///
-/// In this example, a server is started and `spawn` is used to start a new task
-/// that processes each received connection.
-///
 /// ```no_run
 /// #[monoio::main]
 /// async fn main() {
@@ -382,6 +379,40 @@ where
     join
 }
 
+/// Similar to [`spawn()`], but it catches any panic that occurs in the spawned task.
+///
+/// Note that `.await`ing returned `JoinHandle` now returns a `Result<T::Output>`
+/// rather than `T::Output`.
+///
+/// # Examples
+///
+/// ```no_run
+/// #[monoio::main]
+/// async fn main() {
+///     let handle = monoio::spawn(async {
+///         println!("hello from a background task");
+///     });
+///
+///     // Let the task complete
+///     handle.await.expect("no panic would happen!");
+/// }
+/// ```
+pub fn spawn_catch_unwind<T>(
+    future: T,
+) -> JoinHandle<Result<T::Output, Box<dyn std::any::Any + 'static>>>
+where
+    T: Future + 'static,
+    T::Output: 'static,
+{
+    use futures::{FutureExt, TryFutureExt};
+
+    let future = std::panic::AssertUnwindSafe(future)
+        .catch_unwind()
+        .map_err(|e| e as Box<dyn std::any::Any + 'static>);
+
+    spawn(future)
+}
+
 #[cfg(feature = "sync")]
 unsafe fn spawn_without_static<T>(future: T) -> JoinHandle<T::Output>
 where