io-uring: add sendmsg

src/io_uring/ops/sendmsg.rs

@@ -0,0 +1,130 @@
+use {
+    crate::{
+        io_uring::{
+            pending_result::PendingResult,
+            sys::{io_uring_sqe, IORING_OP_SENDMSG, IOSQE_IO_LINK},
+            IoUring, IoUringData, IoUringError, Task,
+        },
+        time::Time,
+        utils::{buf::Buf, vec_ext::UninitVecExt},
+    },
+    std::{
+        mem::{self, MaybeUninit},
+        ptr,
+        rc::Rc,
+    },
+    uapi::{c, OwnedFd},
+};
+
+impl IoUring {
+    pub async fn sendmsg(
+        &self,
+        fd: &Rc<OwnedFd>,
+        buf: Buf,
+        fds: Vec<Rc<OwnedFd>>,
+        timeout: Option<Time>,
+    ) -> Result<usize, IoUringError> {
+        self.ring.check_destroyed()?;
+        let id = self.ring.id();
+        let pr = self.ring.pending_results.acquire();
+        {
+            let mut st = self.ring.cached_sendmsg.pop().unwrap_or_default();
+            st.fds = fds;
+            if st.fds.len() > 0 {
+                let mut fd_ids = self.ring.fd_ids_scratch.borrow_mut();
+                fd_ids.clear();
+                fd_ids.extend(st.fds.iter().map(|f| f.raw()));
+                let space = uapi::cmsg_space(mem::size_of_val(&fd_ids[..]));
+                st.cmsg.clear();
+                st.cmsg.reserve(space);
+                st.cmsg.set_len_safe(space);
+                let hdr = c::cmsghdr {
+                    cmsg_len: 0,
+                    cmsg_level: c::SOL_SOCKET,
+                    cmsg_type: c::SCM_RIGHTS,
+                };
+                uapi::cmsg_write(&mut &mut st.cmsg[..], hdr, &fd_ids[..]).unwrap();
+                st.msghdr.msg_control = st.cmsg.as_ptr() as _;
+                st.msghdr.msg_controllen = st.cmsg.len() as _;
+            } else {
+                st.msghdr.msg_control = ptr::null_mut();
+                st.msghdr.msg_controllen = 0;
+            }
+            st.id = id.id;
+            st.fd = fd.raw();
+            st.iovec.iov_base = buf.as_ptr() as _;
+            st.iovec.iov_len = buf.len() as _;
+            st.msghdr.msg_iov = &st.iovec as *const _ as _;
+            st.msghdr.msg_iovlen = 1;
+            st.data = Some(SendmsgTaskData {
+                _fd: fd.clone(),
+                _buf: buf,
+                res: pr.clone(),
+            });
+            st.has_timeout = timeout.is_some();
+            self.ring.schedule(st);
+            if let Some(timeout) = timeout {
+                self.schedule_timeout(timeout);
+            }
+        }
+        Ok(pr.await? as _)
+    }
+}
+
+struct SendmsgTaskData {
+    _fd: Rc<OwnedFd>,
+    _buf: Buf,
+    res: PendingResult,
+}
+
+pub struct SendmsgTask {
+    id: u64,
+    iovec: c::iovec,
+    msghdr: c::msghdr,
+    fd: i32,
+    has_timeout: bool,
+    fds: Vec<Rc<OwnedFd>>,
+    cmsg: Vec<MaybeUninit<u8>>,
+    data: Option<SendmsgTaskData>,
+}
+
+impl Default for SendmsgTask {
+    fn default() -> Self {
+        unsafe {
+            SendmsgTask {
+                id: 0,
+                iovec: MaybeUninit::zeroed().assume_init(),
+                msghdr: MaybeUninit::zeroed().assume_init(),
+                fd: 0,
+                has_timeout: false,
+                fds: vec![],
+                cmsg: vec![],
+                data: None,
+            }
+        }
+    }
+}
+
+unsafe impl Task for SendmsgTask {
+    fn id(&self) -> u64 {
+        self.id
+    }
+
+    fn complete(mut self: Box<Self>, ring: &IoUringData, res: i32) {
+        self.fds.clear();
+        if let Some(data) = self.data.take() {
+            data.res.complete(res);
+        }
+        ring.cached_sendmsg.push(self);
+    }
+
+    fn encode(&self, sqe: &mut io_uring_sqe) {
+        sqe.opcode = IORING_OP_SENDMSG;
+        sqe.fd = self.fd;
+        sqe.u2.addr = &self.msghdr as *const _ as _;
+        sqe.u3.msg_flags = c::MSG_NOSIGNAL as _;
+        if self.has_timeout {
+            sqe.flags = IOSQE_IO_LINK;
+        }
+    }
+}

