Uring cmd

Cargo.toml

@@ -16,14 +16,22 @@ keywords = ["async", "fs", "io-uring"]
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
+[features]
+# Enable 128-byte submission queue entries. This requires Linux >= 5.19.
+sqe128 = []
+# Enable 32-byte completion queue entries. This requires Linux >= 5.19.
+cqe32 = []
+
 [dependencies]
 tokio = { version = "1.2", features = ["net", "rt", "sync"] }
 slab = "0.4.2"
 libc = "0.2.80"
-io-uring = "0.5.13"
+io-uring = "0.6.0"
 socket2 = { version = "0.4.4", features = ["all"] }
 bytes = { version = "1.0", optional = true }
-futures-util = { version = "0.3.26", default-features = false, features = ["std"] }
+futures-util = { version = "0.3.26", default-features = false, features = [
+    "std",
+] }
 
 [dev-dependencies]
 tempfile = "3.2.0"

src/buf/fixed/handle.rs

@@ -82,9 +82,7 @@ unsafe impl IoBufMut for FixedBuf {
     }
 
     unsafe fn set_init(&mut self, pos: usize) {
-        if self.buf.init_len < pos {
-            self.buf.init_len = pos
-        }
+        self.buf.init_len = pos
     }
 }
 

src/buf/io_buf_mut.rs

@@ -40,9 +40,7 @@ unsafe impl IoBufMut for Vec<u8> {
     }
 
     unsafe fn set_init(&mut self, init_len: usize) {
-        if self.len() < init_len {
-            self.set_len(init_len);
-        }
+        self.set_len(init_len);
     }
 }
 
@@ -53,8 +51,6 @@ unsafe impl IoBufMut for bytes::BytesMut {
     }
 
     unsafe fn set_init(&mut self, init_len: usize) {
-        if self.len() < init_len {
-            self.set_len(init_len);
-        }
+        self.set_len(init_len);
     }
 }

src/fs/file.rs

@@ -182,6 +182,27 @@ impl File {
         op.await
     }
 
+    /// A file/device-specific 16-byte command, akin (but not equivalent) to ioctl(2).
+    ///
+    /// Marked as unsafe, as it is the callers responsibility to ensure that any data which is required
+    /// to be stable is either passed in as the metadata argument, or otherwise stable through the life of the operation
+    pub async unsafe fn uring_cmd16<T: std::marker::Unpin>(
+        &self,
+        cmd_op: u32,
+        cmd: [u8; 16],
+        metadata: T,
+    ) -> (io::Result<u32>, T) {
+        let op = Op::uring_cmd16(&self.fd, cmd_op, cmd, metadata).unwrap();
+        op.await
+    }
+
+    #[cfg(feature = "sqe128")]
+    /// A file/device-specific 80-byte command, akin (but not equivalent) to ioctl(2).
+    pub async fn uring_cmd80(&self, cmd_op: u32, cmd: [u8; 80]) -> io::Result<u32> {
+        let op = Op::uring_cmd80(&self.fd, cmd_op, cmd).unwrap();
+        op.await
+    }
+
     /// Read some bytes at the specified offset from the file into the specified
     /// array of buffers, returning how many bytes were read.
     ///

src/io/fallocate.rs

