Change poll_recv to ReadBuf

tokio-rs · Oct 19, 2020 · 00a16fc · 00a16fc
1 parent b260e0c
commit 00a16fc
Showing 1 changed file with 19 additions and 7 deletions.
diff --git a/tokio/src/net/udp/socket.rs b/tokio/src/net/udp/socket.rs
@@ -1,4 +1,4 @@
-use crate::io::PollEvented;
+use crate::io::{PollEvented, ReadBuf};
 use crate::net::{to_socket_addrs, ToSocketAddrs};
 
 use std::convert::TryFrom;
@@ -355,15 +355,23 @@ impl UdpSocket {
     /// This function may encounter any standard I/O error except `WouldBlock`.
     ///
     /// [`connect`]: method@Self::connect
-    pub fn poll_recv(&self, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<io::Result<usize>> {
+    pub fn poll_recv(
+        &self,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<usize>> {
         loop {
             let ev = ready!(self.io.poll_read_ready(cx))?;
 
-            match self.io.get_ref().recv(buf) {
+            match self.io.get_ref().recv(buf.initialize_unfilled()) {
                 Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                     self.io.clear_readiness(ev);
                 }
-                x => return Poll::Ready(x),
+                Err(e) => return Poll::Ready(Err(e)),
+                Ok(n) => {
+                    buf.advance(n);
+                    return Poll::Ready(Ok(n));
+                }
             }
         }
     }
@@ -514,16 +522,20 @@ impl UdpSocket {
     pub fn poll_recv_from(
         &self,
         cx: &mut Context<'_>,
-        buf: &mut [u8],
+        buf: &mut ReadBuf<'_>,
     ) -> Poll<io::Result<(usize, SocketAddr)>> {
         loop {
             let ev = ready!(self.io.poll_read_ready(cx))?;
 
-            match self.io.get_ref().recv_from(buf) {
+            match self.io.get_ref().recv_from(buf.initialize_unfilled()) {
                 Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                     self.io.clear_readiness(ev);
                 }
-                x => return Poll::Ready(x),
+                Err(e) => return Poll::Ready(Err(e)),
+                Ok(n) => {
+                    buf.advance(n);
+                    return Poll::Ready(Ok(n));
+                }
             }
         }
     }