Merge pull request #167 from Berrysoft/dev/rawop-drop

feat(driver): drop completed `RawOp` on `Proactor` dropping

Author: George-Miao

compio-driver/src/fusion/mod.rs

@@ -17,7 +17,7 @@ pub use poll::{Decision, OpCode as PollOpCode};
 use slab::Slab;
 
 pub(crate) use crate::unix::RawOp;
-use crate::{Entry, ProactorBuilder};
+use crate::{Entry, OutEntries, ProactorBuilder};
 
 mod driver_type {
     use std::sync::atomic::{AtomicU8, Ordering};
@@ -153,12 +153,11 @@ impl Driver {
     pub unsafe fn poll(
         &mut self,
         timeout: Option<Duration>,
-        entries: &mut impl Extend<Entry>,
-        registry: &mut Slab<RawOp>,
+        entries: OutEntries<impl Extend<Entry>>,
     ) -> io::Result<()> {
         match &mut self.fuse {
-            FuseDriver::Poll(driver) => driver.poll(timeout, entries, registry),
-            FuseDriver::IoUring(driver) => driver.poll(timeout, entries, registry),
+            FuseDriver::Poll(driver) => driver.poll(timeout, entries),
+            FuseDriver::IoUring(driver) => driver.poll(timeout, entries),
         }
     }
 

compio-driver/src/iocp/mod.rs

@@ -35,7 +35,7 @@ use windows_sys::Win32::{
     },
 };
 
-use crate::{syscall, AsyncifyPool, Entry, ProactorBuilder};
+use crate::{syscall, AsyncifyPool, Entry, OutEntries, ProactorBuilder};
 
 pub(crate) mod op;
 
@@ -330,8 +330,7 @@ impl Driver {
     pub unsafe fn poll(
         &mut self,
         timeout: Option<Duration>,
-        entries: &mut impl Extend<Entry>,
-        _registry: &mut Slab<RawOp>,
+        mut entries: OutEntries<impl Extend<Entry>>,
     ) -> io::Result<()> {
         instrument!(compio_log::Level::TRACE, "poll", ?timeout);
         // Prevent stack growth.
@@ -430,28 +429,56 @@ impl<T> Overlapped<T> {
     }
 }
 
-pub(crate) struct RawOp(NonNull<Overlapped<dyn OpCode>>);
+pub(crate) struct RawOp {
+    op: NonNull<Overlapped<dyn OpCode>>,
+    // The two flags here are manual reference counting. The driver holds the strong ref until it
+    // completes; the runtime holds the strong ref until the future is dropped.
+    completed: bool,
+    cancelled: bool,
+}
 
 impl RawOp {
     pub(crate) fn new(user_data: usize, op: impl OpCode + 'static) -> Self {
         let op = Overlapped::new(user_data, op);
         let op = Box::new(op) as Box<Overlapped<dyn OpCode>>;
-        Self(unsafe { NonNull::new_unchecked(Box::into_raw(op)) })
+        Self {
+            op: unsafe { NonNull::new_unchecked(Box::into_raw(op)) },
+            completed: false,
+            cancelled: false,
+        }
     }
 
     pub fn as_op_pin(&mut self) -> Pin<&mut dyn OpCode> {
-        unsafe { Pin::new_unchecked(&mut self.0.as_mut().op) }
+        unsafe { Pin::new_unchecked(&mut self.op.as_mut().op) }
     }
 
     pub fn as_mut_ptr(&mut self) -> *mut Overlapped<dyn OpCode> {
-        self.0.as_ptr()
+        self.op.as_ptr()
+    }
+
+    pub fn set_completed(&mut self) -> bool {
+        self.completed = true;
+        self.cancelled
+    }
+
+    pub fn set_cancelled(&mut self) -> bool {
+        self.cancelled = true;
+        self.completed
     }
 
     /// # Safety
     /// The caller should ensure the correct type.
     pub unsafe fn into_inner<T: OpCode>(self) -> T {
         let this = ManuallyDrop::new(self);
-        let this: Box<Overlapped<T>> = Box::from_raw(this.0.cast().as_ptr());
+        let this: Box<Overlapped<T>> = Box::from_raw(this.op.cast().as_ptr());
         this.op
     }
 }
+
+impl Drop for RawOp {
+    fn drop(&mut self) {
+        if self.completed {
+            let _ = unsafe { Box::from_raw(self.op.as_ptr()) };
+        }
+    }
+}

compio-driver/src/iour/mod.rs

@@ -18,7 +18,7 @@ use io_uring::{
 pub(crate) use libc::{sockaddr_storage, socklen_t};
 use slab::Slab;
 
-use crate::{syscall, AsyncifyPool, Entry, ProactorBuilder};
+use crate::{syscall, AsyncifyPool, Entry, OutEntries, ProactorBuilder};
 
 pub(crate) mod op;
 pub(crate) use crate::unix::RawOp;
@@ -217,8 +217,7 @@ impl Driver {
     pub unsafe fn poll(
         &mut self,
         timeout: Option<Duration>,
-        entries: &mut impl Extend<Entry>,
-        _registry: &mut Slab<RawOp>,
+        mut entries: OutEntries<impl Extend<Entry>>,
     ) -> io::Result<()> {
         instrument!(compio_log::Level::TRACE, "poll", ?timeout);
         if !self.notifier_registered {
@@ -239,7 +238,7 @@ impl Driver {
 
             self.submit_auto(timeout, ended)?;
 
-            self.poll_entries(entries);
+            self.poll_entries(&mut entries);
 
             if ended {
                 trace!("polling ended");

compio-driver/src/lib.rs

@@ -197,9 +197,18 @@ impl Proactor {
     /// but just don't return from [`Proactor::poll`]. Therefore, although an
     /// operation is cancelled, you should not reuse its `user_data`.
     ///
-    /// It is well-defined to cancel before polling. If the submitted operation
+    /// It is *safe* to cancel before polling. If the submitted operation
     /// contains a cancelled user-defined data, the operation will be ignored.
+    /// However, to make the operation dropped correctly, you should cancel
+    /// after push.
     pub fn cancel(&mut self, user_data: usize) {
+        if let Some(op) = self.ops.get_mut(user_data) {
+            if op.set_cancelled() {
+                // The op is completed.
+                self.ops.remove(user_data);
+                return;
+            }
+        }
         self.driver.cancel(user_data, &mut self.ops);
     }
 
@@ -227,7 +236,8 @@ impl Proactor {
         entries: &mut impl Extend<Entry>,
     ) -> io::Result<()> {
         unsafe {
-            self.driver.poll(timeout, entries, &mut self.ops)?;
+            self.driver
+                .poll(timeout, OutEntries::new(entries, &mut self.ops))?;
         }
         Ok(())
     }
@@ -320,6 +330,35 @@ impl Entry {
     }
 }
 
+// The output entries need to be marked as `completed`. If an entry has been
+// marked as `cancelled`, it will be removed from the registry.
+struct OutEntries<'a, 'b, E> {
+    entries: &'b mut E,
+    registry: &'a mut Slab<RawOp>,
+}
+
+impl<'a, 'b, E> OutEntries<'a, 'b, E> {
+    pub fn new(entries: &'b mut E, registry: &'a mut Slab<RawOp>) -> Self {
+        Self { entries, registry }
+    }
+
+    #[allow(dead_code)]
+    pub fn registry(&mut self) -> &mut Slab<RawOp> {
+        self.registry
+    }
+}
+
+impl<E: Extend<Entry>> Extend<Entry> for OutEntries<'_, '_, E> {
+    fn extend<T: IntoIterator<Item = Entry>>(&mut self, iter: T) {
+        self.entries.extend(iter.into_iter().map(|e| {
+            if self.registry[e.user_data()].set_completed() {
+                self.registry.remove(e.user_data());
+            }
+            e
+        }))
+    }
+}
+
 #[derive(Debug, Clone)]
 enum ThreadPoolBuilder {
     Create { limit: usize, recv_limit: Duration },

compio-driver/src/poll/mod.rs

@@ -19,7 +19,7 @@ pub(crate) use libc::{sockaddr_storage, socklen_t};
 use polling::{Event, Events, Poller};
 use slab::Slab;
 
-use crate::{syscall, AsyncifyPool, Entry, ProactorBuilder};
+use crate::{syscall, AsyncifyPool, Entry, OutEntries, ProactorBuilder};
 
 pub(crate) mod op;
 
@@ -268,8 +268,7 @@ impl Driver {
     pub unsafe fn poll(
         &mut self,
         timeout: Option<Duration>,
-        entries: &mut impl Extend<Entry>,
-        registry: &mut Slab<RawOp>,
+        mut entries: OutEntries<impl Extend<Entry>>,
     ) -> io::Result<()> {
         self.poll.wait(&mut self.events, timeout)?;
         if self.events.is_empty() && self.pool_completed.is_empty() && timeout.is_some() {
@@ -288,7 +287,7 @@ impl Driver {
             if self.cancelled.remove(&user_data) {
                 entries.extend(Some(entry_cancelled(user_data)));
             } else {
-                let op = registry[user_data].as_pin();
+                let op = entries.registry()[user_data].as_pin();
                 let res = match op.on_event(&event) {
                     Poll::Pending => {
                         // The operation should go back to the front.

compio-driver/src/unix/mod.rs

@@ -7,22 +7,50 @@ use std::{mem::ManuallyDrop, pin::Pin, ptr::NonNull};
 
 use crate::OpCode;
 
-pub(crate) struct RawOp(NonNull<dyn OpCode>);
+pub(crate) struct RawOp {
+    op: NonNull<dyn OpCode>,
+    // The two flags here are manual reference counting. The driver holds the strong ref until it
+    // completes; the runtime holds the strong ref until the future is dropped.
+    completed: bool,
+    cancelled: bool,
+}
 
 impl RawOp {
     pub(crate) fn new(_user_data: usize, op: impl OpCode + 'static) -> Self {
         let op = Box::new(op);
-        Self(unsafe { NonNull::new_unchecked(Box::into_raw(op as Box<dyn OpCode>)) })
+        Self {
+            op: unsafe { NonNull::new_unchecked(Box::into_raw(op as Box<dyn OpCode>)) },
+            completed: false,
+            cancelled: false,
+        }
     }
 
     pub fn as_pin(&mut self) -> Pin<&mut dyn OpCode> {
-        unsafe { Pin::new_unchecked(self.0.as_mut()) }
+        unsafe { Pin::new_unchecked(self.op.as_mut()) }
+    }
+
+    pub fn set_completed(&mut self) -> bool {
+        self.completed = true;
+        self.cancelled
+    }
+
+    pub fn set_cancelled(&mut self) -> bool {
+        self.cancelled = true;
+        self.completed
     }
 
     /// # Safety
     /// The caller should ensure the correct type.
     pub unsafe fn into_inner<T: OpCode>(self) -> T {
         let this = ManuallyDrop::new(self);
-        *Box::from_raw(this.0.cast().as_ptr())
+        *Box::from_raw(this.op.cast().as_ptr())
+    }
+}
+
+impl Drop for RawOp {
+    fn drop(&mut self) {
+        if self.completed {
+            let _ = unsafe { Box::from_raw(self.op.as_ptr()) };
+        }
     }
 }

compio-runtime/src/runtime/mod.rs

@@ -143,8 +143,8 @@ impl RuntimeInner {
     }
 
     pub fn cancel_op<T>(&self, user_data: Key<T>) {
+        self.op_runtime.borrow_mut().remove(*user_data);
         self.driver.borrow_mut().cancel(*user_data);
-        self.op_runtime.borrow_mut().cancel(*user_data);
     }
 
     #[cfg(feature = "time")]
@@ -165,7 +165,7 @@ impl RuntimeInner {
             let res = self
                 .driver
                 .borrow_mut()
-                .pop(&mut op.entry.into_iter())
+                .pop(&mut op.into_completed().into_iter())
                 .next()
                 .expect("the result should have come");
             Poll::Ready(res.map_buffer(|op| unsafe { op.into_op::<T>() }))
@@ -204,9 +204,7 @@ impl RuntimeInner {
             Ok(_) => {
                 debug!("poll driver ok, entries: {}", entries.len());
                 for entry in entries {
-                    self.op_runtime
-                        .borrow_mut()
-                        .update_result(entry.user_data(), entry);
+                    self.op_runtime.borrow_mut().update_result(entry);
                 }
             }
             Err(e) => match e.kind() {

compio-runtime/src/runtime/op.rs

@@ -11,10 +11,20 @@ use compio_driver::{Entry, OpCode};
 use crate::{key::Key, Runtime};
 
 #[derive(Default)]
-pub(crate) struct RegisteredOp {
-    pub waker: Option<Waker>,
-    pub entry: Option<Entry>,
-    pub cancelled: bool,
+pub(crate) enum RegisteredOp {
+    #[default]
+    Default,
+    Submitted(Waker),
+    Completed(Entry),
+}
+
+impl RegisteredOp {
+    pub fn into_completed(self) -> Option<Entry> {
+        match self {
+            Self::Completed(op) => Some(op),
+            _ => None,
+        }
+    }
 }
 
 #[derive(Default)]
@@ -24,29 +34,24 @@ pub(crate) struct OpRuntime {
 
 impl OpRuntime {
     pub fn update_waker(&mut self, key: usize, waker: Waker) {
-        self.ops.entry(key).or_default().waker = Some(waker);
+        *self.ops.entry(key).or_default() = RegisteredOp::Submitted(waker)
     }
 
-    pub fn update_result(&mut self, key: usize, entry: Entry) {
+    pub fn update_result(&mut self, entry: Entry) {
+        let key = entry.user_data();
         let op = self.ops.entry(key).or_default();
-        if let Some(waker) = op.waker.take() {
-            waker.wake();
-        }
-        op.entry = Some(entry);
-        if op.cancelled {
-            self.remove(key);
+        match op {
+            RegisteredOp::Default => *op = RegisteredOp::Completed(entry),
+            RegisteredOp::Submitted(waker) => {
+                waker.wake_by_ref();
+                *op = RegisteredOp::Completed(entry);
+            }
+            RegisteredOp::Completed(res) => *res = entry,
         }
     }
 
     pub fn has_result(&mut self, key: usize) -> bool {
-        self.ops
-            .get_mut(&key)
-            .map(|op| op.entry.is_some())
-            .unwrap_or_default()
-    }
-
-    pub fn cancel(&mut self, key: usize) {
-        self.ops.entry(key).or_default().cancelled = true;
+        matches!(self.ops.get_mut(&key), Some(RegisteredOp::Completed(_)))
     }
 
     pub fn remove(&mut self, key: usize) -> RegisteredOp {