Add io_uring flags available in linux-raw-sys (#1115)

* Add io_uring flags available in linux-raw-sys

* Make IORING_REGISTER_USE_REGISTERED_RING const

* Add io_uring_register_with and IoringRegisterFlags

IORING_REGISTER_USE_REGISTERED_RING must be ored into `opcode` when
calling `io_uring_register` with a registered ring fd. The value of this
flag is 1u32 << 31, and does not fit into the `IoringRegisterOp` enum.
To support this use case, add a new function `io_uring_register_with`
that takes a bitflags flags argument.

Also adds a minimal test that calls `io_uring_register` with a
registered ring fd. This operation isn't supported on older kernels, so
we ensure the feature is supported before completing the test.

Author: ryanseipp

src/backend/libc/io_uring/syscalls.rs

@@ -4,7 +4,7 @@ use crate::backend::c;
 use crate::backend::conv::{borrowed_fd, ret_owned_fd, ret_u32};
 use crate::fd::{BorrowedFd, OwnedFd};
 use crate::io;
-use crate::io_uring::{io_uring_params, IoringEnterFlags, IoringRegisterOp};
+use crate::io_uring::{io_uring_params, IoringEnterFlags, IoringRegisterFlags, IoringRegisterOp};
 
 #[inline]
 pub(crate) fn io_uring_setup(entries: u32, params: &mut io_uring_params) -> io::Result<OwnedFd> {
@@ -40,6 +40,30 @@ pub(crate) unsafe fn io_uring_register(
     ))
 }
 
+#[inline]
+pub(crate) unsafe fn io_uring_register_with(
+    fd: BorrowedFd<'_>,
+    opcode: IoringRegisterOp,
+    flags: IoringRegisterFlags,
+    arg: *const c::c_void,
+    nr_args: u32,
+) -> io::Result<u32> {
+    syscall! {
+        fn io_uring_register(
+            fd: c::c_uint,
+            opcode: c::c_uint,
+            arg: *const c::c_void,
+            nr_args: c::c_uint
+        ) via SYS_io_uring_register -> c::c_int
+    }
+    ret_u32(io_uring_register(
+        borrowed_fd(fd) as _,
+        (opcode as u32) | bitflags_bits!(flags),
+        arg,
+        nr_args,
+    ))
+}
+
 #[inline]
 pub(crate) unsafe fn io_uring_enter(
     fd: BorrowedFd<'_>,

src/backend/linux_raw/io_uring/syscalls.rs

@@ -8,7 +8,7 @@
 use crate::backend::conv::{by_mut, c_uint, pass_usize, ret_c_uint, ret_owned_fd};
 use crate::fd::{BorrowedFd, OwnedFd};
 use crate::io;
-use crate::io_uring::{io_uring_params, IoringEnterFlags, IoringRegisterOp};
+use crate::io_uring::{io_uring_params, IoringEnterFlags, IoringRegisterFlags, IoringRegisterOp};
 use core::ffi::c_void;
 
 #[inline]
@@ -38,6 +38,23 @@ pub(crate) unsafe fn io_uring_register(
     ))
 }
 
