add Buffer pool support

Cargo.toml

@@ -49,12 +49,14 @@ nix = "0.29.0"
 once_cell = "1.18.0"
 os_pipe = "1.1.4"
 paste = "1.0.14"
+rand = "0.8.5"
 slab = "0.4.9"
 socket2 = "0.5.6"
 tempfile = "3.8.1"
 tokio = "1.33.0"
 widestring = "1.0.2"
 windows-sys = "0.52.0"
+pin-project-lite = "0.2.14"
 
 [profile.bench]
 debug = true

compio-driver/Cargo.toml

@@ -42,6 +42,7 @@ socket2 = { workspace = true }
 [target.'cfg(windows)'.dependencies]
 aligned-array = "1.0.1"
 once_cell = { workspace = true }
+pin-project-lite = { workspace = true }
 windows-sys = { workspace = true, features = [
     "Win32_Foundation",
     "Win32_Networking_WinSock",
@@ -58,12 +59,16 @@ windows-sys = { workspace = true, features = [
 # Linux specific dependencies
 [target.'cfg(target_os = "linux")'.dependencies]
 io-uring = { version = "0.6.2", optional = true }
+io_uring_buf_ring = { version = "0.1.0", optional = true }
 polling = { version = "3.3.0", optional = true }
 os_pipe = { workspace = true, optional = true }
 paste = { workspace = true }
+pin-project-lite = { workspace = true }
+slab = { version = "0.4.9", optional = true }
 
 # Other platform dependencies
 [target.'cfg(all(not(target_os = "linux"), unix))'.dependencies]
+pin-project-lite = { workspace = true }
 polling = "3.3.0"
 os_pipe = { workspace = true }
 
@@ -77,11 +82,13 @@ compio-buf = { workspace = true, features = ["arrayvec"] }
 
 [features]
 default = ["io-uring"]
+io-uring = ["dep:io-uring"]
 polling = ["dep:polling", "dep:os_pipe"]
 
 io-uring-sqe128 = []
 io-uring-cqe32 = []
 io-uring-socket = []
+io-uring-buf-ring = ["dep:io_uring_buf_ring", "dep:slab"]
 
 iocp-global = []
 

compio-driver/src/fallback_buffer_pool.rs

@@ -0,0 +1,117 @@
+use std::{
+    borrow::{Borrow, BorrowMut},
+    cell::RefCell,
+    collections::VecDeque,
+    fmt::{Debug, Formatter},
+    mem::ManuallyDrop,
+    ops::{Deref, DerefMut},
+};
+
+use compio_buf::{IntoInner, Slice};
+
+/// Buffer pool
+///
+/// A buffer pool to allow user no need to specify a specific buffer to do the
+/// IO operation
+pub struct BufferPool {
+    buffers: RefCell<VecDeque<Vec<u8>>>,
+}
+
+impl Debug for BufferPool {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("BufferPool").finish_non_exhaustive()
+    }
+}
+
+impl BufferPool {
+    pub(crate) fn new(buffer_len: u16, buffer_size: usize) -> Self {
+        let buffers = (0..buffer_len.next_power_of_two())
+            .map(|_| Vec::with_capacity(buffer_size))
+            .collect();
+
+        Self {
+            buffers: RefCell::new(buffers),
+        }
+    }
+
+    pub(crate) fn get_buffer(&self) -> Option<Vec<u8>> {
+        self.buffers.borrow_mut().pop_front()
+    }
+
+    pub(crate) fn add_buffer(&self, mut buffer: Vec<u8>) {
+        buffer.clear();
+        self.buffers.borrow_mut().push_back(buffer)
+    }
+}
+
+/// Buffer borrowed from buffer pool
+///
+/// When IO operation finish, user will obtain a `BorrowedBuffer` to access the
+/// filled data
+pub struct BorrowedBuffer<'a> {
+    buffer: ManuallyDrop<Slice<Vec<u8>>>,
+    pool: &'a BufferPool,
+}
+
+impl<'a> BorrowedBuffer<'a> {
+    pub(crate) fn new(buffer: Slice<Vec<u8>>, pool: &'a BufferPool) -> Self {
+        Self {
+            buffer: ManuallyDrop::new(buffer),
+            pool,
+        }
+    }
+}
+
+impl Debug for BorrowedBuffer<'_> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("BorrowedBuffer").finish_non_exhaustive()
+    }
+}
+
+impl Drop for BorrowedBuffer<'_> {
+    fn drop(&mut self) {
+        let buffer = unsafe {
+            // Safety: we won't take self.buffer again
+            ManuallyDrop::take(&mut self.buffer)
+        };
+        self.pool.add_buffer(buffer.into_inner());
+    }
+}
+
+impl Deref for BorrowedBuffer<'_> {
+    type Target = [u8];
+
+    fn deref(&self) -> &Self::Target {
+        self.buffer.deref()
+    }
+}
+
+impl DerefMut for BorrowedBuffer<'_> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.buffer.deref_mut()
+    }
+}
+
+impl AsRef<[u8]> for BorrowedBuffer<'_> {
+    fn as_ref(&self) -> &[u8] {
+        self.deref()
+    }
+}
+
+impl AsMut<[u8]> for BorrowedBuffer<'_> {
+    fn as_mut(&mut self) -> &mut [u8] {
+        self.deref_mut()
+    }
+}
+
+impl Borrow<[u8]> for BorrowedBuffer<'_> {
+    fn borrow(&self) -> &[u8] {
+        self.deref()
+    }
+}
+
+impl BorrowMut<[u8]> for BorrowedBuffer<'_> {
+    fn borrow_mut(&mut self) -> &mut [u8] {
+        self.deref_mut()
+    }
+}

