feat(net): add recv_from_managed (#710)

Author: Sherlock-Holo

compio-net/src/socket.rs

@@ -10,9 +10,9 @@ use compio_driver::op::CreateSocket;
 use compio_driver::{
     AsRawFd, ToSharedFd, impl_raw_fd,
     op::{
-        Accept, BufResultExt, CloseSocket, Connect, Recv, RecvFrom, RecvFromVectored, RecvManaged,
-        RecvMsg, RecvResultExt, RecvVectored, ResultTakeBuffer, Send, SendMsg, SendTo,
-        SendToVectored, SendVectored, ShutdownSocket, VecBufResultExt,
+        Accept, BufResultExt, CloseSocket, Connect, Recv, RecvFrom, RecvFromManaged,
+        RecvFromVectored, RecvManaged, RecvMsg, RecvResultExt, RecvVectored, ResultTakeBuffer,
+        Send, SendMsg, SendTo, SendToVectored, SendVectored, ShutdownSocket, VecBufResultExt,
     },
     syscall,
 };
@@ -177,6 +177,21 @@ impl Socket {
             .take_buffer(buffer_pool)
     }
 
+    pub async fn recv_from_managed<'a>(
+        &self,
+        buffer_pool: &'a BufferPool,
+        len: usize,
+        flags: i32,
+    ) -> io::Result<(BorrowedBuffer<'a>, SockAddr)> {
+        let fd = self.to_shared_fd();
+        let buffer_pool = buffer_pool.try_inner()?;
+        let op = RecvFromManaged::new(fd, buffer_pool, len, flags)?;
+        compio_runtime::submit(op)
+            .with_extra()
+            .await
+            .take_buffer(buffer_pool)
+    }
+
     pub async fn send<T: IoBuf>(&self, buffer: T, flags: i32) -> BufResult<usize, T> {
         let fd = self.to_shared_fd();
         let op = Send::new(fd, buffer, flags);

compio-net/src/udp.rs

@@ -219,6 +219,22 @@ impl UdpSocket {
         self.inner.recv_managed(buffer_pool, len, 0).await
     }
 
+    /// Read some bytes from this source with [`BufferPool`] and return
+    /// a [`BorrowedBuffer`] with the sender address.
+    ///
+    /// If `len` == 0, will use [`BufferPool`] inner buffer size as the max len,
+    /// if `len` > 0, `min(len, inner buffer size)` will be the read max len
+    pub async fn recv_from_managed<'a>(
+        &self,
+        buffer_pool: &'a BufferPool,
+        len: usize,
+    ) -> io::Result<(BorrowedBuffer<'a>, SocketAddr)> {
+        self.inner
+            .recv_from_managed(buffer_pool, len, 0)
+            .await
+            .map(|(buffer, addr)| (buffer, addr.as_socket().expect("should be SocketAddr")))
+    }
+
     /// Sends some data to the socket from the buffer, returning the original
     /// buffer and quantity of data sent.
     pub async fn send<T: IoBuf>(&self, buffer: T) -> BufResult<usize, T> {

compio-net/tests/buffer_pool.rs

@@ -58,6 +58,24 @@ async fn test_udp_read_buffer_pool() {
     );
 }
 
+#[compio_macros::test]
+async fn test_udp_recv_from_buffer_pool() {
+    let listener = UdpSocket::bind((Ipv6Addr::LOCALHOST, 0)).await.unwrap();
+    let listener_addr = listener.local_addr().unwrap();
+    let connected = UdpSocket::bind((Ipv6Addr::LOCALHOST, 0)).await.unwrap();
+    let connected_addr = connected.local_addr().unwrap();
+
+    compio_runtime::spawn(async move {
+        connected.send_to(b"test", listener_addr).await.unwrap();
+    })
+    .detach();
+
+    let buffer_pool = BufferPool::new(1, 4).unwrap();
+    let (buffer, addr) = listener.recv_from_managed(&buffer_pool, 0).await.unwrap();
+    assert_eq!(buffer.as_ref(), b"test");
+    assert_eq!(addr, connected_addr);
+}
+
 #[compio_macros::test]
 async fn test_uds_recv_buffer_pool() {
     let dir = tempfile::Builder::new()