+#[inline]
+pub(crate) unsafe fn io_uring_register_with(
+    fd: BorrowedFd<'_>,
+    opcode: IoringRegisterOp,
+    flags: IoringRegisterFlags,
+    arg: *const c_void,
+    nr_args: u32,
+) -> io::Result<u32> {
+    ret_c_uint(syscall_readonly!(
+        __NR_io_uring_register,
+        fd,
+        c_uint((opcode as u32) | bitflags_bits!(flags)),
+        arg,
+        c_uint(nr_args)
+    ))
+}
+
 #[inline]
 pub(crate) unsafe fn io_uring_enter(
     fd: BorrowedFd<'_>,

src/io_uring.rs

@@ -35,7 +35,9 @@ use linux_raw_sys::net;
 pub use crate::event::epoll::{
     Event as EpollEvent, EventData as EpollEventData, EventFlags as EpollEventFlags,
 };
-pub use crate::fs::{Advice, AtFlags, Mode, OFlags, RenameFlags, ResolveFlags, Statx, StatxFlags};
+pub use crate::fs::{
+    Advice, AtFlags, Mode, OFlags, RenameFlags, ResolveFlags, Statx, StatxFlags, XattrFlags,
+};
 pub use crate::io::ReadWriteFlags;
 pub use crate::net::{RecvFlags, SendFlags, SocketFlags};
 pub use crate::timespec::Timespec;
@@ -115,6 +117,30 @@ pub unsafe fn io_uring_register<Fd: AsFd>(
     backend::io_uring::syscalls::io_uring_register(fd.as_fd(), opcode, arg, nr_args)
 }
 
+/// `io_uring_register_with(fd, opcode, flags, arg, nr_args)`—Register files or
+/// user buffers for asynchronous I/O.
+///
+/// # Safety
+///
+/// io_uring operates on raw pointers and raw file descriptors. Users are
+/// responsible for ensuring that memory and resources are only accessed in
+/// valid ways.
+///
+/// # References
+///  - [Linux]
+///
+/// [Linux]: https://man.archlinux.org/man/io_uring_register.2.en
+#[inline]
+pub unsafe fn io_uring_register_with<Fd: AsFd>(
+    fd: Fd,
+    opcode: IoringRegisterOp,
+    flags: IoringRegisterFlags,
+    arg: *const c_void,
+    nr_args: u32,
+) -> io::Result<u32> {
+    backend::io_uring::syscalls::io_uring_register_with(fd.as_fd(), opcode, flags, arg, nr_args)
+}
+
 /// `io_uring_enter(fd, to_submit, min_complete, flags, arg, size)`—Initiate
 /// and/or complete asynchronous I/O.
 ///
@@ -260,6 +286,19 @@ pub enum IoringRegisterOp {
     RegisterFileAllocRange = sys::IORING_REGISTER_FILE_ALLOC_RANGE as _,
 }
 
+bitflags::bitflags! {
+    /// `IORING_REGISTER_*` flags for use with [`io_uring_register_with`].
+    #[repr(transparent)]
+    #[derive(Default, Copy, Clone, Eq, PartialEq, Hash, Debug)]
+    pub struct IoringRegisterFlags: u32 {
+        /// `IORING_REGISTER_USE_REGISTERED_RING`
+        const USE_REGISTERED_RING = sys::IORING_REGISTER_USE_REGISTERED_RING as u32;
+
+        /// <https://docs.rs/bitflags/*/bitflags/#externally-defined-flags>
+        const _ = !0;
+    }
+}
+
 /// `IORING_OP_*` constants for use with [`io_uring_sqe`].
 #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
 #[repr(u8)]
@@ -511,6 +550,9 @@ bitflags::bitflags! {
         /// `IORING_SETUP_REGISTERED_FD_ONLY`
         const REGISTERED_FD_ONLY = sys::IORING_SETUP_REGISTERED_FD_ONLY;
 
+        /// `IORING_SETUP_NO_SQARRAY`
+        const NO_SQARRAY = sys::IORING_SETUP_NO_SQARRAY;
+
         /// <https://docs.rs/bitflags/*/bitflags/#externally-defined-flags>
         const _ = !0;
     }
@@ -712,6 +754,9 @@ bitflags::bitflags! {
         /// `IORING_FEAT_LINKED_FILE`
         const LINKED_FILE = sys::IORING_FEAT_LINKED_FILE;
 
+        /// `IORING_FEAT_REG_REG_RING`
+        const REG_REG_RING = sys::IORING_FEAT_REG_REG_RING;
+
         /// <https://docs.rs/bitflags/*/bitflags/#externally-defined-flags>
         const _ = !0;
     }
@@ -1128,6 +1173,7 @@ pub union op_flags_union {
     pub rename_flags: RenameFlags,
     pub unlink_flags: AtFlags,
     pub hardlink_flags: AtFlags,
+    pub xattr_flags: XattrFlags,
     pub msg_ring_flags: IoringMsgringFlags,
 }
 

tests/io_uring/main.rs

@@ -0,0 +1,4 @@
+#[cfg(linux_kernel)]
+#[cfg(feature = "io_uring")]
+#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
+mod register;

tests/io_uring/register.rs

@@ -0,0 +1,108 @@
+use libc::c_void;
+use rustix::{
+    fd::{AsFd, AsRawFd, BorrowedFd},
+    io::Result,
+    io_uring::{
+        io_uring_params, io_uring_register_with, io_uring_rsrc_update, io_uring_setup,
+        IoringFeatureFlags, IoringRegisterFlags, IoringRegisterOp,
+    },
+};
+
+fn do_register<FD>(
+    fd: FD,
+    registered_fd: bool,
+    opcode: IoringRegisterOp,
+    arg: *const c_void,
+    arg_nr: u32,
+) -> Result<()>
+where
+    FD: AsFd,
+{
+    let flags = if registered_fd {
+        IoringRegisterFlags::USE_REGISTERED_RING
+    } else {
+        IoringRegisterFlags::default()
+    };
+
+    unsafe {
+        io_uring_register_with(fd, opcode, flags, arg, arg_nr)?;
+    }
+
+    Ok(())
+}
+
+fn register_ring<'a>(fd: BorrowedFd<'a>) -> Result<BorrowedFd<'a>> {
+    let update = io_uring_rsrc_update {
+        data: fd.as_raw_fd() as u64,
+        offset: u32::MAX,
+        resv: 0,
+    };
+
+    do_register(
+        fd,
+        false,
+        IoringRegisterOp::RegisterRingFds,
+        (&update) as *const io_uring_rsrc_update as *const c_void,
+        1,
+    )?;
+
+    let registered_fd = unsafe { BorrowedFd::borrow_raw(update.offset as i32) };
+    Ok(registered_fd)
+}
+
+fn unregister_ring<FD>(fd: FD) -> Result<()>
+where
+    FD: AsRawFd + AsFd,
+{
+    let update = io_uring_rsrc_update {
+        offset: fd.as_raw_fd() as u32,
+        data: 0,
+        resv: 0,
+    };
+
+    do_register(
+        fd,
+        true,
+        IoringRegisterOp::UnregisterRingFds,
+        (&update) as *const io_uring_rsrc_update as *const c_void,
+        1,
+    )?;
+
+    Ok(())
+}
+
+/// Set bounded and unbounded async kernel worker counts to 0, to test registering with registered
+/// ring fd.
+fn register_iowq_max_workers<FD>(fd: FD) -> Result<()>
+where
+    FD: AsFd,
+{
+    let iowq_max_workers = [0u32; 2];
+    do_register(
+        fd,
+        true,
+        IoringRegisterOp::RegisterIowqMaxWorkers,
+        (&iowq_max_workers) as *const [u32; 2] as *const c_void,
+        2,
+    )?;
+
+    Ok(())
+}
+
+#[test]
+fn test_io_uring_register_with() {
+    let mut params = io_uring_params::default();
+    let ring_fd = io_uring_setup(4, &mut params).unwrap();
+    assert_eq!(params.sq_entries, 4);
+    assert_eq!(params.cq_entries, 8);
+
+    if !params.features.contains(IoringFeatureFlags::REG_REG_RING) {
+        // Kernel does not support `io_uring_register` with a registered ring fd
+        return;
+    }
+
+    let ring_fd = register_ring(ring_fd.as_fd()).unwrap();
+    let register_result = register_iowq_max_workers(ring_fd);
+    let _ = unregister_ring(ring_fd).unwrap();
+    register_result.unwrap();
+}