@@ -25,8 +25,8 @@ impl Op<Fallocate> {
             x.handle().expect("not in a runtime context").submit_op(
                 Fallocate { fd: fd.clone() },
                 |fallocate| {
-                    opcode::Fallocate64::new(types::Fd(fallocate.fd.raw_fd()), len as _)
-                        .offset64(offset as _)
+                    opcode::Fallocate::new(types::Fd(fallocate.fd.raw_fd()), len as _)
+                        .offset(offset as _)
                         .mode(flags)
                         .build()
                 },

src/io/mod.rs

@@ -56,3 +56,5 @@ mod writev;
 
 mod writev_all;
 pub(crate) use writev_all::writev_at_all;
+
+mod uring_cmd;

src/io/send_zc.rs

@@ -53,6 +53,6 @@ impl<T> Completable for SendZc<T> {
 impl<T> Updateable for SendZc<T> {
     fn update(&mut self, cqe: CqeResult) {
         // uring send_zc promises there will be no error on CQE's marked more
-        self.bytes += *cqe.result.as_ref().unwrap() as usize;
+        self.bytes += cqe.result.unwrap_or_default() as usize;
     }
 }

src/io/sendmsg_zc.rs

@@ -113,6 +113,6 @@ impl<T, U> Completable for SendMsgZc<T, U> {
 impl<T, U> Updateable for SendMsgZc<T, U> {
     fn update(&mut self, cqe: CqeResult) {
         // uring send_zc promises there will be no error on CQE's marked more
-        self.bytes += *cqe.result.as_ref().unwrap() as usize;
+        self.bytes += cqe.result.unwrap_or_default() as usize;
     }
 }

src/io/uring_cmd.rs

@@ -0,0 +1,87 @@
+use crate::io::SharedFd;
+use crate::runtime::driver::op::{Completable, CqeResult, Op};
+use crate::runtime::CONTEXT;
+use std::io;
+
+pub(crate) struct UringCmd16<T> {
+    /// Holds a strong ref to the FD, preventing the file from being closed
+    /// while the operation is in-flight.
+    #[allow(dead_code)]
+    fd: SharedFd,
+    // Holds any required metadata stable for the lifetime of the operation
+    metadata: T,
+}
+
+impl<T> Op<UringCmd16<T>> {
+    /// A file/device-specific 16-byte command, akin (but not equivalent) to ioctl
+    pub(crate) fn uring_cmd16(
+        fd: &SharedFd,
+        cmd_op: u32,
+        cmd: [u8; 16],
+        metadata: T,
+    ) -> io::Result<Op<UringCmd16<T>>> {
+        use io_uring::{opcode, types};
+
+        CONTEXT.with(|x| {
+            x.handle().expect("Not in a runtime context").submit_op(
+                UringCmd16 {
+                    fd: fd.clone(),
+                    metadata,
+                },
+                |_| {
+                    opcode::UringCmd16::new(types::Fd(fd.raw_fd()), cmd_op)
+                        .cmd(cmd)
+                        .build()
+                },
+            )
+        })
+    }
+}
+
+impl<T> Completable for UringCmd16<T> {
+    type Output = (std::result::Result<u32, io::Error>, T);
+
+    fn complete(self, cqe: CqeResult) -> Self::Output {
+        (cqe.result, self.metadata)
+    }
+}
+
+#[cfg(feature = "sqe128")]
+pub(crate) struct UringCmd80 {
+    /// Holds a strong ref to the FD, preventing the file from being closed
+    /// while the operation is in-flight.
+    #[allow(dead_code)]
+    fd: SharedFd,
+}
+
+#[cfg(feature = "sqe128")]
+impl Op<UringCmd80> {
+    /// A file/device-specific 80-byte command, akin (but not equivalent) to ioctl
+    pub(crate) fn uring_cmd80(
+        fd: &SharedFd,
+        cmd_op: u32,
+        cmd: [u8; 80],
+    ) -> io::Result<Op<UringCmd80>> {
+        use io_uring::{opcode, types};
+
+        CONTEXT.with(|x| {
+            x.handle().expect("Not in a runtime context").submit_op(
+                UringCmd80 { fd: fd.clone() },
+                |_| {
+                    opcode::UringCmd80::new(types::Fd(fd.raw_fd()), cmd_op)
+                        .cmd(cmd)
+                        .build()
+                },
+            )
+        })
+    }
+}
+
+#[cfg(feature = "sqe128")]
+impl Completable for UringCmd80 {
+    type Output = io::Result<u32>;
+
+    fn complete(self, cqe: CqeResult) -> Self::Output {
+        cqe.result
+    }
+}

src/io/writev_all.rs

@@ -134,7 +134,7 @@ impl<T: BoundedBuf> Op<WritevAll<T>> {
                 // So this wouldn't need to be a function. Just pass in the entry.
                 |write| {
                     opcode::Writev::new(types::Fd(write.fd.raw_fd()), iovs_ptr, iovs_len)
-                        .offset64(offset as _)
+                        .offset(offset as _)
                         .build()
                 },
             )

src/lib.rs

@@ -84,6 +84,7 @@ pub use runtime::spawn;
 pub use runtime::Runtime;
 
 use crate::runtime::driver::op::Op;
+use crate::runtime::driver::{CEntry, SEntry};
 use std::future::Future;
 
 /// Starts an `io_uring` enabled Tokio runtime.
@@ -154,7 +155,7 @@ pub fn start<F: Future>(future: F) -> F::Output {
 /// This function is provided to avoid requiring the user of this crate from
 /// having to use the io_uring crate as well. Refer to Builder::start example
 /// for its intended usage.
-pub fn uring_builder() -> io_uring::Builder {
+pub fn uring_builder() -> io_uring::Builder<SEntry, CEntry> {
     io_uring::IoUring::builder()
 }
 
@@ -163,7 +164,8 @@ pub fn uring_builder() -> io_uring::Builder {
 // #[derive(Clone, Default)]
 pub struct Builder {
     entries: u32,
-    urb: io_uring::Builder,
+    max_workers: [u32; 2],
+    urb: io_uring::Builder<SEntry, CEntry>,
 }
 
 /// Constructs a [`Builder`] with default settings.
@@ -175,6 +177,7 @@ pub struct Builder {
 pub fn builder() -> Builder {
     Builder {
         entries: 256,
+        max_workers: [0; 2],
         urb: io_uring::IoUring::builder(),
     }
 }
@@ -191,11 +194,23 @@ impl Builder {
         self
     }
 
+    /// Get and/or set the limit for number of io_uring worker threads per NUMA
+    /// node. `bounded` holds the limit for bounded workers, which process I/O
+    /// operations expected to be bound in time, that is I/O on regular files or
+    /// block devices. While `unbounded` holds the limit for unbounded workers,
+    /// which carry out I/O operations that can never complete, for instance I/O
+    /// on sockets. Setting `None` leaves the default
+    pub fn max_workers(&mut self, bounded: Option<u32>, unbounded: Option<u32>) -> &mut Self {
+        self.max_workers = [bounded.unwrap_or_default(), unbounded.unwrap_or_default()];
+        self
+    }
+
     /// Replaces the default [`io_uring::Builder`], which controls the settings for the
     /// inner `io_uring` API.
     ///
-    /// Refer to the [`io_uring::Builder`] documentation for all the supported methods.
-    pub fn uring_builder(&mut self, b: &io_uring::Builder) -> &mut Self {
+    /// Refer to the Builder start method for an example.
+    /// Refer to the io_uring::builder documentation for all the supported methods.
+    pub fn uring_builder(&mut self, b: &io_uring::Builder<SEntry, CEntry>) -> &mut Self {
         self.urb = b.clone();
         self
     }

src/runtime/driver/handle.rs

@@ -22,7 +22,7 @@ use std::task::{Context, Poll};
 
 use crate::buf::fixed::FixedBuffers;
 use crate::runtime::driver::op::{Completable, MultiCQEFuture, Op, Updateable};
-use crate::runtime::driver::Driver;
+use crate::runtime::driver::{Driver, SEntry};
 
 #[derive(Clone)]
 pub(crate) struct Handle {
@@ -67,10 +67,11 @@ impl Handle {
         self.inner.borrow_mut().submit_op_2(sqe)
     }
 
-    pub(crate) fn submit_op<T, S, F>(&self, data: T, f: F) -> io::Result<Op<T, S>>
+    pub(crate) fn submit_op<T, S, F, A>(&self, data: T, f: F) -> io::Result<Op<T, S>>
     where
         T: Completable,
-        F: FnOnce(&mut T) -> squeue::Entry,
+        A: Into<SEntry>,
+        F: FnOnce(&mut T) -> A,
     {
         self.inner.borrow_mut().submit_op(data, f, self.into())
     }