compio-driver/src/fusion/buffer_pool.rs

@@ -0,0 +1,145 @@
+use std::{
+    borrow::{Borrow, BorrowMut},
+    fmt::{Debug, Formatter},
+    ops::{Deref, DerefMut},
+};
+
+/// Buffer pool
+///
+/// A buffer pool to allow user no need to specify a specific buffer to do the
+/// IO operation
+pub struct BufferPool {
+    inner: BufferPollInner,
+}
+
+impl Debug for BufferPool {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("BufferPool").finish_non_exhaustive()
+    }
+}
+
+impl BufferPool {
+    pub(crate) fn new_io_uring(buffer_pool: super::iour::buffer_pool::BufferPool) -> Self {
+        Self {
+            inner: BufferPollInner::IoUring(buffer_pool),
+        }
+    }
+
+    pub(crate) fn as_io_uring(&self) -> &super::iour::buffer_pool::BufferPool {
+        match &self.inner {
+            BufferPollInner::IoUring(inner) => inner,
+            BufferPollInner::Poll(_) => panic!("BufferPool type is not poll type"),
+        }
+    }
+
+    pub(crate) fn as_poll(&self) -> &crate::fallback_buffer_pool::BufferPool {
+        match &self.inner {
+            BufferPollInner::Poll(inner) => inner,
+            BufferPollInner::IoUring(_) => panic!("BufferPool type is not io-uring type"),
+        }
+    }
+
+    pub(crate) fn new_poll(buffer_pool: crate::fallback_buffer_pool::BufferPool) -> Self {
+        Self {
+            inner: BufferPollInner::Poll(buffer_pool),
+        }
+    }
+
+    pub(crate) fn into_poll(self) -> crate::fallback_buffer_pool::BufferPool {
+        match self.inner {
+            BufferPollInner::IoUring(_) => {
+                panic!("BufferPool type is not io-uring type")
+            }
+            BufferPollInner::Poll(inner) => inner,
+        }
+    }
+
+    pub(crate) fn into_io_uring(self) -> super::iour::buffer_pool::BufferPool {
+        match self.inner {
+            BufferPollInner::IoUring(inner) => inner,
+            BufferPollInner::Poll(_) => panic!("BufferPool type is not poll type"),
+        }
+    }
+}
+
+enum BufferPollInner {
+    IoUring(super::iour::buffer_pool::BufferPool),
+    Poll(crate::fallback_buffer_pool::BufferPool),
+}
+
+/// Buffer borrowed from buffer pool
+///
+/// When IO operation finish, user will obtain a `BorrowedBuffer` to access the
+/// filled data
+pub struct BorrowedBuffer<'a> {
+    inner: BorrowedBufferInner<'a>,
+}
+
+impl<'a> BorrowedBuffer<'a> {
+    pub(crate) fn new_io_uring(buffer: super::iour::buffer_pool::BorrowedBuffer<'a>) -> Self {
+        Self {
+            inner: BorrowedBufferInner::IoUring(buffer),
+        }
+    }
+
+    pub(crate) fn new_poll(buffer: crate::fallback_buffer_pool::BorrowedBuffer<'a>) -> Self {
+        Self {
+            inner: BorrowedBufferInner::Poll(buffer),
+        }
+    }
+}
+
+enum BorrowedBufferInner<'a> {
+    IoUring(super::iour::buffer_pool::BorrowedBuffer<'a>),
+    Poll(crate::fallback_buffer_pool::BorrowedBuffer<'a>),
+}
+
+impl Debug for BorrowedBuffer<'_> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("BorrowedBuffer").finish_non_exhaustive()
+    }
+}
+
+impl Deref for BorrowedBuffer<'_> {
+    type Target = [u8];
+
+    fn deref(&self) -> &Self::Target {
+        match &self.inner {
+            BorrowedBufferInner::IoUring(inner) => inner.deref(),
+            BorrowedBufferInner::Poll(inner) => inner.deref(),
+        }
+    }
+}
+
+impl DerefMut for BorrowedBuffer<'_> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        match &mut self.inner {
+            BorrowedBufferInner::IoUring(inner) => inner.deref_mut(),
+            BorrowedBufferInner::Poll(inner) => inner.deref_mut(),
+        }
+    }
+}
+
+impl AsRef<[u8]> for BorrowedBuffer<'_> {
+    fn as_ref(&self) -> &[u8] {
+        self.deref()
+    }
+}
+
+impl AsMut<[u8]> for BorrowedBuffer<'_> {
+    fn as_mut(&mut self) -> &mut [u8] {
+        self.deref_mut()
+    }
+}
+
+impl Borrow<[u8]> for BorrowedBuffer<'_> {
+    fn borrow(&self) -> &[u8] {
+        self.deref()
+    }
+}
+
+impl BorrowMut<[u8]> for BorrowedBuffer<'_> {
+    fn borrow_mut(&mut self) -> &mut [u8] {
+        self.deref_mut()
+    }
+}

