lib: add register_buffers* functions (#313)

* lib: add `register_buffers_*` family of functions

* add tests and docs

* chore: update tests and names

---------

Co-authored-by: Fathy Boundjadj <[email protected]>

Author: saiintbrisson

io-uring-test/src/main.rs

@@ -75,8 +75,9 @@ fn test<S: squeue::EntryMarker, C: cqueue::EntryMarker>(
 
     // register
     tests::register::test_register_files_sparse(&mut ring, &test)?;
-    tests::register_buf_ring::test_register_buf_ring(&mut ring, &test)?;
     tests::register_buffers::test_register_buffers(&mut ring, &test)?;
+    tests::register_buffers::test_register_buffers_update(&mut ring, &test)?;
+    tests::register_buf_ring::test_register_buf_ring(&mut ring, &test)?;
     tests::register_sync_cancel::test_register_sync_cancel(&mut ring, &test)?;
     tests::register_sync_cancel::test_register_sync_cancel_unsubmitted(&mut ring, &test)?;
     tests::register_sync_cancel::test_register_sync_cancel_any(&mut ring, &test)?;

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

@@ -8,24 +8,70 @@ use io_uring::{
 };
 use libc::iovec;
 use std::{
+    fs::File,
     io::{self, IoSliceMut},
+    io::{Error, Write},
     ops::DerefMut,
     os::fd::AsRawFd,
+    os::fd::FromRawFd,
 };
 
-/// Creates shared buffers to be registered into the [`IoUring`] instance, and then tests that they
-/// can actually be used through the [`WriteFixed`] and [`ReadFixed`] opcodes.
+use io_uring::{opcode, types};
+use libc::EFAULT;
+
 pub fn test_register_buffers<S: squeue::EntryMarker, C: cqueue::EntryMarker>(
     ring: &mut IoUring<S, C>,
     test: &Test,
-) -> io::Result<()> {
+) -> anyhow::Result<()> {
+    _test_register_buffers(
+        ring,
+        test,
+        "register_buffers",
+        None,
+        |ring, iovecs, _| unsafe { ring.submitter().register_buffers(&iovecs) },
+    )?;
+    _test_register_buffers(
+        ring,
+        test,
+        "register_buffers2",
+        ring.params().is_feature_resource_tagging(),
+        |ring, iovecs, tags| unsafe { ring.submitter().register_buffers2(iovecs, tags) },
+    )?;
+    _test_register_buffers(
+        ring,
+        test,
+        "register_buffers_sparse",
+        ring.params().is_feature_resource_tagging(),
+        |ring, iovecs, _| {
+            let submitter = ring.submitter();
+            submitter.register_buffers_sparse(iovecs.len() as _)?;
+            unsafe { submitter.register_buffers_update(0, iovecs, None) }
+        },
+    )?;
+
+    return Ok(());
+}
+
+fn _test_register_buffers<
+    S: squeue::EntryMarker,
+    C: cqueue::EntryMarker,
+    F: FnMut(&mut IoUring<S, C>, &[iovec], &[u64]) -> io::Result<()>,
+    P: Into<Option<bool>>,
+>(
+    ring: &mut IoUring<S, C>,
+    test: &Test,
+    name: &str,
+    probe: P,
+    mut register: F,
+) -> anyhow::Result<()> {
     require!(
         test;
         test.probe.is_supported(WriteFixed::CODE);
         test.probe.is_supported(ReadFixed::CODE);
+        probe.into().unwrap_or(true);
     );
 
-    println!("test register_buffers");
+    println!("test {name}");
 
     const BUF_SIZE: usize = 1 << 12; // Page size
     const BUFFERS: usize = 8; // The maximum number of squeue entries (in main.rs)
@@ -56,11 +102,11 @@ pub fn test_register_buffers<S: squeue::EntryMarker, C: cqueue::EntryMarker>(
     // to cast these as `slices` is valid for this entire function
     let iovecs: &[iovec] =
         unsafe { std::slice::from_raw_parts(slices.as_ptr().cast(), slices.len()) };
+    let tags = vec![0; slices.len()];
 
-    let submitter = ring.submitter();
     // Safety: Since `iovecs` is derived from valid `IoSliceMut`s, this upholds the safety contract
     // of `register_buffers` that the buffers are valid until the buffers are unregistered
-    unsafe { submitter.register_buffers(iovecs)? };
+    register(ring, iovecs, &tags).unwrap();
 
     // Prepare writing the buffers out to the file
     (0..BUFFERS).for_each(|index| {
@@ -141,3 +187,165 @@ pub fn test_register_buffers<S: squeue::EntryMarker, C: cqueue::EntryMarker>(
 
     Ok(())
 }
+
+const BUFFER_TAG: u64 = 0xbadcafe;
+const TIMEOUT_TAG: u64 = 0xbadf00d;
+
+pub fn test_register_buffers_update<S: squeue::EntryMarker, C: cqueue::EntryMarker>(
+    ring: &mut IoUring<S, C>,
+    test: &Test,
+) -> anyhow::Result<()> {
+    require!(
+        test;
+        test.probe.is_supported(opcode::ReadFixed::CODE);
+        ring.params().is_feature_resource_tagging();
+    );
+
+    println!("test register_buffers_update");
+
+    let (read, mut write) = create_pipe()?;
+    let mut buf = [0xde, 0xed, 0xbe, 0xef];
+    let mut tags = [BUFFER_TAG];
+    let iovecs = [libc::iovec {
+        iov_len: buf.len(),
+        iov_base: buf.as_mut_ptr() as _,
+    }];
+    let timeout = types::Timespec::new().nsec(50 * 1_000_000);
+    let timeout = opcode::Timeout::new(&timeout as _)
+        .build()
+        .user_data(TIMEOUT_TAG)
+        .into();
+    let read_sqe = opcode::ReadFixed::new(
+        types::Fd(read.as_raw_fd()),
+        buf.as_mut_ptr(),
+        buf.len() as _,
+        5,
+    )
+    .build()
+    .user_data(42)
+    .into();
+
+    // Register a buffer table and then immediately unregister it
+    ring.submitter().register_buffers_sparse(1)?;
+    ring.submitter().unregister_buffers()?;
+
+    // Push a timeout of 50ms
+    unsafe { ring.submission().push(&timeout).unwrap() };
+
+    // We should not receive any other entries than the timeout
+    check_only_timeout(ring)?;
+
+    // Register a sparse buffer table of 10 elements
+    ring.submitter().register_buffers_sparse(10)?;
+
+    // Try read the pipe using a sparse buffer
+    let cqe = {
+        write.write("yo".as_bytes())?;
+
+        unsafe { ring.submission().push(&read_sqe).unwrap() };
+
+        ring.submit_and_wait(1)?;
+
+        ring.completion().next().unwrap().into()
+    };
+
+    // We should get the correct user_data
+    if cqe.user_data() != 42 {
+        return Err(anyhow::anyhow!("unexpected completion queue entry"));
+    }
+
+    // EFAULT is to be expected with incorrect fixed buffers
+    if cqe.result() != -EFAULT {
+        return Err(anyhow::anyhow!("unexpected read result: {}", cqe.result()));
+    }
+
+    // Register a buffer at the index 5
+    unsafe {
+        ring.submitter()
+            .register_buffers_update(5, &iovecs, Some(&tags))?;
+
+        // Push a timeout of 50ms
+        ring.submission().push(&timeout).unwrap();
+    }
+
+    // We should not receive any other entries than the timeout
+    check_only_timeout(ring)?;
+
+    // Register a buffer at the same index 5, but this time with an empty tag.
+    let cqe = {
+        tags[0] = 0;
+
+        unsafe {
+            ring.submitter()
+                .register_buffers_update(5, &iovecs, Some(&tags))?;
+        }
+        ring.submit_and_wait(1)?;
+
+        ring.completion().next().unwrap().into()
+    };
+
+    // We should get a cqe with the first tag because we registered a
+    // new buffer at an index where a buffer was already registered.
+    if cqe.user_data() != BUFFER_TAG {
+        return Err(anyhow::anyhow!(
+            "expected completion queue to contain a buffer unregistered event"
+        ));
+    }
+
+    // Try reading now that the buffer is registered at index 5
+    let cqe = {
+        unsafe {
+            ring.submission().push(&read_sqe).unwrap();
+        }
+
+        ring.submit_and_wait(1)?;
+
+        ring.completion().next().unwrap().into()
+    };
+
+    // We should get the correct user_data
+    if cqe.user_data() != 42 {
+        return Err(anyhow::anyhow!("unexpected completion queue entry"));
+    }
+
+    // We should read exactly two bytes
+    if cqe.result() != 2 {
+        return Err(anyhow::anyhow!("unexpected read result: {}", cqe.result()));
+    }
+
+    // The first two bytes of `buf` should be "yo"
+    if &buf[0..2] != "yo".as_bytes() {
+        return Err(anyhow::anyhow!("unexpected read buffer data: {:x?}", &buf));
+    }
+
+    return Ok(());
+}
+
+/// Create a pipe and return both ends as RAII `File` handles
+fn create_pipe() -> io::Result<(File, File)> {
+    let mut fds = [-1, -1];
+
+    unsafe {
+        if libc::pipe(fds.as_mut_ptr()) == -1 {
+            Err(Error::last_os_error())
+        } else {
+            Ok((File::from_raw_fd(fds[0]), File::from_raw_fd(fds[1])))
+        }
+    }
+}
+
+/// Submit sqes and asserts the only cqe is a timeout entry
+fn check_only_timeout<S: squeue::EntryMarker, C: cqueue::EntryMarker>(
+    ring: &mut IoUring<S, C>,
+) -> Result<(), anyhow::Error> {
+    ring.submit_and_wait(1)?;
+
+    if Into::<cqueue::Entry>::into(ring.completion().next().unwrap()).user_data() == TIMEOUT_TAG {
+        // There should not be any more entries in the queue
+        if ring.completion().next().is_none() {
+            return Ok(());
+        }
+    }
+
+    Err(anyhow::anyhow!("unexpected completion queue entry"))
+}

src/lib.rs

@@ -354,7 +354,6 @@ impl<S: squeue::EntryMarker, C: cqueue::EntryMarker> Builder<S, C> {
     /// events. You are only able to [register restrictions](Submitter::register_restrictions) when
     /// the rings are disabled due to concurrency issues. You can enable the rings with
     /// [`Submitter::register_enable_rings`]. Available since 5.10.
-
     pub fn setup_r_disabled(&mut self) -> &mut Self {
         self.params.flags |= sys::IORING_SETUP_R_DISABLED;
         self

src/submit.rs

@@ -193,6 +193,108 @@ impl<'a> Submitter<'a> {
         .map(drop)
     }
 
+    /// Update a range of fixed buffers starting at `offset`.
+    ///
+    /// This is required to use buffers registered using
+    /// [`register_buffers_sparse`](Self::register_buffers_sparse),
+    /// although it can be also be used with [`register_buffers`](Self::register_buffers).
+    ///
+    /// See [`register_buffers2`](Self::register_buffers2)
+    /// for more information about resource tagging.
+    ///
+    /// Available since Linux 5.13.
+    ///
+    /// # Safety
+    ///
+    /// Developers must ensure that the `iov_base` and `iov_len` values are valid and will
+    /// be valid until buffers are unregistered or the ring destroyed, otherwise undefined
+    /// behaviour may occur.
+    pub unsafe fn register_buffers_update(
+        &self,
+        offset: u32,
+        bufs: &[libc::iovec],
+        tags: Option<&[u64]>,
+    ) -> io::Result<()> {
+        let nr = tags
+            .as_ref()
+            .map_or(bufs.len(), |tags| bufs.len().min(tags.len()));
+
+        let rr = sys::io_uring_rsrc_update2 {
+            nr: nr as _,
+            data: bufs.as_ptr() as _,
+            tags: tags.map(|tags| tags.as_ptr() as _).unwrap_or(0),
+            offset,
+            ..Default::default()
+        };
+
+        execute(
+            self.fd.as_raw_fd(),
+            sys::IORING_REGISTER_BUFFERS_UPDATE,
+            cast_ptr::<sys::io_uring_rsrc_update2>(&rr).cast(),
+            std::mem::size_of::<sys::io_uring_rsrc_update2>() as _,
+        )
+        .map(drop)
+    }
+
+    /// Variant of [`register_buffers`](Self::register_buffers)
+    /// with resource tagging.
+    ///
+    /// `tags` should be the same length as `bufs` and contain the
+    /// tag value corresponding to the buffer at the same index.
+    ///
+    /// If a tag is zero, then tagging for this particular resource
+    /// (a buffer in this case) is disabled. Otherwise, after the
+    /// resource had been unregistered and it's not used anymore,
+    /// a CQE will be posted with `user_data` set to the specified
+    /// tag and all other fields zeroed.
+    ///
+    /// Available since Linux 5.13.
+    ///
+    /// # Safety
+    ///
+    /// Developers must ensure that the `iov_base` and `iov_len` values are valid and will
+    /// be valid until buffers are unregistered or the ring destroyed, otherwise undefined
+    /// behaviour may occur.
+    pub unsafe fn register_buffers2(&self, bufs: &[libc::iovec], tags: &[u64]) -> io::Result<()> {
+        let rr = sys::io_uring_rsrc_register {
+            nr: bufs.len().min(tags.len()) as _,
+            data: bufs.as_ptr() as _,
+            tags: tags.as_ptr() as _,
+            ..Default::default()
+        };
+        execute(
+            self.fd.as_raw_fd(),
+            sys::IORING_REGISTER_BUFFERS2,
+            cast_ptr::<sys::io_uring_rsrc_register>(&rr).cast(),
+            std::mem::size_of::<sys::io_uring_rsrc_register>() as _,
+        )
+        .map(drop)
+    }
+
+    /// Registers an empty table of nr fixed buffers buffers.
+    ///
+    /// These must be updated before use, using eg.
+    /// [`register_buffers_update`](Self::register_buffers_update).
+    ///
+    /// See [`register_buffers`](Self::register_buffers)
+    /// for more information about fixed buffers.
+    ///
+    /// Available since Linux 5.13.
+    pub fn register_buffers_sparse(&self, nr: u32) -> io::Result<()> {
+        let rr = sys::io_uring_rsrc_register {
+            nr,
+            flags: sys::IORING_RSRC_REGISTER_SPARSE,
+            ..Default::default()
+        };
+        execute(
+            self.fd.as_raw_fd(),
+            sys::IORING_REGISTER_BUFFERS2,
+            cast_ptr::<sys::io_uring_rsrc_register>(&rr).cast(),
+            std::mem::size_of::<sys::io_uring_rsrc_register>() as _,
+        )
+        .map(drop)
+    }
+
     /// Registers an empty file table of nr_files number of file descriptors. The sparse variant is
     /// available in kernels 5.19 and later.
     ///
@@ -331,6 +433,8 @@ impl<'a> Submitter<'a> {
     ///
     /// You do not need to explicitly call this before dropping the [`IoUring`](crate::IoUring), as
     /// it will be cleaned up by the kernel automatically.
+    ///
+    /// Available since Linux 5.1.
     pub fn unregister_buffers(&self) -> io::Result<()> {
         execute(
             self.fd.as_raw_fd(),