src/io_uring/ops/timeout.rs

@@ -0,0 +1,54 @@
+use {
+    crate::{
+        io_uring::{
+            sys::{io_uring_sqe, IORING_OP_LINK_TIMEOUT, IORING_TIMEOUT_ABS},
+            IoUring, IoUringData, Task,
+        },
+        time::Time,
+    },
+    uapi::c,
+};
+
+#[allow(non_camel_case_types)]
+#[repr(C)]
+#[derive(Default)]
+struct timespec64 {
+    tv_sec: i64,
+    tv_nsec: c::c_long,
+}
+
+#[derive(Default)]
+pub struct TimeoutTask {
+    id: u64,
+    timespec: timespec64,
+}
+
+impl IoUring {
+    pub(super) fn schedule_timeout(&self, timeout: Time) {
+        let id = self.ring.id_raw();
+        {
+            let mut to = self.ring.cached_timeouts.pop().unwrap_or_default();
+            to.id = id;
+            to.timespec.tv_sec = timeout.0.tv_sec as _;
+            to.timespec.tv_nsec = timeout.0.tv_nsec as _;
+            self.ring.schedule(to);
+        }
+    }
+}
+
+unsafe impl Task for TimeoutTask {
+    fn id(&self) -> u64 {
+        self.id
+    }
+
+    fn complete(self: Box<Self>, ring: &IoUringData, _res: i32) {
+        ring.cached_timeouts.push(self);
+    }
+
+    fn encode(&self, sqe: &mut io_uring_sqe) {
+        sqe.opcode = IORING_OP_LINK_TIMEOUT;
+        sqe.u2.addr = &self.timespec as *const _ as _;
+        sqe.len = 1;
+        sqe.u3.timeout_flags = IORING_TIMEOUT_ABS;
+    }
+}

src/io_uring/ops/write.rs

@@ -1,26 +1,22 @@
 use {
-    crate::io_uring::{
-        ops::TaskResult,
-        pending_result::PendingResult,
-        sys::{io_uring_sqe, IORING_OP_WRITE},
-        IoUring, IoUringData, Task,
+    crate::{
+        io_uring::{
+            ops::TaskResult,
+            pending_result::PendingResult,
+            sys::{io_uring_sqe, IORING_OP_WRITE},
+            IoUring, IoUringData, Task,
+        },
+        utils::buf::Buf,
     },
     std::{
         cell::{Cell, RefCell},
-        ops::Range,
         rc::Rc,
     },
     uapi::OwnedFd,
 };
 
 impl IoUring {
-    pub async fn write(
-        &self,
-        fd: &Rc<OwnedFd>,
-        buf: &Rc<Vec<u8>>,
-        offset: usize,
-        n: usize,
-    ) -> TaskResult<usize> {
+    pub async fn write(&self, fd: &Rc<OwnedFd>, buf: Buf) -> TaskResult<usize> {
         self.ring.check_destroyed()?;
         let id = self.ring.id();
         let pr = self.ring.pending_results.acquire();
@@ -34,8 +30,7 @@ impl IoUring {
             pw.id.set(id.id);
             *pw.data.borrow_mut() = Some(WriteTaskData {
                 fd: fd.clone(),
-                buf: buf.clone(),
-                range: offset..offset + n,
+                buf,
                 res: pr.clone(),
             });
             self.ring.schedule(pw);
@@ -46,8 +41,7 @@ impl IoUring {
 
 struct WriteTaskData {
     fd: Rc<OwnedFd>,
-    buf: Rc<Vec<u8>>,
-    range: Range<usize>,
+    buf: Buf,
     res: PendingResult,
 }
 
@@ -74,8 +68,8 @@ unsafe impl Task for WriteTask {
         sqe.opcode = IORING_OP_WRITE;
         sqe.fd = data.fd.raw();
         sqe.u1.off = !0;
-        sqe.u2.addr = data.buf[data.range.clone()].as_ptr() as _;
+        sqe.u2.addr = data.buf.as_ptr() as _;
         sqe.u3.rw_flags = 0;
-        sqe.len = data.range.len() as _;
+        sqe.len = data.buf.len() as _;
     }
 }