feat(driver): recvmsg multi (#842)

* feat(driver): recvmsg multi

* refactor(driver): hide impl details

Author: Berrysoft

compio-driver/src/op/managed.rs

@@ -1,10 +1,13 @@
 use std::io;
 
-use compio_buf::IntoInner;
+use compio_buf::{IntoInner, IoBuf, IoBufMut, SetLen};
 use socket2::SockAddr;
 
 use super::{Read, ReadAt, Recv, RecvFrom};
-use crate::{AsFd, BufferPool, BufferRef, op::TakeBuffer};
+use crate::{
+    AsFd, BufferPool, BufferRef,
+    op::{RecvMsg, TakeBuffer},
+};
 
 /// Read a file at specified position into managed buffer.
 pub struct ReadManagedAt<S> {
@@ -97,9 +100,158 @@ impl<S: AsFd> TakeBuffer for RecvFromManaged<S> {
     }
 }
 
+/// Receive data into managed buffer, and ancillary data into control buffer.
+pub struct RecvMsgManaged<C: IoBufMut, S: AsFd> {
+    pub(crate) op: RecvMsg<[BufferRef; 1], C, S>,
+}
+
+impl<C: IoBufMut, S: AsFd> RecvMsgManaged<C, S> {
+    /// Create [`RecvMsgManaged`].
+    pub fn new(fd: S, pool: &BufferPool, len: usize, control: C, flags: i32) -> io::Result<Self> {
+        Ok(Self {
+            op: RecvMsg::new(fd, [pool.pop()?.with_capacity(len)], control, flags),
+        })
+    }
+}
+
+impl<C: IoBufMut, S: AsFd> TakeBuffer for RecvMsgManaged<C, S> {
+    type Buffer = ((BufferRef, C), Option<SockAddr>, usize);
+
+    fn take_buffer(self) -> Option<Self::Buffer> {
+        let (([buf], control), addr, len) = self.op.into_inner();
+        Some(((buf, control), addr, len))
+    }
+}
+
 /// Read a file at specified position into multiple managed buffers.
 pub type ReadMultiAt<S> = ReadManagedAt<S>;
 /// Read a file into multiple managed buffers.
 pub type ReadMulti<S> = ReadManaged<S>;
 /// Receive data from remote into multiple managed buffers.
 pub type RecvMulti<S> = RecvManaged<S>;
+
+/// Result of [`RecvFromMulti`].
+pub struct RecvFromMultiResult {
+    buffer: BufferRef,
+    addr: Option<SockAddr>,
+}
+
+impl RecvFromMultiResult {
+    #[doc(hidden)]
+    pub unsafe fn new(_: BufferRef) -> Self {
+        unreachable!("should not be called directly")
+    }
+
+    /// Get the payload data.
+    pub fn data(&self) -> &[u8] {
+        self.buffer.as_init()
+    }
+
+    /// Get the source address if applicable.
+    pub fn addr(&self) -> Option<SockAddr> {
+        self.addr.clone()
+    }
+}
+
+impl IntoInner for RecvFromMultiResult {
+    type Inner = BufferRef;
+
+    fn into_inner(self) -> Self::Inner {
+        self.buffer
+    }
+}
+
+/// Receive data and source address multi times into multiple managed buffers.
+pub struct RecvFromMulti<S: AsFd> {
+    pub(crate) op: RecvFromManaged<S>,
+    pub(crate) len: usize,
+}
+
+impl<S: AsFd> RecvFromMulti<S> {
+    /// Create [`RecvFromMulti`].
+    pub fn new(fd: S, pool: &BufferPool, flags: i32) -> io::Result<Self> {
+        Ok(Self {
+            op: RecvFromManaged::new(fd, pool, 0, flags)?,
+            len: 0,
+        })
+    }
+}
+
+impl<S: AsFd> TakeBuffer for RecvFromMulti<S> {
+    type Buffer = RecvFromMultiResult;
+
+    fn take_buffer(self) -> Option<Self::Buffer> {
+        let (mut buffer, addr) = self.op.take_buffer()?;
+        unsafe { buffer.advance_to(self.len) };
+        Some(RecvFromMultiResult { buffer, addr })
+    }
+}
+
+/// Result of [`RecvMsgMulti`].
+pub struct RecvMsgMultiResult {
+    buffer: BufferRef,
+    control: BufferRef,
+    addr: Option<SockAddr>,
+}
+
+impl RecvMsgMultiResult {
+    #[doc(hidden)]
+    pub unsafe fn new(_: BufferRef, _: usize) -> Self {
+        unreachable!("should not be called directly")
+    }
+
+    /// Get the payload data.
+    pub fn data(&self) -> &[u8] {
+        self.buffer.as_init()
+    }
+
+    /// Get the source address if applicable.
+    pub fn addr(&self) -> Option<SockAddr> {
+        self.addr.clone()
+    }
+
+    /// Get the ancillary data.
+    pub fn ancillary(&self) -> &[u8] {
+        self.control.as_init()
+    }
+}
+
+impl IntoInner for RecvMsgMultiResult {
+    type Inner = BufferRef;
+
+    fn into_inner(self) -> Self::Inner {
+        self.buffer
+    }
+}
+
+/// Receive data, ancillary data and source address multi times into multiple
+/// managed buffers.
+pub struct RecvMsgMulti<S: AsFd> {
+    pub(crate) op: RecvMsgManaged<BufferRef, S>,
+    pub(crate) len: usize,
+}
+
+impl<S: AsFd> RecvMsgMulti<S> {
+    /// Create [`RecvMsgMulti`].
+    pub fn new(fd: S, pool: &BufferPool, control_len: usize, flags: i32) -> io::Result<Self> {
+        Ok(Self {
+            op: RecvMsgManaged::new(fd, pool, 0, pool.pop()?.with_capacity(control_len), flags)?,
+            len: 0,
+        })
+    }
+}
+
+impl<S: AsFd> TakeBuffer for RecvMsgMulti<S> {
+    type Buffer = RecvMsgMultiResult;
+
+    fn take_buffer(self) -> Option<Self::Buffer> {
+        let ((mut buffer, mut control), addr, control_len) = self.op.take_buffer()?;
+        unsafe { buffer.advance_to(self.len) };
+        unsafe { control.advance_to(control_len) };
+        Some(RecvMsgMultiResult {
+            buffer,
+            control,
+            addr,
+        })
+    }
+}

compio-driver/src/op/mod.rs

@@ -14,7 +14,8 @@ pub(crate) mod managed;
 
 #[cfg(not(io_uring))]
 pub use managed::{
-    ReadManaged, ReadManagedAt, ReadMulti, ReadMultiAt, RecvFromManaged, RecvManaged, RecvMulti,
+    ReadManaged, ReadManagedAt, ReadMulti, ReadMultiAt, RecvFromManaged, RecvFromMulti,
+    RecvFromMultiResult, RecvManaged, RecvMsgManaged, RecvMsgMulti, RecvMsgMultiResult, RecvMulti,
 };
 
 #[cfg(linux_all)]
@@ -33,7 +34,8 @@ pub use crate::sys::op::{
 pub use crate::sys::op::{ConnectNamedPipe, DeviceIoControl};
 #[cfg(io_uring)]
 pub use crate::sys::op::{
-    ReadManaged, ReadManagedAt, ReadMulti, ReadMultiAt, RecvFromManaged, RecvManaged, RecvMulti,
+    ReadManaged, ReadManagedAt, ReadMulti, ReadMultiAt, RecvFromManaged, RecvFromMulti,
+    RecvFromMultiResult, RecvManaged, RecvMsgManaged, RecvMsgMulti, RecvMsgMultiResult, RecvMulti,
 };
 use crate::{BufferRef, OwnedFd, SharedFd};
 

compio-driver/src/sys/fusion/op.rs

@@ -205,8 +205,8 @@ macro_rules! mop {
 
                 fn take_buffer(self) -> Option<$inner> {
                     match self.inner {
-                        [< $name Inner >]::IoUring(op) => op.take_buffer(),
-                        [< $name Inner >]::Poll(op) => op.take_buffer(),
+                        [< $name Inner >]::IoUring(op) => op.take_buffer().map(Into::into),
+                        [< $name Inner >]::Poll(op) => op.take_buffer().map(Into::into),
                     }
                 }
             }
@@ -272,6 +272,157 @@ mop!(<S: AsFd> ReadManagedAt(fd: S, offset: u64, pool: &BufferPool, len: usize)
 mop!(<S: AsFd> ReadManaged(fd: S, pool: &BufferPool, len: usize) with pool);
 mop!(<S: AsFd> RecvManaged(fd: S, pool: &BufferPool, len: usize, flags: i32) with pool);
 mop!(<S: AsFd> RecvFromManaged(fd: S, pool: &BufferPool, len: usize, flags: i32) with pool; (BufferRef, Option<SockAddr>));
+mop!(<C: IoBufMut, S: AsFd> RecvMsgManaged(fd: S, pool: &BufferPool, len: usize, control: C, flags: i32) with pool; ((BufferRef, C), Option<SockAddr>, usize));
 mop!(<S: AsFd> ReadMultiAt(fd: S, offset: u64, pool: &BufferPool, len: usize) with pool);
 mop!(<S: AsFd> ReadMulti(fd: S, pool: &BufferPool, len: usize) with pool);
 mop!(<S: AsFd> RecvMulti(fd: S, pool: &BufferPool, len: usize, flags: i32) with pool);
+mop!(<S: AsFd> RecvFromMulti(fd: S, pool: &BufferPool, flags: i32) with pool; RecvFromMultiResult);
+mop!(<S: AsFd> RecvMsgMulti(fd: S, pool: &BufferPool, control_len: usize, flags: i32) with pool; RecvMsgMultiResult);
+
+enum RecvFromMultiResultInner {
+    Poll(crate::op::managed::RecvFromMultiResult),
+    IoUring(iour::RecvFromMultiResult),
+}
+
+/// Result of [`RecvFromMulti`].
+pub struct RecvFromMultiResult {
+    inner: RecvFromMultiResultInner,
+}
+
+impl From<crate::op::managed::RecvFromMultiResult> for RecvFromMultiResult {
+    fn from(result: crate::op::managed::RecvFromMultiResult) -> Self {
+        Self {
+            inner: RecvFromMultiResultInner::Poll(result),
+        }
+    }
+}
+
+impl From<iour::RecvFromMultiResult> for RecvFromMultiResult {
+    fn from(result: iour::RecvFromMultiResult) -> Self {
+        Self {
+            inner: RecvFromMultiResultInner::IoUring(result),
+        }
+    }
+}
+
+impl RecvFromMultiResult {
+    /// Create [`RecvFromMultiResult`] from a buffer received from
+    /// [`RecvFromMulti`]. It should be used for io-uring only.
+    ///
+    /// # Safety
+    ///
+    /// The buffer must be received from [`RecvFromMulti`] or have the same
+    /// format as the buffer received from [`RecvFromMulti`].
+    pub unsafe fn new(buffer: BufferRef) -> Self {
+        Self {
+            inner: RecvFromMultiResultInner::IoUring(unsafe {
+                iour::RecvFromMultiResult::new(buffer)
+            }),
+        }
+    }
+
+    /// Get the payload data.
+    pub fn data(&self) -> &[u8] {
+        match &self.inner {
+            RecvFromMultiResultInner::Poll(result) => result.data(),
+            RecvFromMultiResultInner::IoUring(result) => result.data(),
+        }
+    }
+
+    /// Get the source address if applicable.
+    pub fn addr(&self) -> Option<SockAddr> {
+        match &self.inner {
+            RecvFromMultiResultInner::Poll(result) => result.addr(),
+            RecvFromMultiResultInner::IoUring(result) => result.addr(),
+        }
+    }
+}
+
+impl IntoInner for RecvFromMultiResult {
+    type Inner = BufferRef;
+
+    fn into_inner(self) -> Self::Inner {
+        match self.inner {
+            RecvFromMultiResultInner::Poll(result) => result.into_inner(),
+            RecvFromMultiResultInner::IoUring(result) => result.into_inner(),
+        }
+    }
+}
+
+enum RecvMsgMultiResultInner {
+    Poll(crate::op::managed::RecvMsgMultiResult),
+    IoUring(iour::RecvMsgMultiResult),
+}
+
+/// Result of [`RecvMsgMulti`].
+pub struct RecvMsgMultiResult {
+    inner: RecvMsgMultiResultInner,
+}
+
+impl From<crate::op::managed::RecvMsgMultiResult> for RecvMsgMultiResult {
+    fn from(result: crate::op::managed::RecvMsgMultiResult) -> Self {
+        Self {
+            inner: RecvMsgMultiResultInner::Poll(result),
+        }
+    }
+}
+
+impl From<iour::RecvMsgMultiResult> for RecvMsgMultiResult {
+    fn from(result: iour::RecvMsgMultiResult) -> Self {
+        Self {
+            inner: RecvMsgMultiResultInner::IoUring(result),
+        }
+    }
+}
+
+impl RecvMsgMultiResult {
+    /// Create [`RecvMsgMultiResult`] from a buffer received from
+    /// [`RecvMsgMulti`]. It should be used for io-uring only.
+    ///
+    /// # Safety
+    ///
+    /// The buffer must be received from [`RecvMsgMulti`] or have the same
+    /// format as the buffer received from [`RecvMsgMulti`].
+    pub unsafe fn new(buffer: BufferRef, clen: usize) -> Self {
+        Self {
+            inner: RecvMsgMultiResultInner::IoUring(unsafe {
+                iour::RecvMsgMultiResult::new(buffer, clen)
+            }),
+        }
+    }
+
+    /// Get the payload data.
+    pub fn data(&self) -> &[u8] {
+        match &self.inner {
+            RecvMsgMultiResultInner::Poll(result) => result.data(),
+            RecvMsgMultiResultInner::IoUring(result) => result.data(),
+        }
+    }
+
+    /// Get the ancillary data.
+    pub fn ancillary(&self) -> &[u8] {
+        match &self.inner {
+            RecvMsgMultiResultInner::Poll(result) => result.ancillary(),
+            RecvMsgMultiResultInner::IoUring(result) => result.ancillary(),
+        }
+    }
+
+    /// Get the source address if applicable.
+    pub fn addr(&self) -> Option<SockAddr> {
+        match &self.inner {
+            RecvMsgMultiResultInner::Poll(result) => result.addr(),
+            RecvMsgMultiResultInner::IoUring(result) => result.addr(),
+        }
+    }
+}
+
+impl IntoInner for RecvMsgMultiResult {
+    type Inner = BufferRef;
+
+    fn into_inner(self) -> Self::Inner {
+        match self.inner {
+            RecvMsgMultiResultInner::Poll(result) => result.into_inner(),
+            RecvMsgMultiResultInner::IoUring(result) => result.into_inner(),
+        }
+    }
+}