refactor(driver): record multishot results in ops (#748)

Berrysoft · web-flow · commit 3d5dba933ac9 · 2026-03-10T22:18:58.000+08:00
* refactor(driver): record multishot results in ops

* docs(driver,iour): fixme

* docs(driver,iour): grammar

* fix(driver,iour): rename
diff --git a/compio-driver/src/buffer_pool/fallback.rs b/compio-driver/src/buffer_pool/fallback.rs
@@ -15,6 +15,8 @@ use std::{
 };
 
 use compio_buf::{IntoInner, IoBuf, IoBufMut, SetLen, Slice};
+#[cfg(not(fusion))]
+pub use {BufferPool as FallbackBufferPool, OwnedBuffer as FallbackOwnedBuffer};
 
 struct BufferPoolInner {
     buffers: RefCell<VecDeque<Vec<u8>>>,
diff --git a/compio-driver/src/buffer_pool/fusion.rs b/compio-driver/src/buffer_pool/fusion.rs
@@ -5,8 +5,9 @@ use std::{
 };
 
 pub use fallback::BufferPool as FallbackBufferPool;
-pub(crate) use fallback::OwnedBuffer;
+pub(crate) use fallback::OwnedBuffer as FallbackOwnedBuffer;
 pub use iour::BufferPool as IoUringBufferPool;
+pub(crate) use iour::OwnedBuffer as IoUringOwnedBuffer;
 
 use super::{fallback, iour};
 
diff --git a/compio-driver/src/buffer_pool/iour.rs b/compio-driver/src/buffer_pool/iour.rs
@@ -7,16 +7,32 @@ use std::{
     fmt::{Debug, Formatter},
     io,
     ops::{Deref, DerefMut},
+    rc::Rc,
 };
 
 use io_uring_buf_ring::IoUringBufRing;
+#[cfg(not(fusion))]
+pub use {BufferPool as IoUringBufferPool, OwnedBuffer as IoUringOwnedBuffer};
+
+struct BufferPoolInner {
+    buf_ring: IoUringBufRing<Vec<u8>>,
+}
+
+impl BufferPoolInner {
+    fn reuse_buffer(&self, buffer_id: u16) {
+        // SAFETY: 0 is always valid length. We just want to get the buffer once and
+        // return it immediately.
+        unsafe { self.buf_ring.get_buf(buffer_id, 0) };
+    }
+}
 
 /// Buffer pool
 ///
 /// A buffer pool to allow user no need to specify a specific buffer to do the
 /// IO operation
+#[derive(Clone)]
 pub struct BufferPool {
-    buf_ring: IoUringBufRing<Vec<u8>>,
+    inner: Rc<BufferPoolInner>,
 }
 
 impl Debug for BufferPool {
@@ -27,33 +43,67 @@ impl Debug for BufferPool {
 
 impl BufferPool {
     pub(crate) fn new(buf_ring: IoUringBufRing<Vec<u8>>) -> Self {
-        Self { buf_ring }
+        Self {
+            inner: Rc::new(BufferPoolInner { buf_ring }),
+        }
     }
 
     pub(crate) fn buffer_group(&self) -> u16 {
-        self.buf_ring.buffer_group()
+        self.inner.buf_ring.buffer_group()
+    }
+
+    pub(crate) fn into_inner(self) -> Result<IoUringBufRing<Vec<u8>>, Self> {
+        Rc::try_unwrap(self.inner)
+            .map(|inner| inner.buf_ring)
+            .map_err(|inner| Self { inner })
     }
 
-    pub(crate) fn into_inner(self) -> IoUringBufRing<Vec<u8>> {
-        self.buf_ring
+    #[doc(hidden)]
+    pub unsafe fn get_buffer(&self, buffer_id: u16, available_len: usize) -> OwnedBuffer {
+        OwnedBuffer {
+            pool: self.inner.clone(),
+            params: Some((buffer_id, available_len)),
+        }
     }
 
     /// ## Safety
-    /// * `available_len` should be the returned value from the op.
-    pub(crate) unsafe fn get_buffer(
+    /// * `len` should be the returned value from the op.
+    pub(crate) unsafe fn create_proxy(
         &self,
-        buffer_id: u16,
-        available_len: usize,
+        slice: OwnedBuffer,
+        len: usize,
     ) -> io::Result<BorrowedBuffer<'_>> {
-        unsafe { self.buf_ring.get_buf(buffer_id, available_len) }
+        let Some((buffer_id, available_len)) = slice.leak() else {
+            return Err(io::Error::other("no buffer selected"));
+        };
+        debug_assert_eq!(available_len, len);
+        unsafe { self.inner.buf_ring.get_buf(buffer_id, available_len) }
             .map(BorrowedBuffer)
             .ok_or_else(|| io::Error::other(format!("cannot find buffer {buffer_id}")))
     }
 
     pub(crate) fn reuse_buffer(&self, buffer_id: u16) {
-        // SAFETY: 0 is always valid length. We just want to get the buffer once and
-        // return it immediately.
-        unsafe { self.buf_ring.get_buf(buffer_id, 0) };
+        self.inner.reuse_buffer(buffer_id);
+    }
+}
+
+#[doc(hidden)]
+pub struct OwnedBuffer {
+    pool: Rc<BufferPoolInner>,
+    params: Option<(u16, usize)>,
+}
+
+impl OwnedBuffer {
+    pub fn leak(mut self) -> Option<(u16, usize)> {
+        self.params.take()
+    }
+}
+
+impl Drop for OwnedBuffer {
+    fn drop(&mut self) {
+        if let Some((buffer_id, _)) = self.params {
+            self.pool.reuse_buffer(buffer_id);
+        }
     }
 }
 
diff --git a/compio-driver/src/key.rs b/compio-driver/src/key.rs
@@ -279,11 +279,12 @@ impl ErasedKey {
     pub(crate) fn set_result(&self, res: io::Result<usize>) {
         let mut this = self.borrow();
         #[cfg(io_uring)]
-        if let Ok(res) = res
-            && this.extra.is_iour()
         {
-            unsafe {
-                Pin::new_unchecked(&mut this.op).set_result(res);
+            let this = &mut *this;
+            if this.extra.is_iour() {
+                unsafe {
+                    Pin::new_unchecked(&mut this.op).set_result(&res, &this.extra);
+                }
             }
         }
         if let PushEntry::Pending(Some(w)) =
diff --git a/compio-driver/src/lib.rs b/compio-driver/src/lib.rs
@@ -233,7 +233,6 @@ impl Proactor {
         instrument!(compio_log::Level::DEBUG, "pop", ?key);
         if key.has_result() {
             self.cancel.remove(&key);
-            self.driver.cleanup_multishot(&key);
             PushEntry::Ready(key.take_result())
         } else {
             PushEntry::Pending(key)
@@ -253,7 +252,6 @@ impl Proactor {
         instrument!(compio_log::Level::DEBUG, "pop", ?key);
         if key.has_result() {
             self.cancel.remove(&key);
-            self.driver.cleanup_multishot(&key);
             let extra = key.swap_extra(self.default_extra());
             let res = key.take_result();
             PushEntry::Ready((res, extra))
diff --git a/compio-driver/src/op.rs b/compio-driver/src/op.rs
@@ -495,10 +495,10 @@ pub(crate) mod managed {
     use socket2::SockAddr;
 
     use super::{Read, ReadAt, Recv, RecvFrom};
-    use crate::{AsFd, BorrowedBuffer, BufferPool, OwnedBuffer, TakeBuffer};
+    use crate::{AsFd, BorrowedBuffer, BufferPool, FallbackOwnedBuffer, TakeBuffer};
 
     fn take_buffer(
-        slice: OwnedBuffer,
+        slice: FallbackOwnedBuffer,
         buffer_pool: &BufferPool,
         result: io::Result<usize>,
     ) -> io::Result<BorrowedBuffer<'_>> {
@@ -516,7 +516,7 @@ pub(crate) mod managed {
         /// Read a file at specified position into managed buffer.
         pub struct ReadManagedAt<S> {
             #[pin]
-            pub(crate) op: ReadAt<OwnedBuffer, S>,
+            pub(crate) op: ReadAt<FallbackOwnedBuffer, S>,
         }
     }
 
@@ -549,7 +549,7 @@ pub(crate) mod managed {
         /// Read a file into managed buffer.
         pub struct ReadManaged<S> {
             #[pin]
-            pub(crate) op: Read<OwnedBuffer, S>,
+            pub(crate) op: Read<FallbackOwnedBuffer, S>,
         }
     }
 
@@ -585,7 +585,7 @@ pub(crate) mod managed {
         /// use [`ReadManaged`].
         pub struct RecvManaged<S> {
             #[pin]
-            pub(crate) op: Recv<OwnedBuffer, S>,
+            pub(crate) op: Recv<FallbackOwnedBuffer, S>,
         }
     }
 
@@ -618,7 +618,7 @@ pub(crate) mod managed {
         /// Receive data and source address into managed buffer.
         pub struct RecvFromManaged<S: AsFd> {
             #[pin]
-            pub(crate) op: RecvFrom<OwnedBuffer, S>,
+            pub(crate) op: RecvFrom<FallbackOwnedBuffer, S>,
         }
     }
 
diff --git a/compio-driver/src/sys/fusion/mod.rs b/compio-driver/src/sys/fusion/mod.rs
@@ -157,13 +157,6 @@ impl Driver {
             FuseDriver::IoUring(driver) => driver.pop_multishot(key),
         }
     }
-
-    pub fn cleanup_multishot(&mut self, key: &ErasedKey) {
-        match &mut self.fuse {
-            FuseDriver::Poll(driver) => driver.cleanup_multishot(key),
-            FuseDriver::IoUring(driver) => driver.cleanup_multishot(key),
-        }
-    }
 }
 
 impl AsRawFd for Driver {
diff --git a/compio-driver/src/sys/fusion/op.rs b/compio-driver/src/sys/fusion/op.rs
@@ -97,6 +97,26 @@ macro_rules! op {
             fn create_entry(self: std::pin::Pin<&mut Self>) -> OpEntry {
                 unsafe { self.map_unchecked_mut(|x| x.inner.iour() ) }.create_entry()
             }
+
+            fn create_entry_fallback(self: std::pin::Pin<&mut Self>) -> OpEntry {
+                unsafe { self.map_unchecked_mut(|x| x.inner.iour() ) }.create_entry_fallback()
+            }
+
+            fn call_blocking(self: std::pin::Pin<&mut Self>) -> std::io::Result<usize> {
+                unsafe { self.map_unchecked_mut(|x| x.inner.iour() ) }.call_blocking()
+            }
+
+            unsafe fn set_result(self: std::pin::Pin<&mut Self>, result: &std::io::Result<usize>, extra: &crate::Extra) {
+                unsafe { self.map_unchecked_mut(|x| x.inner.iour() ).set_result(result, extra) }
+            }
+
+            unsafe fn push_multishot(self: std::pin::Pin<&mut Self>, result: std::io::Result<usize>, extra: crate::Extra) {
+                unsafe { self.map_unchecked_mut(|x| x.inner.iour() ).push_multishot(result, extra) }
+            }
+
+            fn pop_multishot(self: std::pin::Pin<&mut Self>) -> Option<BufResult<usize, crate::Extra>> {
+                unsafe { self.map_unchecked_mut(|x| x.inner.iour() ) }.pop_multishot()
+            }
         }
     };
 }
@@ -202,6 +222,26 @@ macro_rules! mop {
             fn create_entry(self: std::pin::Pin<&mut Self>) -> OpEntry {
                 unsafe { self.map_unchecked_mut(|x| x.inner.iour() ) }.create_entry()
             }
+
+            fn create_entry_fallback(self: std::pin::Pin<&mut Self>) -> OpEntry {
+                unsafe { self.map_unchecked_mut(|x| x.inner.iour() ) }.create_entry_fallback()
+            }
+
+            fn call_blocking(self: std::pin::Pin<&mut Self>) -> std::io::Result<usize> {
+                unsafe { self.map_unchecked_mut(|x| x.inner.iour() ) }.call_blocking()
+            }
+
+            unsafe fn set_result(self: std::pin::Pin<&mut Self>, result: &std::io::Result<usize>, extra: &crate::Extra) {
+                unsafe { self.map_unchecked_mut(|x| x.inner.iour() ).set_result(result, extra) }
+            }
+
+            unsafe fn push_multishot(self: std::pin::Pin<&mut Self>, result: std::io::Result<usize>, extra: crate::Extra) {
+                unsafe { self.map_unchecked_mut(|x| x.inner.iour() ).push_multishot(result, extra) }
+            }
+
+            fn pop_multishot(self: std::pin::Pin<&mut Self>) -> Option<BufResult<usize, crate::Extra>> {
+                unsafe { self.map_unchecked_mut(|x| x.inner.iour() ) }.pop_multishot()
+            }
         }
     };
 }
diff --git a/compio-driver/src/sys/iocp/mod.rs b/compio-driver/src/sys/iocp/mod.rs
@@ -502,8 +502,6 @@ impl Driver {
     pub fn pop_multishot(&mut self, _: &ErasedKey) -> Option<BufResult<usize, crate::sys::Extra>> {
         None
     }
-
-    pub fn cleanup_multishot(&mut self, _: &ErasedKey) {}
 }
 
 impl AsRawFd for Driver {
diff --git a/compio-driver/src/sys/iour/mod.rs b/compio-driver/src/sys/iour/mod.rs
diff --git a/compio-driver/src/sys/iour/op.rs b/compio-driver/src/sys/iour/op.rs
diff --git a/compio-driver/src/sys/poll/mod.rs b/compio-driver/src/sys/poll/mod.rs
diff --git a/compio-driver/src/sys/stub/mod.rs b/compio-driver/src/sys/stub/mod.rs

Original file line number	Diff line number	Diff line change
`@@ -279,11 +279,12 @@ impl ErasedKey {`
`279`	`279`	`pub(crate) fn set_result(&self, res: io::Result<usize>) {`
`280`	`280`	`let mut this = self.borrow();`
`281`	`281`	`#[cfg(io_uring)]`
`282`		`- if let Ok(res) = res`
`283`		`- && this.extra.is_iour()`
`284`	`282`	`{`
`285`		`- unsafe {`
`286`		`- Pin::new_unchecked(&mut this.op).set_result(res);`
	`283`	`+ let this = &mut *this;`
	`284`	`+ if this.extra.is_iour() {`
	`285`	`+ unsafe {`
	`286`	`+ Pin::new_unchecked(&mut this.op).set_result(&res, &this.extra);`
	`287`	`+ }`
`287`	`288`	`}`
`288`	`289`	`}`
`289`	`290`	`if let PushEntry::Pending(Some(w)) =`
Original file line number	Diff line number	Diff line change
`@@ -495,10 +495,10 @@ pub(crate) mod managed {`
`495`	`495`	`use socket2::SockAddr;`
`496`	`496`
`497`	`497`	`use super::{Read, ReadAt, Recv, RecvFrom};`
`498`		`- use crate::{AsFd, BorrowedBuffer, BufferPool, OwnedBuffer, TakeBuffer};`
	`498`	`+ use crate::{AsFd, BorrowedBuffer, BufferPool, FallbackOwnedBuffer, TakeBuffer};`
`499`	`499`
`500`	`500`	`fn take_buffer(`
`501`		`- slice: OwnedBuffer,`
	`501`	`+ slice: FallbackOwnedBuffer,`
`502`	`502`	`buffer_pool: &BufferPool,`
`503`	`503`	`result: io::Result<usize>,`
`504`	`504`	`) -> io::Result<BorrowedBuffer<'_>> {`
`@@ -516,7 +516,7 @@ pub(crate) mod managed {`
`516`	`516`	`/// Read a file at specified position into managed buffer.`
`517`	`517`	`pub struct ReadManagedAt<S> {`
`518`	`518`	`#[pin]`
`519`		`- pub(crate) op: ReadAt<OwnedBuffer, S>,`
	`519`	`+ pub(crate) op: ReadAt<FallbackOwnedBuffer, S>,`
`520`	`520`	`}`
`521`	`521`	`}`
`522`	`522`
`@@ -549,7 +549,7 @@ pub(crate) mod managed {`
`549`	`549`	`/// Read a file into managed buffer.`
`550`	`550`	`pub struct ReadManaged<S> {`
`551`	`551`	`#[pin]`
`552`		`- pub(crate) op: Read<OwnedBuffer, S>,`
	`552`	`+ pub(crate) op: Read<FallbackOwnedBuffer, S>,`
`553`	`553`	`}`
`554`	`554`	`}`
`555`	`555`
`@@ -585,7 +585,7 @@ pub(crate) mod managed {`
`585`	`585`	/// use [`ReadManaged`].
`586`	`586`	`pub struct RecvManaged<S> {`
`587`	`587`	`#[pin]`
`588`		`- pub(crate) op: Recv<OwnedBuffer, S>,`
	`588`	`+ pub(crate) op: Recv<FallbackOwnedBuffer, S>,`
`589`	`589`	`}`
`590`	`590`	`}`
`591`	`591`
`@@ -618,7 +618,7 @@ pub(crate) mod managed {`
`618`	`618`	`/// Receive data and source address into managed buffer.`
`619`	`619`	`pub struct RecvFromManaged<S: AsFd> {`
`620`	`620`	`#[pin]`
`621`		`- pub(crate) op: RecvFrom<OwnedBuffer, S>,`
	`621`	`+ pub(crate) op: RecvFrom<FallbackOwnedBuffer, S>,`
`622`	`622`	`}`
`623`	`623`	`}`
`624`	`624`
Original file line number	Diff line number	Diff line change
`@@ -157,13 +157,6 @@ impl Driver {`
`157`	`157`	`FuseDriver::IoUring(driver) => driver.pop_multishot(key),`
`158`	`158`	`}`
`159`	`159`	`}`
`160`		`-`
`161`		`- pub fn cleanup_multishot(&mut self, key: &ErasedKey) {`
`162`		`- match &mut self.fuse {`
`163`		`- FuseDriver::Poll(driver) => driver.cleanup_multishot(key),`
`164`		`- FuseDriver::IoUring(driver) => driver.cleanup_multishot(key),`
`165`		`- }`
`166`		`- }`
`167`	`160`	`}`
`168`	`161`
`169`	`162`	`impl AsRawFd for Driver {`
Original file line number	Diff line number	Diff line change
`@@ -502,8 +502,6 @@ impl Driver {`
`502`	`502`	`pub fn pop_multishot(&mut self, _: &ErasedKey) -> Option<BufResult<usize, crate::sys::Extra>> {`
`503`	`503`	`None`
`504`	`504`	`}`
`505`		`-`
`506`		`- pub fn cleanup_multishot(&mut self, _: &ErasedKey) {}`
`507`	`505`	`}`
`508`	`506`
`509`	`507`	`impl AsRawFd for Driver {`