Convert RecvFrom and SentTo to work with IntoSlice

tokio-rs · Mar 21, 2022 · 8b74f9e · 8b74f9e
1 parent cbdc2b2
commit 8b74f9e
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 22 deletions.
diff --git a/examples/udp_socket.rs b/examples/udp_socket.rs
@@ -15,8 +15,9 @@ fn main() {
 
         let buf = vec![0u8; 128];
 
-        let (result, mut buf) = socket.recv_from(buf).await;
+        let (result, buf) = socket.recv_from(buf).await;
         let (read, socket_addr) = result.unwrap();
+        let mut buf = buf.into_inner();
         buf.resize(read, 0);
         println!("received from {}: {:?}", socket_addr, &buf[..]);
 

diff --git a/src/driver/recv_from.rs b/src/driver/recv_from.rs
@@ -1,5 +1,5 @@
 use crate::{
-    buf::IoBufMut,
+    buf::{IoBufMut, Slice},
     driver::{Op, SharedFd},
     BufResult,
 };
@@ -13,14 +13,14 @@ use std::{
 #[allow(dead_code)]
 pub(crate) struct RecvFrom<T> {
     fd: SharedFd,
-    pub(crate) buf: T,
+    buf: Slice<T>,
     io_slices: Vec<IoSliceMut<'static>>,
     pub(crate) socket_addr: Box<SockAddr>,
     pub(crate) msghdr: Box<libc::msghdr>,
 }
 
 impl<T: IoBufMut> Op<RecvFrom<T>> {
-    pub(crate) fn recv_from(fd: &SharedFd, mut buf: T) -> io::Result<Op<RecvFrom<T>>> {
+    pub(crate) fn recv_from(fd: &SharedFd, mut buf: Slice<T>) -> io::Result<Op<RecvFrom<T>>> {
         use io_uring::{opcode, types};
 
         let mut io_slices = vec![IoSliceMut::new(unsafe {
@@ -53,7 +53,7 @@ impl<T: IoBufMut> Op<RecvFrom<T>> {
         )
     }
 
-    pub(crate) async fn recv(mut self) -> BufResult<(usize, SocketAddr), T> {
+    pub(crate) async fn recv(mut self) -> BufResult<(usize, SocketAddr), Slice<T>> {
         use crate::future::poll_fn;
 
         poll_fn(move |cx| self.poll_recv_from(cx)).await
@@ -62,7 +62,7 @@ impl<T: IoBufMut> Op<RecvFrom<T>> {
     pub(crate) fn poll_recv_from(
         &mut self,
         cx: &mut Context<'_>,
-    ) -> Poll<BufResult<(usize, SocketAddr), T>> {
+    ) -> Poll<BufResult<(usize, SocketAddr), Slice<T>>> {
         use std::future::Future;
         use std::pin::Pin;
 

diff --git a/src/driver/send_to.rs b/src/driver/send_to.rs
@@ -1,4 +1,4 @@
-use crate::buf::IoBuf;
+use crate::buf::{IoBuf, Slice};
 use crate::driver::{Op, SharedFd};
 use crate::BufResult;
 use socket2::SockAddr;
@@ -9,7 +9,7 @@ use std::{boxed::Box, io, net::SocketAddr};
 pub(crate) struct SendTo<T> {
     #[allow(dead_code)]
     fd: SharedFd,
-    pub(crate) buf: T,
+    buf: Slice<T>,
     #[allow(dead_code)]
     io_slices: Vec<IoSlice<'static>>,
     #[allow(dead_code)]
@@ -20,7 +20,7 @@ pub(crate) struct SendTo<T> {
 impl<T: IoBuf> Op<SendTo<T>> {
     pub(crate) fn send_to(
         fd: &SharedFd,
-        buf: T,
+        buf: Slice<T>,
         socket_addr: SocketAddr,
     ) -> io::Result<Op<SendTo<T>>> {
         use io_uring::{opcode, types};
@@ -55,13 +55,13 @@ impl<T: IoBuf> Op<SendTo<T>> {
         )
     }
 
-    pub(crate) async fn send(mut self) -> BufResult<usize, T> {
+    pub(crate) async fn send(mut self) -> BufResult<usize, Slice<T>> {
         use crate::future::poll_fn;
 
         poll_fn(move |cx| self.poll_send(cx)).await
     }
 
-    pub(crate) fn poll_send(&mut self, cx: &mut Context<'_>) -> Poll<BufResult<usize, T>> {
+    pub(crate) fn poll_send(&mut self, cx: &mut Context<'_>) -> Poll<BufResult<usize, Slice<T>>> {
         use std::future::Future;
         use std::pin::Pin;
 

diff --git a/src/driver/socket.rs b/src/driver/socket.rs
@@ -1,5 +1,5 @@
 use crate::{
-    buf::{IntoSlice, IoBuf, IoBufMut, Slice},
+    buf::{IntoSlice, IoBufMut, Slice},
     driver::{Op, SharedFd},
 };
 use std::{
@@ -47,12 +47,12 @@ impl Socket {
         op.write().await
     }
 
-    pub(crate) async fn send_to<T: IoBuf>(
+    pub(crate) async fn send_to<T: IntoSlice>(
         &self,
         buf: T,
         socket_addr: SocketAddr,
-    ) -> crate::BufResult<usize, T> {
-        let op = Op::send_to(&self.fd, buf, socket_addr).unwrap();
+    ) -> crate::BufResult<usize, Slice<T::Buf>> {
+        let op = Op::send_to(&self.fd, buf.into_full_slice(), socket_addr).unwrap();
         op.send().await
     }
 
@@ -65,11 +65,15 @@ impl Socket {
         op.read().await
     }
 
-    pub(crate) async fn recv_from<T: IoBufMut>(
+    pub(crate) async fn recv_from<T>(
         &self,
         buf: T,
-    ) -> crate::BufResult<(usize, SocketAddr), T> {
-        let op = Op::recv_from(&self.fd, buf).unwrap();
+    ) -> crate::BufResult<(usize, SocketAddr), Slice<T::Buf>>
+    where
+        T: IntoSlice,
+        T::Buf: IoBufMut,
+    {
+        let op = Op::recv_from(&self.fd, buf.into_full_slice()).unwrap();
         op.recv().await
     }
 

diff --git a/src/net/udp.rs b/src/net/udp.rs
@@ -1,5 +1,5 @@
 use crate::{
-    buf::{IntoSlice, IoBuf, IoBufMut, Slice},
+    buf::{IntoSlice, IoBufMut, Slice},
     driver::Socket,
 };
 use socket2::SockAddr;
@@ -105,17 +105,21 @@ impl UdpSocket {
 
     /// Sends data on the socket to the given address. On success, returns the
     /// number of bytes written.
-    pub async fn send_to<T: IoBuf>(
+    pub async fn send_to<T: IntoSlice>(
         &self,
         buf: T,
         socket_addr: SocketAddr,
-    ) -> crate::BufResult<usize, T> {
+    ) -> crate::BufResult<usize, Slice<T::Buf>> {
         self.inner.send_to(buf, socket_addr).await
     }
 
     /// Receives a single datagram message on the socket. On success, returns
     /// the number of bytes read and the origin.
-    pub async fn recv_from<T: IoBufMut>(&self, buf: T) -> crate::BufResult<(usize, SocketAddr), T> {
+    pub async fn recv_from<T>(&self, buf: T) -> crate::BufResult<(usize, SocketAddr), Slice<T::Buf>>
+    where
+        T: IntoSlice,
+        T::Buf: IoBufMut,
+    {
         self.inner.recv_from(buf).await
     }