Add safer abstractions for buffer rings

Signed-off-by: Alex Saveau <[email protected]>

io-uring-test/src/tests/register_buf_ring.rs

@@ -189,7 +189,7 @@ impl InnerBufRing {
         // it is unregistered. The backing store is an AnonymousMmap which remains valid until it
         // is dropped which in this case, is when Self is dropped.
         let res = unsafe {
-            ring.submitter().register_buf_ring(
+            ring.submitter().register_buf_ring_unchecked(
                 self.ring_start.as_ptr() as _,
                 self.ring_entries(),
                 bgid,

src/buf_ring.rs

@@ -0,0 +1,141 @@
+use std::mem::ManuallyDrop;
+use std::ops::{Deref, DerefMut};
+use std::{convert::TryFrom, marker::PhantomData, num::Wrapping, slice, sync::atomic};
+
+use crate::{sys, util::Mmap};
+
+pub struct BufRing {
+    ring: Mmap,
+    ring_entries: u16,
+    entry_size: u32,
+}
+
+impl BufRing {
+    pub(crate) fn init(ring: Mmap, ring_entries: u16, entry_size: u32) -> Self {
+        let mut this = Self {
+            ring,
+            ring_entries,
+            entry_size,
+        };
+
+        {
+            let mut s = this.submissions();
+            for i in 0u16..ring_entries {
+                let buf = unsafe { s._recycle_by_index(i) };
+                buf.len = entry_size;
+            }
+        }
+
+        this
+    }
+
+    pub fn submissions(&mut self) -> BufRingSubmissions<'_> {
+        let ring_ptr = self.ring.as_mut_ptr().cast::<sys::io_uring_buf>();
+        let tail_ptr = unsafe { self.ring.offset(8 + 4 + 2) };
+        let ring_entries = usize::from(self.ring_entries);
+        BufRingSubmissions {
+            ring_ptr,
+            buf_ptr: unsafe { ring_ptr.add(ring_entries).cast() },
+            tail_ptr: tail_ptr.cast::<atomic::AtomicU16>(),
+
+            tail: Wrapping(usize::from(unsafe { *tail_ptr.cast::<u16>() })),
+            tail_mask: ring_entries - 1,
+            entry_size: usize::try_from(self.entry_size).unwrap(),
+
+            _marker: PhantomData,
+        }
+    }
+}
+
+pub struct BufRingSubmissions<'ctx> {
+    ring_ptr: *mut sys::io_uring_buf,
+    buf_ptr: *mut libc::c_void,
+    tail_ptr: *const atomic::AtomicU16,
+
+    tail: Wrapping<usize>,
+    tail_mask: usize,
+    entry_size: usize,
+
+    _marker: PhantomData<&'ctx ()>,
+}
+
+impl<'a> BufRingSubmissions<'a> {
+    pub fn sync(&mut self) {
+        unsafe { &*self.tail_ptr }.store(self.tail.0 as u16, atomic::Ordering::Release);
+    }
+
+    pub unsafe fn get(&mut self, flags: u32, len: usize) -> Buf<'_, 'a> {
+        let index = Self::flags_to_index(flags);
+        let buf = unsafe { self.buf_ptr.add(usize::from(index) * self.entry_size) };
+        Buf {
+            buf: unsafe { slice::from_raw_parts_mut(buf.cast(), len) },
+            index,
+            submissions: self,
+        }
+    }
+
+    pub unsafe fn recycle(&mut self, flags: u32) {
+        self.recycle_by_index(Self::flags_to_index(flags));
+    }
+
+    pub unsafe fn recycle_by_index(&mut self, index: u16) {
+        self._recycle_by_index(index);
+    }
+
+    unsafe fn _recycle_by_index(&mut self, index: u16) -> &mut sys::io_uring_buf {
+        let uindex = usize::from(index);
+        {
+            let next_buf = unsafe { &mut *self.ring_ptr.add(self.tail.0 & self.tail_mask) };
+            next_buf.addr = unsafe { self.buf_ptr.add(uindex * self.entry_size) } as u64;
+            next_buf.bid = index;
+        }
+        self.tail += &1;
+
+        unsafe { &mut *self.ring_ptr.add(uindex) }
+    }
+
+    fn flags_to_index(flags: u32) -> u16 {
+        u16::try_from(flags >> sys::IORING_CQE_BUFFER_SHIFT).unwrap()
+    }
+}
+
+impl Drop for BufRingSubmissions<'_> {
+    fn drop(&mut self) {
+        self.sync()
+    }
+}
+
+pub struct Buf<'a, 'b> {
+    buf: &'a mut [u8],
+    index: u16,
+    submissions: &'a mut BufRingSubmissions<'b>,
+}
+
+impl Deref for Buf<'_, '_> {
+    type Target = [u8];
+
+    fn deref(&self) -> &Self::Target {
+        self.buf
+    }
+}
+
+impl DerefMut for Buf<'_, '_> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.buf
+    }
+}
+
+impl Buf<'_, '_> {
+    pub fn into_index(self) -> u16 {
+        let me = ManuallyDrop::new(self);
+        me.index
+    }
+}
+
+impl Drop for Buf<'_, '_> {
+    fn drop(&mut self) {
+        unsafe {
+            self.submissions.recycle_by_index(self.index);
+        }
+    }
+}

