refactor(driver): handle panicking

Author: George-Miao

compio-driver/src/asyncify.rs

@@ -76,7 +76,7 @@ fn worker(
         counter.fetch_add(1, Ordering::AcqRel);
         let _guard = CounterGuard(counter);
         while let Ok(f) = receiver.recv_timeout(timeout) {
-            f.run();
+            f.run()
         }
     }
 }

compio-driver/src/lib.rs

@@ -25,6 +25,7 @@ use compio_buf::BufResult;
 use compio_log::instrument;
 
 mod macros;
+mod panic;
 
 mod key;
 pub use key::Key;
@@ -49,7 +50,7 @@ pub use cancel::*;
 
 mod control;
 
-use crate::{key::ErasedKey, op::OpCodeFlag};
+use crate::{key::ErasedKey, op::OpCodeFlag, panic::resume_unwind_io};
 
 mod sys_slice;
 
@@ -276,12 +277,15 @@ impl Proactor {
     ///
     /// # Panics
     ///
-    /// This function will panic if the [`Key`] is not unique.
+    /// This function will panic if the [`Key`] is not unique or if the
+    /// operation is blocking and it panicked in the thread pool.
     pub fn pop<T: OpCode>(&mut self, key: Key<T>) -> PushEntry<Key<T>, BufResult<usize, T>> {
         instrument!(compio_log::Level::DEBUG, "pop", ?key);
         if key.has_result() {
             self.cancel.remove(&key);
-            PushEntry::Ready(key.take_result())
+            let (res, buf) = key.take_result().into_parts();
+            
+            PushEntry::Ready(BufResult(resume_unwind_io(res), buf))
         } else {
             PushEntry::Pending(key)
         }
@@ -292,7 +296,8 @@ impl Proactor {
     ///
     /// # Panics
     ///
-    /// This function will panic if the [`Key`] is not unique.
+    /// This function will panic if the [`Key`] is not unique or if the
+    /// operation is blocking and it panicked in the thread pool.
     pub fn pop_with_extra<T: OpCode>(
         &mut self,
         key: Key<T>,
@@ -301,8 +306,8 @@ impl Proactor {
         if key.has_result() {
             self.cancel.remove(&key);
             let extra = key.swap_extra(self.default_extra());
-            let res = key.take_result();
-            PushEntry::Ready((res, extra))
+            let (res, buf) = key.take_result().into_parts();
+            PushEntry::Ready((BufResult(resume_unwind_io(res), buf), extra))
         } else {
             PushEntry::Pending(key)
         }

compio-driver/src/panic.rs

@@ -0,0 +1,47 @@
+use std::{
+    any::Any,
+    error::Error,
+    fmt::{Debug, Display},
+    io,
+    panic::{UnwindSafe, catch_unwind, resume_unwind},
+};
+
+pub(crate) struct Panic(Box<dyn Any + Send>);
+
+// Panic is unconditionally `Sync`
+unsafe impl Sync for Panic {}
+
+impl Display for Panic {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        Debug::fmt(self, f)
+    }
+}
+
+impl Debug for Panic {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_tuple("Panic").finish_non_exhaustive()
+    }
+}
+
+impl Error for Panic {}
+
+impl From<Panic> for io::Error {
+    fn from(value: Panic) -> Self {
+        Self::other(value)
+    }
+}
+
+pub(crate) fn catch_unwind_io<F, R>(f: F) -> io::Result<R>
+where
+    F: FnOnce() -> io::Result<R> + UnwindSafe,
+{
+    catch_unwind(f).map_err(|err| io::Error::from(Panic(err)))?
+}
+
+pub(crate) fn resume_unwind_io<T>(res: io::Result<T>) -> io::Result<T> {
+    let Err(e) = res else { return res };
+    match e.downcast::<Panic>() {
+        Ok(p) => resume_unwind(p.0),
+        Err(e) => Err(e),
+    }
+}

compio-driver/src/sys/iocp/mod.rs

@@ -6,6 +6,7 @@ use std::{
         AsHandle, AsRawHandle, AsRawSocket, AsSocket, BorrowedHandle, BorrowedSocket, OwnedHandle,
         OwnedSocket,
     },
+    panic::AssertUnwindSafe,
     sync::Arc,
     task::{Poll, Wake, Waker},
     time::Duration,
@@ -18,7 +19,10 @@ use windows_sys::Win32::{
     System::IO::OVERLAPPED,
 };
 
-use crate::{AsyncifyPool, DriverType, Entry, ErasedKey, ProactorBuilder, control::Carrier};
+use crate::{
+    AsyncifyPool, DriverType, Entry, ErasedKey, ProactorBuilder, control::Carrier,
+    panic::catch_unwind_io,
+};
 
 pub(crate) mod op;
 
@@ -483,7 +487,7 @@ impl Driver {
 
         let mut closure = move || {
             let op = key.as_mut();
-            let res = op.operate_blocking();
+            let res = catch_unwind_io(AssertUnwindSafe(|| op.operate_blocking()));
             let optr = op.extra_mut().optr();
             // key will be unfronzen in `create_entry` when the result is ready
             notify.port.post(res, optr).ok();
@@ -524,7 +528,7 @@ impl Driver {
             entry
         } else {
             let key = entry.into_key();
-            let result = key.borrow().operate_blocking();
+            let result = catch_unwind_io(AssertUnwindSafe(|| key.borrow().operate_blocking()));
             Entry::new(key, result)
         };
 

compio-driver/src/sys/iour/mod.rs

@@ -6,6 +6,7 @@ use std::{
     io,
     marker::PhantomData,
     os::fd::FromRawFd,
+    panic::AssertUnwindSafe,
     sync::Arc,
     task::{Poll, Wake, Waker},
     time::Duration,
@@ -39,6 +40,7 @@ use crate::{
     AsyncifyPool, DriverType, Entry, ProactorBuilder,
     control::Carrier,
     key::{BorrowedKey, ErasedKey},
+    panic::catch_unwind_io,
     syscall,
 };
 
@@ -544,7 +546,7 @@ impl Driver {
         // SAFETY: we're submitting into the driver, so it's safe to freeze here.
         let mut key = unsafe { key.freeze() };
         let mut closure = move || {
-            let res = key.as_mut().carrier.call_blocking();
+            let res = catch_unwind_io(AssertUnwindSafe(|| key.as_mut().carrier.call_blocking()));
             let _ = completed.send(Entry::new(key.into_inner(), res));
             waker.wake();
         };

compio-driver/src/sys/poll/mod.rs

@@ -7,6 +7,7 @@ use std::{
     collections::{HashMap, VecDeque},
     io,
     num::NonZeroUsize,
+    panic::AssertUnwindSafe,
     sync::Arc,
     task::{Poll, Wake, Waker},
     time::Duration,
@@ -20,7 +21,7 @@ use smallvec::SmallVec;
 
 use crate::{
     AsyncifyPool, DriverType, Entry, ErasedKey, ProactorBuilder, control::Carrier,
-    key::BorrowedKey, sys::op::Interest, syscall,
+    key::BorrowedKey, panic::catch_unwind_io, sys::op::Interest, syscall,
 };
 
 mod extra;
@@ -588,11 +589,11 @@ impl Driver {
         let mut key = unsafe { key.freeze() };
 
         let mut closure = move || {
-            let poll = key.as_mut().carrier.operate();
-            let res = match poll {
+            let operate = || match key.as_mut().carrier.operate() {
                 Poll::Pending => unreachable!("this operation is not non-blocking"),
                 Poll::Ready(res) => res,
             };
+            let res = catch_unwind_io(AssertUnwindSafe(operate));
             let _ = completed.send(Entry::new(key.into_inner(), res));
             waker.wake();
         };

compio-driver/tests/asyncify.rs

@@ -1,3 +1,8 @@
+use std::{
+    panic::{AssertUnwindSafe, catch_unwind},
+    time::Duration,
+};
+
 use compio_buf::BufResult;
 use compio_driver::{Key, OpCode, Proactor, PushEntry, op::Asyncify};
 
@@ -12,14 +17,33 @@ fn take_key<T: OpCode, R>(res: PushEntry<Key<T>, R>) -> Key<T> {
 
 fn wait_for<T: OpCode>(driver: &mut Proactor, mut key: Key<T>) -> BufResult<usize, T> {
     loop {
-        driver.poll(None).unwrap();
+        _ = driver.poll(Some(Duration::from_millis(1)));
         match driver.pop(key) {
             PushEntry::Pending(k) => key = k,
             PushEntry::Ready(res) => break res,
         }
     }
 }
 
+#[test]
+fn panic() {
+    let mut driver = Proactor::builder().thread_pool_limit(1).build().unwrap();
+
+    // make panicking less noisy
+    std::panic::set_hook(Box::new(|_| {}));
+
+    let a = take_key(driver.push(Asyncify::new(|| -> BufResult<usize, ()> {
+        panic!("this should not blow up driver's thread pool");
+    })));
+    let b = take_key(driver.push(Asyncify::new(|| BufResult(Ok(1), ()))));
+
+    let res = wait_for(&mut driver, b);
+    assert!(res.0.is_ok_and(|res| res == 1));
+
+    let res = catch_unwind(AssertUnwindSafe(|| wait_for(&mut driver, a)));
+    assert!(res.is_err());
+}
+
 #[test]
 #[should_panic(expected = "the thread pool is needed but no worker thread is running")]
 fn disable() {