Add missing register_files APIs

io-uring-test/src/main.rs

@@ -2,13 +2,28 @@
 mod utils;
 mod tests;
 
+use anyhow::Context;
 use io_uring::{cqueue, squeue, IoUring, Probe};
 use std::cell::Cell;
 
 pub struct Test {
     probe: Probe,
     target: Option<String>,
     count: Cell<usize>,
+    event_fd: libc::c_int,
+}
+
+impl Test {
+    /// Reset the eventfd counter, this can be used to prevent
+    /// previous operations affecting the test.
+    fn reset_eventfd_counter(&self) {
+        // Reset the event fd state without blocking.
+        unsafe {
+            let _ = libc::eventfd_write(self.event_fd, 1);
+            let mut val: u64 = 0;
+            let _ = libc::eventfd_read(self.event_fd, (&mut val) as *mut _);
+        };
+    }
 }
 
 fn main() -> anyhow::Result<()> {
@@ -59,8 +74,18 @@ fn test<S: squeue::EntryMarker, C: cqueue::EntryMarker>(
     println!("probe: {:?}", probe);
     println!();
 
+    // Used for waiting on events, this is more reliable for CI & general use cases
+    // than doing a full event loop.
+    // Since this is just for testing, we don't really need to close the event fd explicitly,
+    // it'll just be cleaned up at the end.
+    let event_fd = unsafe { libc::eventfd(0, 0) };
+    ring.submitter()
+        .register_eventfd(event_fd)
+        .context("register event fd")?;
+
     let test = Test {
         probe,
+        event_fd,
         target: std::env::args().nth(1),
         count: Cell::new(0),
     };
@@ -75,6 +100,8 @@ fn test<S: squeue::EntryMarker, C: cqueue::EntryMarker>(
 
     // register
     tests::register::test_register_files_sparse(&mut ring, &test)?;
+    tests::register::test_register_files_tags(&mut ring, &test)?;
+    tests::register::test_register_files_update_tag(&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)?;

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

@@ -1,5 +1,8 @@
 use crate::Test;
+use anyhow::{bail, Context};
+use io_uring::cqueue::Entry;
 use io_uring::{cqueue, opcode, squeue, IoUring};
+use std::os::fd::AsRawFd;
 
 pub fn test_register_files_sparse<S: squeue::EntryMarker, C: cqueue::EntryMarker>(
     ring: &mut IoUring<S, C>,
@@ -68,3 +71,146 @@ pub fn test_register_files_sparse<S: squeue::EntryMarker, C: cqueue::EntryMarker
 
     Ok(())
 }
+
+pub fn test_register_files_tags<S: squeue::EntryMarker, C: cqueue::EntryMarker>(
+    ring: &mut IoUring<S, C>,
+    test: &Test,
+) -> anyhow::Result<()> {
+    // register_files_sparse was introduced in kernel 5.13, we need to check if resource tagging
+    // is available.
+    require!(
+        test;
+        ring.params().is_feature_resource_tagging();
+    );
+
+    println!("test register_files_tags");
+    test.reset_eventfd_counter();
+
+    let (submitter, _, mut completion) = ring.split();
+
+    let file1 = tempfile::tempfile()?;
+    let file2 = tempfile::tempfile()?;
+    let file3 = tempfile::tempfile()?;
+
+    let descriptors = &[file1.as_raw_fd(), file2.as_raw_fd(), file3.as_raw_fd()];
+
+    let flags = &[0, 1, 2];
+
+    submitter
+        .register_files_tags(descriptors, flags)
+        .context("register_files_tags failed")?;
+
+    // See that same call again, with any value, will fail because a direct table cannot be built
+    // over an existing one.
+    if let Ok(()) = submitter.register_files_tags(descriptors, flags) {
+        bail!("register_files_tags should not have succeeded twice in a row");
+    }
+
+    // See that the direct table can be removed.
+    submitter
+        .unregister_files()
+        .context("unregister_files failed")?;
+
+    // See that a second attempt to remove the direct table would fail.
+    if let Ok(()) = submitter.unregister_files() {
+        bail!("should not have succeeded twice in a row");
+    }
+
+    // Wait for the event fd to complete indicating we have completions.
+    unsafe {
+        let mut value: u64 = 0;
+        let _ = libc::eventfd_read(test.event_fd, (&mut value) as *mut _);
+    }
+
+    // Check that we get completion events for unregistering files with non-zero tags.
+    completion.sync();
+
+    let cqes = completion.map(Into::into).collect::<Vec<Entry>>();
+
+    if cqes.len() != 2 {
+        bail!("incorrect number of completion events: {cqes:?}")
+    }
+
+    for event in cqes {
+        if event.flags() != 0 || event.result() != 0 {
+            bail!(
+                "completion events from unregistering tagged \
+                files should have zeroed flags and result fields"
+            );
+        }
+
+        let tag = event.user_data();
+        if !(1..3).contains(&tag) {
+            bail!("completion event user data does not contain one of the possible tag values")
+        }
+    }
+
+    Ok(())
+}
+
+pub fn test_register_files_update_tag<S: squeue::EntryMarker, C: cqueue::EntryMarker>(
+    ring: &mut IoUring<S, C>,
+    test: &Test,
+) -> anyhow::Result<()> {
+    // register_files_update_tag was introduced in kernel 5.13, we need to check if resource tagging
+    // is available.
+    require!(
+        test;
+        ring.params().is_feature_resource_tagging();
+    );
+
+    let (submitter, _, mut completion) = ring.split();
+
+    println!("test register_files_update_tag");
+    test.reset_eventfd_counter();
+
+    let descriptors = &[-1, -1, -1];
+    let flags = &[0, 0, 0];
+    submitter
+        .register_files_tags(descriptors, flags)
+        .context("register_files_tags failed")?;
+
+    let file = tempfile::tempfile().context("create temp file")?;
+
+    // Update slot `1` with a tagged file
+    submitter
+        .register_files_update_tag(1, &[file.as_raw_fd()], &[1])
+        .context("register_files_update_tag failed")?;
+
+    // Update slot `2` with an untagged file
+    submitter
+        .register_files_update_tag(1, &[file.as_raw_fd()], &[0])
+        .context("register_files_update_tag failed")?;
+
+    submitter
+        .unregister_files()
+        .context("unregister_files failed")?;
+
+    // Wait for the event fd to complete indicating we have completions.
+    unsafe {
+        let mut value: u64 = 0;
+        let _ = libc::eventfd_read(test.event_fd, (&mut value) as *mut _);
+    }
+
+    // Check that we get completion events for unregistering files with non-zero tags.
+    completion.sync();
+
+    let cqes = completion.map(Into::into).collect::<Vec<Entry>>();
+
+    if cqes.len() != 1 {
+        bail!("incorrect number of completion events: {cqes:?}")
+    }
+
+    if cqes[0].result() != 0 || cqes[0].flags() != 0 {
+        bail!(
+            "completion events from unregistering tagged \
+            files should have zeroed flags and result fields"
+        );
+    }
+
+    if cqes[0].user_data() != 1 {
+        bail!("completion event user data does not contain tag of registered file");
+    }
+
+    Ok(())
+}

src/submit.rs

@@ -323,8 +323,12 @@ impl<'a> Submitter<'a> {
     /// Each fd may be -1, in which case it is considered "sparse", and can be filled in later with
     /// [`register_files_update`](Self::register_files_update).
     ///
-    /// Note that this will wait for the ring to idle; it will only return once all active requests
-    /// are complete. Use [`register_files_update`](Self::register_files_update) to avoid this.
+    /// Note that before 5.13 registering buffers would wait for the ring to idle.  
+    /// If the application currently has requests in-flight, the registration will
+    /// wait for those to finish before proceeding.
+    ///
+    /// You can use [`register_files_update`](Self::register_files_update) to execute
+    /// this operation asynchronously.
     pub fn register_files(&self, fds: &[RawFd]) -> io::Result<()> {
         execute(
             self.fd.as_raw_fd(),
@@ -335,6 +339,35 @@ impl<'a> Submitter<'a> {
         .map(drop)
     }
 
+    /// Variant of [`register_files`](Self::register_files)
+    /// with resource tagging.
+    ///
+    /// Each fd may be -1, in which case it is considered "sparse", and can be filled in later with
+    /// [`register_files_update`](Self::register_files_update).
+    ///
+    /// `tags` should be the same length as `fds` and contain the
+    /// tag value corresponding to the file descriptor at the same index.
+    ///
+    /// See [`register_buffers2`](Self::register_buffers2)
+    /// for more information about resource tagging.
+    ///
+    /// Available since Linux 5.13.
+    pub fn register_files_tags(&self, fds: &[RawFd], tags: &[u64]) -> io::Result<()> {
+        let rr = sys::io_uring_rsrc_register {
+            nr: fds.len().min(tags.len()) as _,
+            data: fds.as_ptr() as _,
+            tags: tags.as_ptr() as _,
+            ..Default::default()
+        };
+        execute(
+            self.fd.as_raw_fd(),
+            sys::IORING_REGISTER_FILES2,
+            cast_ptr::<sys::io_uring_rsrc_register>(&rr).cast(),
+            mem::size_of::<sys::io_uring_rsrc_register>() as _,
+        )
+        .map(drop)
+    }
+
     /// This operation replaces existing files in the registered file set with new ones,
     /// either turning a sparse entry (one where fd is equal to -1) into a real one, removing an existing entry (new one is set to -1),
     /// or replacing an existing entry with a new existing entry. The `offset` parameter specifies
@@ -357,6 +390,38 @@ impl<'a> Submitter<'a> {
         Ok(ret as _)
     }
 
+    /// Variant of [`register_files_update`](Self::register_files_update)
+    /// with resource tagging.
+    ///
+    /// `tags` should be the same length as `fds` and contain the
+    /// tag value corresponding to the file descriptor at the same index.
+    ///
+    /// See [`register_buffers2`](Self::register_buffers2)
+    /// for more information about resource tagging.
+    ///
+    /// Available since Linux 5.13.
+    pub fn register_files_update_tag(
+        &self,
+        offset: u32,
+        fds: &[RawFd],
+        tags: &[u64],
+    ) -> io::Result<()> {
+        let rr = sys::io_uring_rsrc_update2 {
+            offset,
+            nr: fds.len().min(tags.len()) as _,
+            data: fds.as_ptr() as _,
+            tags: tags.as_ptr() as _,
+            ..Default::default()
+        };
+        execute(
+            self.fd.as_raw_fd(),
+            sys::IORING_REGISTER_FILES_UPDATE2,
+            cast_ptr::<sys::io_uring_rsrc_update2>(&rr).cast(),
+            mem::size_of::<sys::io_uring_rsrc_update2>() as _,
+        )
+        .map(drop)
+    }
+
     /// Register an eventfd created by [`eventfd`](libc::eventfd) with the io_uring instance.
     pub fn register_eventfd(&self, eventfd: RawFd) -> io::Result<()> {
         execute(