FixedBufPool: a dynamic fixed buffer collection (tokio-rs#190)

Make the internal infrastructure for fixed buffers more flexible,
allowing `FixedBuf` handles to be backed by alternative collections.
Add one new collection as `FixedBufPool`, enabling dynamic retrieval of
free buffers.

Author: mzabaluev

Cargo.toml

@@ -32,6 +32,7 @@ futures = "0.3.25"
 criterion = "0.4.0"
 # we use joinset in our tests
 tokio = "1.21.0"
+nix = "0.26.1"
 
 [package.metadata.docs.rs]
 all-features = true

src/buf/fixed/buffers.rs

@@ -1,113 +1,22 @@
-use libc::{iovec, UIO_MAXIOV};
-use std::cmp;
-use std::mem;
-use std::ptr;
-use std::slice;
+use libc::iovec;
 
-// Internal state shared by FixedBufRegistry and FixedBuf handles.
-pub(crate) struct FixedBuffers {
-    // Pointer to an allocated array of iovec records referencing
-    // the allocated buffers. The number of initialized records is the
-    // same as the length of the states array.
-    raw_bufs: ptr::NonNull<iovec>,
-    // State information on the buffers. Indices in this array correspond to
-    // the indices in the array at raw_bufs.
-    states: Vec<BufState>,
-    // Original capacity of raw_bufs as a Vec.
-    orig_cap: usize,
-}
-
-// State information of a buffer in the registry,
-enum BufState {
-    // The buffer is not in use.
-    // The field records the length of the initialized part.
-    Free { init_len: usize },
-    // The buffer is checked out.
-    // Its data are logically owned by the FixedBuf handle,
-    // which also keeps track of the length of the initialized part.
-    CheckedOut,
-}
-
-impl FixedBuffers {
-    pub(crate) fn new(bufs: impl Iterator<Item = Vec<u8>>) -> Self {
-        let bufs = bufs.take(cmp::min(UIO_MAXIOV as usize, 65_536));
-        let (size_hint, _) = bufs.size_hint();
-        let mut iovecs = Vec::with_capacity(size_hint);
-        let mut states = Vec::with_capacity(size_hint);
-        for mut buf in bufs {
-            iovecs.push(iovec {
-                iov_base: buf.as_mut_ptr() as *mut _,
-                iov_len: buf.capacity(),
-            });
-            states.push(BufState::Free {
-                init_len: buf.len(),
-            });
-            mem::forget(buf);
-        }
-        debug_assert_eq!(iovecs.len(), states.len());
-        // Safety: Vec::as_mut_ptr never returns null
-        let raw_bufs = unsafe { ptr::NonNull::new_unchecked(iovecs.as_mut_ptr()) };
-        let orig_cap = iovecs.capacity();
-        mem::forget(iovecs);
-        FixedBuffers {
-            raw_bufs,
-            states,
-            orig_cap,
-        }
-    }
-
-    // If the indexed buffer is free, changes its state to checked out and
-    // returns its data. If the buffer is already checked out, returns None.
-    pub(crate) fn check_out(&mut self, index: usize) -> Option<(iovec, usize)> {
-        let iovecs_ptr = self.raw_bufs;
-        self.states.get_mut(index).and_then(|state| match *state {
-            BufState::Free { init_len } => {
-                *state = BufState::CheckedOut;
-                // Safety: the allocated array under the pointer is valid
-                // for the lifetime of self, the index is inside the array
-                // as checked by Vec::get_mut above, called on the array of
-                // states that has the same length.
-                let iovec = unsafe { iovecs_ptr.as_ptr().add(index).read() };
-                Some((iovec, init_len))
-            }
-            BufState::CheckedOut => None,
-        })
-    }
-
-    // Sets the indexed buffer's state to free and records the updated length
-    // of its initialized part. The buffer addressed must be in the checked out
-    // state, otherwise this function may panic.
-    pub(crate) fn check_in(&mut self, index: usize, init_len: usize) {
-        let state = self.states.get_mut(index).expect("invalid buffer index");
-        debug_assert!(
-            matches!(state, BufState::CheckedOut),
-            "the buffer must be checked out"
-        );
-        *state = BufState::Free { init_len };
-    }
-
-    pub(crate) fn iovecs(&self) -> &[iovec] {
-        // Safety: the raw_bufs pointer is valid for the lifetime of self,
-        // the slice length is valid by construction.
-        unsafe { slice::from_raw_parts(self.raw_bufs.as_ptr(), self.states.len()) }
-    }
-}
+// Abstracts management of fixed buffers in a buffer registry.
+pub(crate) trait FixedBuffers {
+    // Provides access to the raw buffers as a slice of iovec.
+    fn iovecs(&self) -> &[iovec];
 
-impl Drop for FixedBuffers {
-    fn drop(&mut self) {
-        let iovecs = unsafe {
-            Vec::from_raw_parts(self.raw_bufs.as_ptr(), self.states.len(), self.orig_cap)
-        };
-        for (i, iovec) in iovecs.iter().enumerate() {
-            match self.states[i] {
-                BufState::Free { init_len } => {
-                    let ptr = iovec.iov_base as *mut u8;
-                    let cap = iovec.iov_len;
-                    let v = unsafe { Vec::from_raw_parts(ptr, init_len, cap) };
-                    mem::drop(v);
-                }
-                BufState::CheckedOut => unreachable!("all buffers must be checked in"),
-            }
-        }
-    }
+    /// Sets the indexed buffer's state to free and records the updated length
+    /// of its initialized part.
+    ///
+    /// # Panics
+    ///
+    /// The buffer addressed must be in the checked out state,
+    /// otherwise this function may panic.
+    ///
+    /// # Safety
+    ///
+    /// While the implementation of this method typically does not need to
+    /// do anything unsafe, the caller must ensure that the bytes in the buffer
+    /// are initialized up to the specified length.
+    unsafe fn check_in(&mut self, buf_index: u16, init_len: usize);
 }

src/buf/fixed/handle.rs

@@ -8,45 +8,59 @@ use std::mem::ManuallyDrop;
 use std::ops::{Deref, DerefMut};
 use std::rc::Rc;
 
+// Data to construct a `FixedBuf` handle from.
+pub(super) struct CheckedOutBuf {
+    // Pointer and size of the buffer.
+    pub iovec: iovec,
+    // Length of the initialized part.
+    pub init_len: usize,
+    // Buffer index.
+    pub index: u16,
+}
+
 /// A unique handle to a memory buffer that can be pre-registered with
 /// the kernel for `io-uring` operations.
 ///
-/// `FixedBuf` handles can be obtained from a [`FixedBufRegistry`] collection.
+/// `FixedBuf` handles can be obtained from a collection of fixed buffers,
+/// either [`FixedBufRegistry`] or [`FixedBufPool`].
 /// For each buffer, only a single `FixedBuf` handle can be either used by the
 /// application code or owned by an I/O operation at any given time,
 /// thus avoiding data races between `io-uring` operations in flight and
 /// the application accessing buffer data.
 ///
 /// [`FixedBufRegistry`]: super::FixedBufRegistry
+/// [`FixedBufPool`]:     super::FixedBufPool
 ///
 pub struct FixedBuf {
-    registry: Rc<RefCell<FixedBuffers>>,
+    registry: Rc<RefCell<dyn FixedBuffers>>,
     buf: ManuallyDrop<Vec<u8>>,
     index: u16,
 }
 
 impl Drop for FixedBuf {
     fn drop(&mut self) {
         let mut registry = self.registry.borrow_mut();
-        debug_assert_eq!(
-            registry.iovecs()[self.index as usize].iov_base as *const u8,
-            self.buf.as_ptr()
-        );
-        debug_assert_eq!(
-            registry.iovecs()[self.index as usize].iov_len,
-            self.buf.capacity()
-        );
-        registry.check_in(self.index as usize, self.buf.len());
+        // Safety: the length of the initialized data in the buffer has been
+        // maintained accordingly to the safety contracts on
+        // Self::new and IoBufMut.
+        unsafe {
+            registry.check_in(self.index, self.buf.len());
+        }
     }
 }
 
 impl FixedBuf {
-    pub(super) unsafe fn new(
-        registry: Rc<RefCell<FixedBuffers>>,
-        iovec: iovec,
-        init_len: usize,
-        index: u16,
-    ) -> Self {
+    // Safety: Validity constraints must apply to CheckedOutBuf members:
+    // - iovec must refer to an array allocated by Vec<u8>;
+    // - the array will not be deallocated until the buffer is checked in;
+    // - the data in the array must be initialized up to the number of bytes
+    //   given in init_len.
+    pub(super) unsafe fn new(registry: Rc<RefCell<dyn FixedBuffers>>, data: CheckedOutBuf) -> Self {
+        let CheckedOutBuf {
+            iovec,
+            init_len,
+            index,
+        } = data;
         let buf = Vec::from_raw_parts(iovec.iov_base as _, init_len, iovec.iov_len);
         FixedBuf {
             registry,

src/buf/fixed/mod.rs

@@ -4,18 +4,25 @@
 //! the `tokio-uring` runtime. Operations like [`File::read_fixed_at`][rfa] and
 //! [`File::write_fixed_at`][wfa] make use of buffers pre-mapped by
 //! the kernel to reduce per-I/O overhead.
-//! The [`FixedBufRegistry::register`] method is used to register a collection of
-//! buffers with the kernel; it must be called before any of the [`FixedBuf`]
-//! handles to the collection's buffers can be used with I/O operations.
+//!
+//! Two kinds of buffer collections are provided: [`FixedBufRegistry`] and
+//! [`FixedBufPool`], realizing two different patterns of buffer management.
+//! The `register` method on either of these types is used to register a
+//! collection of buffers with the kernel. It must be called before any of
+//! the [`FixedBuf`] handles to the collection's buffers can be used with
+//! I/O operations.
 //!
 //! [rfa]: crate::fs::File::read_fixed_at
 //! [wfa]: crate::fs::File::write_fixed_at
 
-mod buffers;
-pub(crate) use self::buffers::FixedBuffers;
-
 mod handle;
 pub use handle::FixedBuf;
 
+mod buffers;
+pub(crate) use buffers::FixedBuffers;
+
+mod pool;
+pub use pool::FixedBufPool;
+
 mod registry;
 pub use registry::FixedBufRegistry;