compio-driver/src/fusion/mod.rs

@@ -4,6 +4,7 @@ mod poll;
 #[path = "../iour/mod.rs"]
 mod iour;
 
+pub(crate) mod buffer_pool;
 pub(crate) mod op;
 
 #[cfg_attr(all(doc, docsrs), doc(cfg(all())))]
@@ -15,7 +16,7 @@ pub(crate) use iour::{sockaddr_storage, socklen_t};
 pub use iour::{OpCode as IourOpCode, OpEntry};
 pub use poll::{Decision, OpCode as PollOpCode};
 
-use crate::{Key, OutEntries, ProactorBuilder};
+use crate::{BufferPool, Key, OutEntries, ProactorBuilder};
 
 mod driver_type {
     use std::sync::atomic::{AtomicU8, Ordering};
@@ -180,6 +181,31 @@ impl Driver {
         };
         Ok(NotifyHandle::from_fuse(fuse))
     }
+
+    pub fn create_buffer_pool(
+        &mut self,
+        buffer_len: u16,
+        buffer_size: usize,
+    ) -> io::Result<BufferPool> {
+        match &mut self.fuse {
+            FuseDriver::IoUring(driver) => Ok(BufferPool::new_io_uring(
+                driver.create_buffer_pool(buffer_len, buffer_size)?,
+            )),
+            FuseDriver::Poll(driver) => Ok(BufferPool::new_poll(
+                driver.create_buffer_pool(buffer_len, buffer_size)?,
+            )),
+        }
+    }
+
+    /// # Safety
+    ///
+    /// caller must make sure release the buffer pool with correct driver
+    pub unsafe fn release_buffer_pool(&mut self, buffer_pool: BufferPool) -> io::Result<()> {
+        match &mut self.fuse {
+            FuseDriver::Poll(driver) => driver.release_buffer_pool(buffer_pool.into_poll()),
+            FuseDriver::IoUring(driver) => driver.release_buffer_pool(buffer_pool.into_io_uring()),
+        }
+    }
 }
 
 impl AsRawFd for Driver {