src/lib.rs

@@ -6,6 +6,7 @@
 
 #[macro_use]
 mod util;
+pub mod buf_ring;
 pub mod cqueue;
 pub mod opcode;
 pub mod register;

src/submit.rs

@@ -1,11 +1,13 @@
+use std::convert::TryFrom;
 use std::os::unix::io::{AsRawFd, RawFd};
 use std::sync::atomic;
 use std::{io, mem, ptr};
 
+use crate::buf_ring::BufRing;
 use crate::register::{execute, Probe};
 use crate::sys;
 use crate::types::{CancelBuilder, Timespec};
-use crate::util::{cast_ptr, OwnedFd};
+use crate::util::{cast_ptr, Mmap, OwnedFd};
 use crate::Parameters;
 
 use crate::register::Restriction;
@@ -437,6 +439,26 @@ impl<'a> Submitter<'a> {
         .map(drop)
     }
 
+    pub fn register_buf_ring(
+        &self,
+        ring_entries: u16,
+        bgid: u16,
+        entry_size: u32,
+    ) -> io::Result<BufRing> {
+        let bytes = || {
+            usize::from(ring_entries).checked_mul(
+                mem::size_of::<sys::io_uring_buf>()
+                    .checked_add(usize::try_from(entry_size).unwrap())?,
+            )
+        };
+
+        let ring = Mmap::new_anon(bytes().ok_or(io::ErrorKind::InvalidInput)?)?;
+        unsafe {
+            self.register_buf_ring_unchecked(ring.as_mut_ptr() as u64, ring_entries, bgid)?;
+        }
+        Ok(BufRing::init(ring, ring_entries, entry_size))
+    }
+
     /// Register buffer ring for provided buffers.
     ///
     /// Details can be found in the io_uring_register_buf_ring.3 man page.
@@ -451,7 +473,7 @@ impl<'a> Submitter<'a> {
     /// Developers must ensure that the `ring_addr` and its length represented by `ring_entries`
     /// are valid and will be valid until the bgid is unregistered or the ring destroyed,
     /// otherwise undefined behaviour may occur.
-    pub unsafe fn register_buf_ring(
+    pub unsafe fn register_buf_ring_unchecked(
         &self,
         ring_addr: u64,
         ring_entries: u16,

src/util.rs

@@ -37,6 +37,26 @@ impl Mmap {
         }
     }
 
+    pub fn new_anon(len: usize) -> io::Result<Mmap> {
+        unsafe {
+            match libc::mmap(
+                ptr::null_mut(),
+                len,
+                libc::PROT_READ | libc::PROT_WRITE,
+                libc::MAP_ANONYMOUS | libc::MAP_PRIVATE,
+                0,
+                0,
+            ) {
+                libc::MAP_FAILED => Err(io::Error::last_os_error()),
+                addr => {
+                    // here, `mmap` will never return null
+                    let addr = ptr::NonNull::new_unchecked(addr);
+                    Ok(Mmap { addr, len })
+                }
+            }
+        }
+    }
+
     /// Do not make the stored memory accessible by child processes after a `fork`.
     pub fn dontfork(&self) -> io::Result<()> {
         match unsafe { libc::madvise(self.addr.as_ptr(), self.len, libc::MADV_DONTFORK) } {