From d6607c0584ae768e2f6f7d485966e56289f5f3f4 Mon Sep 17 00:00:00 2001
From: dragon-zhang <zhangzicheng@apache.org>
Date: Thu, 21 Nov 2024 14:59:06 +0800
Subject: [PATCH] code polish

---
 core/src/common/macros.rs       | 12 +++++++++---
 core/src/coroutine/suspender.rs | 12 ++++++++++--
 deny.toml                       |  2 +-
 open-coroutine/src/lib.rs       | 15 ++++++++++++---
 4 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/core/src/common/macros.rs b/core/src/common/macros.rs
index 7c40fcb0..16dd86f3 100644
--- a/core/src/common/macros.rs
+++ b/core/src/common/macros.rs
@@ -98,7 +98,8 @@ macro_rules! impl_current_for {
             /// Init the current.
             pub(crate) fn init_current(current: &Self) {
                 $name.with(|s| {
-                    s.borrow_mut()
+                    s.try_borrow_mut()
+                        .unwrap_or_else(|_| panic!("init {} current failed", stringify!($name)))
                         .push_front(core::ptr::from_ref(current).cast::<std::ffi::c_void>());
                 });
             }
@@ -108,7 +109,8 @@ macro_rules! impl_current_for {
             #[allow(unreachable_pub)]
             pub fn current<'current>() -> Option<&'current Self> {
                 $name.with(|s| {
-                    s.borrow()
+                    s.try_borrow()
+                        .unwrap_or_else(|_| panic!("get {} current failed", stringify!($name)))
                         .front()
                         .map(|ptr| unsafe { &*(*ptr).cast::<Self>() })
                 })
@@ -116,7 +118,11 @@ macro_rules! impl_current_for {
 
             /// Clean the current.
             pub(crate) fn clean_current() {
-                $name.with(|s| _ = s.borrow_mut().pop_front());
+                $name.with(|s| {
+                    _ = s.try_borrow_mut()
+                        .unwrap_or_else(|_| panic!("clean {} current failed", stringify!($name)))
+                        .pop_front();
+                });
             }
         }
     };
diff --git a/core/src/coroutine/suspender.rs b/core/src/coroutine/suspender.rs
index 4698db31..edfb8f0c 100644
--- a/core/src/coroutine/suspender.rs
+++ b/core/src/coroutine/suspender.rs
@@ -18,13 +18,21 @@ impl<Param, Yield> Suspender<'_, Param, Yield> {
     /// Delay the execution of the coroutine with an arg until `timestamp`.
     pub fn until_with(&self, arg: Yield, timestamp: u64) -> Param {
         TIMESTAMP.with(|s| {
-            s.borrow_mut().push_front(timestamp);
+            s.try_borrow_mut()
+                .unwrap_or_else(|_| panic!("init TIMESTAMP current failed"))
+                .push_front(timestamp);
         });
         self.suspend_with(arg)
     }
 
     pub(crate) fn timestamp() -> u64 {
-        TIMESTAMP.with(|s| s.borrow_mut().pop_front()).unwrap_or(0)
+        TIMESTAMP
+            .with(|s| {
+                s.try_borrow_mut()
+                    .unwrap_or_else(|_| panic!("get TIMESTAMP current failed"))
+                    .pop_front()
+            })
+            .unwrap_or(0)
     }
 }
 
diff --git a/deny.toml b/deny.toml
index bb7965af..f3c0845c 100644
--- a/deny.toml
+++ b/deny.toml
@@ -3,7 +3,7 @@ allow = [
   "Apache-2.0",
   "Apache-2.0 WITH LLVM-exception",
   "Unlicense",
-  "Unicode-DFS-2016",
+  "Unicode-3.0",
   "MIT"
 ]
 confidence-threshold = 0.95
diff --git a/open-coroutine/src/lib.rs b/open-coroutine/src/lib.rs
index f6d0aec2..0fbb9151 100644
--- a/open-coroutine/src/lib.rs
+++ b/open-coroutine/src/lib.rs
@@ -113,7 +113,16 @@ pub fn task<P: 'static, R: 'static, F: FnOnce(P) -> R>(f: F, param: P) -> JoinHa
         unsafe {
             let ptr = &mut *((input as *mut c_void).cast::<(F, P)>());
             let data = std::ptr::read_unaligned(ptr);
-            let result: &'static mut R = Box::leak(Box::new((data.0)(data.1)));
+            let result: &'static mut std::io::Result<R> = Box::leak(Box::new(
+                std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| (data.0)(data.1)))
+                    .map_err(|e| {
+                        Error::new(
+                            ErrorKind::Other,
+                            e.downcast_ref::<&'static str>()
+                                .map_or("task failed without message", |msg| *msg),
+                        )
+                    }),
+            ));
             std::ptr::from_mut(result).cast::<c_void>() as usize
         }
     }
@@ -140,7 +149,7 @@ impl<R> JoinHandle<R> {
             match ptr.cmp(&0) {
                 Ordering::Less => Err(Error::new(ErrorKind::Other, "timeout join failed")),
                 Ordering::Equal => Ok(None),
-                Ordering::Greater => Ok(Some(std::ptr::read_unaligned(ptr as *mut R))),
+                Ordering::Greater => Ok(Some((*Box::from_raw(ptr as *mut std::io::Result<R>))?)),
             }
         }
     }
@@ -151,7 +160,7 @@ impl<R> JoinHandle<R> {
             match ptr.cmp(&0) {
                 Ordering::Less => Err(Error::new(ErrorKind::Other, "join failed")),
                 Ordering::Equal => Ok(None),
-                Ordering::Greater => Ok(Some(std::ptr::read_unaligned(ptr as *mut R))),
+                Ordering::Greater => Ok(Some((*Box::from_raw(ptr as *mut std::io::Result<R>))?)),
             }
         }
     }