From c2512b61561ed8794aed97263d15a747d4b7b026 Mon Sep 17 00:00:00 2001 From: Sherlock Holo Date: Tue, 16 Jun 2026 19:22:17 +0800 Subject: [PATCH 1/2] feat: add runtime allocate BufferPool support --- compio-driver/src/buffer_pool.rs | 4 -- compio-driver/src/lib.rs | 34 +++++++++++++--- compio-driver/tests/buffer_pool.rs | 63 ++++++++++++++++++++++++++++++ compio-runtime/src/lib.rs | 33 ++++++++++++++++ 4 files changed, 125 insertions(+), 9 deletions(-) diff --git a/compio-driver/src/buffer_pool.rs b/compio-driver/src/buffer_pool.rs index 8859067c8..7c63fcee8 100644 --- a/compio-driver/src/buffer_pool.rs +++ b/compio-driver/src/buffer_pool.rs @@ -193,10 +193,6 @@ impl BufferPoolRoot { shared: Rc::downgrade(&self.shared), } } - - pub(crate) fn is_unique(&self) -> bool { - Rc::strong_count(&self.shared) == 1 - } } impl Debug for BufferPool { diff --git a/compio-driver/src/lib.rs b/compio-driver/src/lib.rs index 4dfca371b..b734cbca1 100644 --- a/compio-driver/src/lib.rs +++ b/compio-driver/src/lib.rs @@ -101,6 +101,7 @@ impl PushEntry { pub struct Proactor { driver: Driver, buffer_pool: BufferPoolState, + extra_pools: Vec, } enum BufferPoolState { @@ -139,11 +140,12 @@ impl BufferPoolState { impl Drop for Proactor { fn drop(&mut self) { - let BufferPoolState::Init(buffer_pool) = &mut self.buffer_pool else { - return; - }; - debug_assert!(buffer_pool.is_unique()); // Just in case. Shouldn't happen - _ = unsafe { buffer_pool.release(&mut self.driver) }; + if let BufferPoolState::Init(buffer_pool) = &mut self.buffer_pool { + _ = unsafe { buffer_pool.release(&mut self.driver) }; + } + for pool in &mut self.extra_pools { + _ = unsafe { pool.release(&mut self.driver) }; + } } } @@ -170,6 +172,7 @@ impl Proactor { buffer_len: builder.buffer_pool_buffer_len, flags: builder.buffer_pool_flag, }, + extra_pools: Vec::new(), }) } @@ -460,6 +463,27 @@ impl Proactor { pub fn buffer_pool(&mut self) -> io::Result { self.buffer_pool.get(&mut self.driver) } + + /// Create a new buffer pool independently. + /// + /// Unlike [`buffer_pool`](Self::buffer_pool) which lazily initializes a + /// single pool configured via [`ProactorBuilder`], this method allows + /// creating additional pools at runtime with custom parameters. + /// + /// The pool will be released when this `Proactor` is dropped. + pub fn create_buffer_pool( + &mut self, + num_of_bufs: u16, + buffer_size: usize, + flags: u16, + ) -> io::Result { + let alloc = BufferAlloc::new::(); + let root = + BufferPoolRoot::new(&mut self.driver, alloc, num_of_bufs, buffer_size, flags)?; + let pool = root.get_pool(); + self.extra_pools.push(root); + Ok(pool) + } } impl AsRawFd for Proactor { diff --git a/compio-driver/tests/buffer_pool.rs b/compio-driver/tests/buffer_pool.rs index cdca2fc0e..ca290d512 100644 --- a/compio-driver/tests/buffer_pool.rs +++ b/compio-driver/tests/buffer_pool.rs @@ -168,3 +168,66 @@ fn buffer_pool_recycle() { drop(pool); drop(buf); } + +#[cfg(any(not(target_os = "linux"), feature = "polling"))] +#[test] +fn create_buffer_pool_multiple() { + use compio_buf::IoBufMut; + use compio_driver::BoxAllocator; + + let mut driver = build_proactor(2, 4096); + + // Default pool + let pool = driver.buffer_pool().unwrap(); + + // Second pool + let pool2 = driver + .create_buffer_pool::(2, 4096, 0) + .unwrap(); + + // Default pool still works + let _buf = pool.pop().unwrap(); + + // Second pool also works + let _buf2 = pool2.pop().unwrap(); + + // Third pool with different size + let pool3 = driver + .create_buffer_pool::(4, 1024, 0) + .unwrap(); + let mut buf3 = pool3.pop().unwrap(); + assert_eq!(buf3.as_uninit().len(), 1024); +} + +#[cfg(any(not(target_os = "linux"), feature = "polling"))] +#[test] +fn create_buffer_pool_independent() { + use compio_buf::IoBufMut; + use compio_driver::BoxAllocator; + + let mut driver = build_proactor(2, 4096); + + let pool_a = driver + .create_buffer_pool::(2, 4096, 0) + .unwrap(); + let pool_b = driver + .create_buffer_pool::(2, 4096, 0) + .unwrap(); + + // Each pool has its own buffers + let mut buf_a1 = pool_a.pop().unwrap(); + let mut buf_a2 = pool_a.pop().unwrap(); + let mut buf_b1 = pool_b.pop().unwrap(); + let mut buf_b2 = pool_b.pop().unwrap(); + + // Verify all pointers are distinct + let p_a1 = buf_a1.as_uninit().as_ptr(); + let p_a2 = buf_a2.as_uninit().as_ptr(); + let p_b1 = buf_b1.as_uninit().as_ptr(); + let p_b2 = buf_b2.as_uninit().as_ptr(); + + assert_ne!(p_a1, p_a2); + assert_ne!(p_b1, p_b2); + // Buffers from different pools may have the same address (reallocation), + // but within the same pool they must be distinct +} diff --git a/compio-runtime/src/lib.rs b/compio-runtime/src/lib.rs index 582d9d2ae..d16ac5c2b 100644 --- a/compio-runtime/src/lib.rs +++ b/compio-runtime/src/lib.rs @@ -295,6 +295,25 @@ impl Runtime { self.driver.borrow_mut().buffer_pool() } + /// Create a new buffer pool independently. + /// + /// Unlike [`buffer_pool`](Self::buffer_pool) which provides the default + /// pool, this allows creating additional buffer pools at runtime with + /// custom allocators and parameters. + /// + /// The pool will be released when the `Runtime`'s inner `Proactor` is + /// dropped. + pub fn create_buffer_pool( + &self, + num_of_bufs: u16, + buffer_size: usize, + flags: u16, + ) -> io::Result { + self.driver + .borrow_mut() + .create_buffer_pool::(num_of_bufs, buffer_size, flags) + } + /// Register file descriptors for fixed-file operations. /// /// This is only supported on io-uring driver, and will return an @@ -566,3 +585,17 @@ pub fn register_files(fds: &[RawFd]) -> io::Result<()> { pub fn unregister_files() -> io::Result<()> { Runtime::with_current(|r| r.unregister_files()) } + +/// Create a new buffer pool for the current runtime. +/// +/// ## Panics +/// +/// This method doesn't create runtime. It tries to obtain the current runtime +/// by [`Runtime::with_current`]. +pub fn create_buffer_pool( + num_of_bufs: u16, + buffer_size: usize, + flags: u16, +) -> io::Result { + Runtime::with_current(|r| r.create_buffer_pool::(num_of_bufs, buffer_size, flags)) +} From 1eedc89961d0a2c50e34aadb3af95b9a8afe6e51 Mon Sep 17 00:00:00 2001 From: Sherlock Holo Date: Tue, 16 Jun 2026 19:58:45 +0800 Subject: [PATCH 2/2] feat: add delay RAII release for BufferPool --- compio-driver/src/buffer_pool.rs | 48 +++++++- compio-driver/src/lib.rs | 76 +++++++++++- compio-driver/src/sys/buffer_pool/fusion.rs | 4 +- compio-driver/src/sys/buffer_pool/iour.rs | 18 ++- compio-driver/src/sys/buffer_pool/mod.rs | 5 +- compio-driver/tests/buffer_pool.rs | 122 ++++++++++++++++++++ 6 files changed, 256 insertions(+), 17 deletions(-) diff --git a/compio-driver/src/buffer_pool.rs b/compio-driver/src/buffer_pool.rs index 7c63fcee8..2aa9c6f23 100644 --- a/compio-driver/src/buffer_pool.rs +++ b/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, + + /// Queue used by extra buffer pools to ask the proactor to release them. + release_queue: Option, +} + +#[derive(Debug)] +struct ReleaseQueue { + buffer_group: u16, + queue: Weak>>, } impl BufferPoolRoot { @@ -132,6 +142,8 @@ impl BufferPoolRoot { num_of_bufs: u16, buffer_size: usize, flags: u16, + buffer_group: u16, + release_queue: Option>>>, ) -> io::Result { 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::>(); - 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) { + 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(); + } + } +} diff --git a/compio-driver/src/lib.rs b/compio-driver/src/lib.rs index b734cbca1..e10c9ed19 100644 --- a/compio-driver/src/lib.rs +++ b/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 PushEntry { pub struct Proactor { driver: Driver, buffer_pool: BufferPoolState, - extra_pools: Vec, + extra_pools: HashMap, + pending_extra_pool_releases: Rc>>, + free_extra_bgids: VecDeque, + 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 { 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 { + 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) -> io::Result<()> { + self.release_unused_extra_pools()?; self.driver.poll(timeout) } @@ -477,11 +523,29 @@ impl Proactor { buffer_size: usize, flags: u16, ) -> io::Result { + self.release_unused_extra_pools()?; let alloc = BufferAlloc::new::(); - 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) } } diff --git a/compio-driver/src/sys/buffer_pool/fusion.rs b/compio-driver/src/sys/buffer_pool/fusion.rs index 7d368f407..3f4b22e92 100644 --- a/compio-driver/src/sys/buffer_pool/fusion.rs +++ b/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 { 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)) diff --git a/compio-driver/src/sys/buffer_pool/iour.rs b/compio-driver/src/sys/buffer_pool/iour.rs index ca94e1f7b..1654c50c7 100644 --- a/compio-driver/src/sys/buffer_pool/iour.rs +++ b/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 { debug_assert!(bufs.len().is_power_of_two()); @@ -57,13 +58,18 @@ impl BufControl { .expect("mmap failed") .cast::(); - 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(()) diff --git a/compio-driver/src/sys/buffer_pool/mod.rs b/compio-driver/src/sys/buffer_pool/mod.rs index 0c4c9d9fc..cc7653cc8 100644 --- a/compio-driver/src/sys/buffer_pool/mod.rs +++ b/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 { #[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)) } diff --git a/compio-driver/tests/buffer_pool.rs b/compio-driver/tests/buffer_pool.rs index ca290d512..3b193984e 100644 --- a/compio-driver/tests/buffer_pool.rs +++ b/compio-driver/tests/buffer_pool.rs @@ -231,3 +231,125 @@ fn create_buffer_pool_independent() { // Buffers from different pools may have the same address (reallocation), // but within the same pool they must be distinct } + +#[cfg(any(not(target_os = "linux"), feature = "polling"))] +#[test] +fn extra_buffer_pool_released_after_last_pool_drop() { + use std::{ + mem::MaybeUninit, + ptr::NonNull, + sync::atomic::{AtomicUsize, Ordering}, + time::Duration, + }; + + use compio_driver::{BoxAllocator, BufferAllocator}; + + static DEALLOCS: AtomicUsize = AtomicUsize::new(0); + + struct CountingAllocator; + + impl BufferAllocator for CountingAllocator { + fn allocate(len: u32) -> NonNull> { + BoxAllocator::allocate(len) + } + + unsafe fn deallocate(ptr: NonNull>, len: u32) { + DEALLOCS.fetch_add(1, Ordering::SeqCst); + unsafe { BoxAllocator::deallocate(ptr, len) }; + } + } + + DEALLOCS.store(0, Ordering::SeqCst); + + let mut driver = build_proactor(1, 4096); + let pool = driver + .create_buffer_pool::(2, 1024, 0) + .unwrap(); + + drop(pool); + assert_eq!(DEALLOCS.load(Ordering::SeqCst), 0); + + _ = driver.poll(Some(Duration::ZERO)); + assert_eq!(DEALLOCS.load(Ordering::SeqCst), 2); +} + +#[cfg(any(not(target_os = "linux"), feature = "polling"))] +#[test] +fn extra_buffer_pool_waits_for_live_buffer_ref() { + use std::{ + mem::MaybeUninit, + ptr::NonNull, + sync::atomic::{AtomicUsize, Ordering}, + time::Duration, + }; + + use compio_driver::{BoxAllocator, BufferAllocator}; + + static DEALLOCS: AtomicUsize = AtomicUsize::new(0); + + struct CountingAllocator; + + impl BufferAllocator for CountingAllocator { + fn allocate(len: u32) -> NonNull> { + BoxAllocator::allocate(len) + } + + unsafe fn deallocate(ptr: NonNull>, len: u32) { + DEALLOCS.fetch_add(1, Ordering::SeqCst); + unsafe { BoxAllocator::deallocate(ptr, len) }; + } + } + + DEALLOCS.store(0, Ordering::SeqCst); + + let mut driver = build_proactor(1, 4096); + let pool = driver + .create_buffer_pool::(1, 1024, 0) + .unwrap(); + let buf = pool.pop().unwrap(); + + drop(pool); + _ = driver.poll(Some(Duration::ZERO)); + assert_eq!(DEALLOCS.load(Ordering::SeqCst), 0); + + drop(buf); + _ = driver.poll(Some(Duration::ZERO)); + assert_eq!(DEALLOCS.load(Ordering::SeqCst), 1); +} + +#[cfg(io_uring)] +#[test] +fn create_buffer_pool_iouring_multiple_groups() { + use compio_buf::IoBuf; + use compio_driver::{BoxAllocator, DriverType}; + + let mut driver = ProactorBuilder::new() + .driver_type(DriverType::IoUring) + .build() + .unwrap(); + if !driver.driver_type().is_iouring() { + return; + } + + let file = std::fs::File::open("Cargo.toml").unwrap(); + let fd = SharedFd::new(file); + driver.attach(fd.as_raw_fd()).unwrap(); + + let pool_a = driver + .create_buffer_pool::(2, 128, 0) + .unwrap(); + let pool_b = driver + .create_buffer_pool::(2, 256, 0) + .unwrap(); + + let op = ReadManagedAt::new(fd.clone(), 0, &pool_a, 32).unwrap(); + let res = push_and_wait(&mut driver, op); + let buffer = unsafe { res.take_buffer() }.unwrap().unwrap(); + assert!(buffer.as_init().starts_with(b"[package]")); + drop(buffer); + + let op = ReadManagedAt::new(fd, 0, &pool_b, 32).unwrap(); + let res = push_and_wait(&mut driver, op); + let buffer = unsafe { res.take_buffer() }.unwrap().unwrap(); + assert!(buffer.as_init().starts_with(b"[package]")); +}