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/key.rs

@@ -60,14 +60,17 @@ impl<C: crate::Carry + ?Sized> RawOp<C> {
     ///
     /// [`Poll::Pending`]: std::task::Poll::Pending
     pub fn operate_blocking(&mut self) -> io::Result<usize> {
-        use std::task::Poll;
+        use std::{panic::AssertUnwindSafe, task::Poll};
+
+        use crate::panic::catch_unwind_io;
 
         let optr = self.extra_mut().optr();
-        let res = unsafe { self.carrier.operate(optr.cast()) };
-        match res {
-            Poll::Pending => unreachable!("this operation is not overlapped"),
-            Poll::Ready(res) => res,
-        }
+        catch_unwind_io(AssertUnwindSafe(|| unsafe {
+            match self.carrier.operate(optr.cast()) {
+                Poll::Pending => unreachable!("this operation is not overlapped"),
+                Poll::Ready(res) => res,
+            }
+        }))
     }
 }
 

compio-driver/src/lib.rs

@@ -26,6 +26,7 @@ use compio_log::instrument;
 
 mod control;
 mod macros;
+mod panic;
 
 mod key;
 pub use key::Key;
@@ -55,6 +56,7 @@ use crate::{
     buffer_pool::{BufferAlloc, BufferPoolRoot},
     key::ErasedKey,
     op::OpCodeFlag,
+    panic::resume_unwind_io,
 };
 
 mod sys_slice;
@@ -212,7 +214,8 @@ impl Proactor {
         }
         self.cancel.remove(&key);
         if key.is_unique() && key.has_result() {
-            Some(key.take_result())
+            let (res, buf) = key.take_result().into_parts();
+            Some(BufResult(resume_unwind_io(res), buf))
         } else {
             self.driver.cancel(key.erase());
             None
@@ -286,12 +289,14 @@ 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)
         }
@@ -302,7 +307,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>,
@@ -311,8 +317,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/cp/global.rs

@@ -103,6 +103,7 @@ impl Port {
         self.global_port.attach(fd)
     }
 
+    #[allow(dead_code)]
     pub fn post(&self, res: io::Result<usize>, optr: *mut Overlapped) -> io::Result<()> {
         self.port.post(res, optr)
     }

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

@@ -83,6 +83,7 @@ impl CompletionPort {
         Ok(())
     }
 
+    #[allow(dead_code)]
     pub fn post(&self, res: io::Result<usize>, optr: *mut Overlapped) -> io::Result<()> {
         if let Some(overlapped) = unsafe { optr.as_mut() } {
             match &res {

compio-driver/src/sys/iocp/cp/multi.rs

@@ -22,6 +22,7 @@ impl Port {
         self.port.attach(fd)
     }
 
+    #[allow(dead_code)]
     pub fn post(&self, res: io::Result<usize>, optr: *mut Overlapped) -> io::Result<()> {
         self.port.post(res, optr)
     }

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

@@ -13,6 +13,7 @@ use std::{
 
 use compio_buf::BufResult;
 use compio_log::{instrument, trace};
+use flume::{Receiver, Sender};
 use windows_sys::Win32::{
     Foundation::{ERROR_CANCELLED, HANDLE},
     System::IO::OVERLAPPED,
@@ -396,6 +397,8 @@ pub(crate) struct Driver {
     notify: Arc<Notify>,
     waits: HashMap<usize, wait::Wait>,
     pool: AsyncifyPool,
+    completed_tx: Sender<Entry>,
+    completed_rx: Receiver<Entry>,
     _local_marker: PhantomData<ErasedKey>,
 }
 
@@ -407,8 +410,12 @@ impl Driver {
         let driver = port.as_raw_handle() as _;
         let overlapped = Overlapped::new(driver);
         let notify = Arc::new(Notify::new(port, overlapped));
+        let (completed_tx, completed_rx) = flume::unbounded();
+
         Ok(Self {
             notify,
+            completed_tx,
+            completed_rx,
             waits: HashMap::default(),
             pool: builder.create_or_get_thread_pool(),
             _local_marker: PhantomData,
@@ -477,16 +484,14 @@ impl Driver {
     }
 
     fn push_blocking(&mut self, key: ErasedKey) {
-        let notify = self.notify.clone();
         // SAFETY: we're submitting into the driver, so it's safe to freeze here.
         let mut key = unsafe { key.freeze() };
+        let tx = self.completed_tx.clone();
 
         let mut closure = move || {
-            let op = key.as_mut();
-            let res = 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();
+            let res = key.as_mut().operate_blocking();
+            let entry = Entry::new(key.into_inner(), res);
+            _ = tx.send(entry);
         };
 
         while let Err(e) = self.pool.dispatch(closure) {
@@ -536,6 +541,10 @@ impl Driver {
 
         let notify = &self.notify.overlapped as *const Overlapped;
 
+        while let Ok(entry) = self.completed_rx.try_recv() {
+            entry.notify();
+        }
+
         for e in self.notify.port.poll(timeout)? {
             if let Some(e) = Self::create_entry(notify, &mut self.waits, e) {
                 e.notify()

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();
         };