From 00a16fc0978f5ae5cbdde3b0c92c92bb71022be0 Mon Sep 17 00:00:00 2001
From: Evan Cameron <cameron.evan@gmail.com>
Date: Mon, 19 Oct 2020 10:09:35 -0400
Subject: [PATCH] Change poll_recv to ReadBuf

---
 tokio/src/net/udp/socket.rs | 26 +++++++++++++++++++-------
 1 file changed, 19 insertions(+), 7 deletions(-)

diff --git a/tokio/src/net/udp/socket.rs b/tokio/src/net/udp/socket.rs
index 8c1bcf42d23..b2d29ba8041 100644
--- 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));
+                }
             }
         }
     }