feat: add delay RAII release for BufferPool

Author: Sherlock-Holo

compio-driver/src/buffer_pool.rs

@@ -1,5 +1,6 @@
 use std::{
-    cell::UnsafeCell,
+    cell::{RefCell, UnsafeCell},
+    collections::VecDeque,
     fmt::Debug,
     io,
     mem::{self, MaybeUninit},
@@ -123,6 +124,15 @@ struct Inner {
 
     /// Buffer pointers
     bufs: Vec<Slot>,
+
+    /// Queue used by extra buffer pools to ask the proactor to release them.
+    release_queue: Option<ReleaseQueue>,
+}
+
+#[derive(Debug)]
+struct ReleaseQueue {
+    buffer_group: u16,
+    queue: Weak<RefCell<VecDeque<u16>>>,
 }
 
 impl BufferPoolRoot {
@@ -132,6 +142,8 @@ impl BufferPoolRoot {
         num_of_bufs: u16,
         buffer_size: usize,
         flags: u16,
+        buffer_group: u16,
+        release_queue: Option<Weak<RefCell<VecDeque<u16>>>>,
     ) -> io::Result<Self> {
         let size: u32 = buffer_size.try_into().map_err(|_| {
             io::Error::new(
@@ -142,7 +154,11 @@ impl BufferPoolRoot {
         let bufs = (0..num_of_bufs.next_power_of_two())
             .map(|_| Some((alloc.allocate)(size)))
             .collect::<Vec<_>>();
-        let ctrl = unsafe { BufControl::new(driver, &bufs, size, flags) }?;
+        let ctrl = unsafe { BufControl::new(driver, &bufs, size, flags, buffer_group) }?;
+        let release_queue = release_queue.map(|queue| ReleaseQueue {
+            buffer_group,
+            queue,
+        });
 
         Ok(Self {
             shared: Shared {
@@ -151,6 +167,7 @@ impl BufferPoolRoot {
                     ctrl,
                     size,
                     bufs,
+                    release_queue,
                 }
                 .into(),
             }
@@ -361,6 +378,24 @@ impl Shared {
     fn len(&self) -> u32 {
         unsafe { self.with(|inner| inner.size) }
     }
+
+    fn queue_release_if_unused(self: &Rc<Self>) {
+        if Rc::weak_count(self) != 1 {
+            return;
+        }
+
+        unsafe {
+            self.with(|inner| {
+                let Some(release_queue) = &inner.release_queue else {
+                    return;
+                };
+
+                if let Some(queue) = release_queue.queue.upgrade() {
+                    queue.borrow_mut().push_back(release_queue.buffer_group);
+                }
+            })
+        }
+    }
 }
 
 impl BufferRef {
@@ -429,8 +464,17 @@ impl Drop for BufferRef {
         if let Some(shared) = self.shared.upgrade() {
             // If the buffer pool is alive, set the pointer back
             shared.reset(self.buffer_id, self.ptr);
+            shared.queue_release_if_unused();
         } else {
             unsafe { (self.alloc.deallocate)(self.ptr, self.full_cap) }
         }
     }
 }
+
+impl Drop for BufferPool {
+    fn drop(&mut self) {
+        if let Some(shared) = self.shared.upgrade() {
+            shared.queue_release_if_unused();
+        }
+    }
+}

compio-driver/src/lib.rs

@@ -15,8 +15,11 @@
 )]
 
 use std::{
+    cell::RefCell,
+    collections::{HashMap, VecDeque},
     io,
     num::NonZero,
+    rc::Rc,
     task::{Poll, Waker},
     time::Duration,
 };
@@ -101,7 +104,10 @@ impl<K, R> PushEntry<K, R> {
 pub struct Proactor {
     driver: Driver,
     buffer_pool: BufferPoolState,
-    extra_pools: Vec<BufferPoolRoot>,
+    extra_pools: HashMap<u16, BufferPoolRoot>,
+    pending_extra_pool_releases: Rc<RefCell<VecDeque<u16>>>,
+    free_extra_bgids: VecDeque<u16>,
+    next_extra_bgid: u16,
 }
 
 enum BufferPoolState {
@@ -115,6 +121,8 @@ enum BufferPoolState {
 }
 
 impl BufferPoolState {
+    const DEFAULT_BUFFER_GROUP: u16 = 1;
+
     fn get(&mut self, driver: &mut Driver) -> io::Result<BufferPool> {
         loop {
             match self {
@@ -130,6 +138,8 @@ impl BufferPoolState {
                         *num_of_bufs,
                         *buffer_len,
                         *flags,
+                        Self::DEFAULT_BUFFER_GROUP,
+                        None,
                     )?);
                 }
                 BufferPoolState::Init(root) => return Ok(root.get_pool()),
@@ -143,7 +153,7 @@ impl Drop for Proactor {
         if let BufferPoolState::Init(buffer_pool) = &mut self.buffer_pool {
             _ = unsafe { buffer_pool.release(&mut self.driver) };
         }
-        for pool in &mut self.extra_pools {
+        for (_, mut pool) in self.extra_pools.drain() {
             _ = unsafe { pool.release(&mut self.driver) };
         }
     }
@@ -172,10 +182,45 @@ impl Proactor {
                 buffer_len: builder.buffer_pool_buffer_len,
                 flags: builder.buffer_pool_flag,
             },
-            extra_pools: Vec::new(),
+            extra_pools: HashMap::new(),
+            pending_extra_pool_releases: Rc::default(),
+            free_extra_bgids: VecDeque::new(),
+            next_extra_bgid: BufferPoolState::DEFAULT_BUFFER_GROUP + 1,
         })
     }
 
+    fn alloc_extra_bgid(&mut self) -> io::Result<u16> {
+        if let Some(buffer_group) = self.free_extra_bgids.pop_front() {
+            return Ok(buffer_group);
+        }
+
+        let buffer_group = self.next_extra_bgid;
+        self.next_extra_bgid = buffer_group
+            .checked_add(1)
+            .ok_or_else(|| io::Error::other("no buffer group id available"))?;
+
+        Ok(buffer_group)
+    }
+
+    fn release_extra_bgid(&mut self, buffer_group: u16) {
+        self.free_extra_bgids.push_back(buffer_group);
+    }
+
+    fn release_unused_extra_pools(&mut self) -> io::Result<()> {
+        while let Some(buffer_group) = { self.pending_extra_pool_releases.borrow_mut().pop_front() }
+        {
+            let Some(mut root) = self.extra_pools.remove(&buffer_group) else {
+                continue;
+            };
+
+            unsafe { root.release(&mut self.driver)? };
+
+            self.release_extra_bgid(buffer_group)
+        }
+
+        Ok(())
+    }
+
     /// Get a default [`Extra`] for underlying driver.
     pub fn default_extra(&self) -> Extra {
         Extra::new(&self.driver)
@@ -295,6 +340,7 @@ impl Proactor {
     /// You need to call [`Proactor::pop`] to get the pushed
     /// operations.
     pub fn poll(&mut self, timeout: Option<Duration>) -> io::Result<()> {
+        self.release_unused_extra_pools()?;
         self.driver.poll(timeout)
     }
 
@@ -477,11 +523,29 @@ impl Proactor {
         buffer_size: usize,
         flags: u16,
     ) -> io::Result<BufferPool> {
+        self.release_unused_extra_pools()?;
         let alloc = BufferAlloc::new::<A>();
-        let root =
-            BufferPoolRoot::new(&mut self.driver, alloc, num_of_bufs, buffer_size, flags)?;
+        let buffer_group = self.alloc_extra_bgid()?;
+        let release_queue = Rc::downgrade(&self.pending_extra_pool_releases);
+        let root = match BufferPoolRoot::new(
+            &mut self.driver,
+            alloc,
+            num_of_bufs,
+            buffer_size,
+            flags,
+            buffer_group,
+            Some(release_queue),
+        ) {
+            Ok(root) => root,
+            Err(err) => {
+                self.release_extra_bgid(buffer_group);
+                return Err(err);
+            }
+        };
+
         let pool = root.get_pool();
-        self.extra_pools.push(root);
+        self.extra_pools.insert(buffer_group, root);
+
         Ok(pool)
     }
 }

compio-driver/src/sys/buffer_pool/fusion.rs

@@ -21,9 +21,11 @@ impl BufControl {
         bufs: &[Slot],
         bufs_len: u32,
         flags: u16,
+        buffer_group: u16,
     ) -> io::Result<Self> {
         let inner = if driver.as_iour().is_some() {
-            let ctrl = unsafe { iour::BufControl::new(driver, bufs, bufs_len, flags) }?;
+            let ctrl =
+                unsafe { iour::BufControl::new(driver, bufs, bufs_len, flags, buffer_group) }?;
             Inner::IoUring(ctrl)
         } else {
             Inner::Fallback(fallback::BufControl::new(bufs))

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

@@ -23,14 +23,14 @@ pub(in crate::sys) struct BufControl {
     len: NonZeroU16,
     /// Total size of the mmap
     size: usize,
+    /// Buffer group registered in io_uring.
+    buffer_group: u16,
 }
 
 assert_not_impl!(BufControl, Send);
 assert_not_impl!(BufControl, Sync);
 
 impl BufControl {
-    const BUF_GROUP: u16 = 1;
-
     /// # Safety
     ///
     /// Caller must ensure the buffers will:
@@ -43,6 +43,7 @@ impl BufControl {
         bufs: &[Slot],
         bufs_len: u32,
         flags: u16,
+        buffer_group: u16,
     ) -> io::Result<Self> {
         debug_assert!(bufs.len().is_power_of_two());
 
@@ -57,13 +58,18 @@ impl BufControl {
             .expect("mmap failed")
             .cast::<BufRingEntry>();
 
-        let mut this = Self { ptr, len, size };
+        let mut this = Self {
+            ptr,
+            len,
+            size,
+            buffer_group,
+        };
 
         unsafe {
             driver.inner().submitter().register_buf_ring_with_flags(
                 ptr.addr().get() as u64,
                 len.get(),
-                Self::BUF_GROUP,
+                buffer_group,
                 flags,
             )
         }?;
@@ -88,7 +94,7 @@ impl BufControl {
 
     /// Get the buffer group id
     pub const fn buffer_group(&self) -> u16 {
-        Self::BUF_GROUP
+        self.buffer_group
     }
 
     /// Reset the buffer and make it available to the kernel
@@ -113,7 +119,7 @@ impl BufControl {
         driver
             .inner()
             .submitter()
-            .unregister_buf_ring(Self::BUF_GROUP)?;
+            .unregister_buf_ring(self.buffer_group)?;
         unsafe { munmap(self.ptr.cast().as_ptr(), self.size) }?;
 
         Ok(())

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

@@ -34,14 +34,15 @@ impl BufControl {
         bufs: &[Slot],
         buf_len: u32,
         flags: u16,
+        buffer_group: u16,
     ) -> io::Result<BufControl> {
         #[cfg(io_uring)]
-        let inner = unsafe { imp::BufControl::new(driver, bufs, buf_len, flags)? };
+        let inner = unsafe { imp::BufControl::new(driver, bufs, buf_len, flags, buffer_group)? };
 
         #[cfg(not(io_uring))]
         let inner = fallback::BufControl::new(bufs);
 
-        _ = (driver, buf_len, flags);
+        _ = (driver, buf_len, flags, buffer_group);
 
         Ok(Self(inner))
     }