udp: update tokio_udp::UdpFramed to std::future (#1370)

tokio-udp/Cargo.toml

@@ -20,12 +20,16 @@ UDP bindings for tokio.
 categories = ["asynchronous"]
 
 [dependencies]
+tokio-codec = { version = "=0.2.0-alpha.1", path = "../tokio-codec" }
+# tokio-io = { version = "0.2.0", path = "../tokio-io" }
 tokio-reactor = { version = "=0.2.0-alpha.1", path = "../tokio-reactor" }
 
+bytes = "0.4.12"
 mio = "0.6.14"
 log = "0.4"
 futures-core-preview = "=0.3.0-alpha.18"
 futures-util-preview = "=0.3.0-alpha.18"
+futures-sink-preview = "=0.3.0-alpha.18"
 
 [dev-dependencies]
 tokio = { version = "=0.2.0-alpha.1", path = "../tokio", default-features = false, features = ["rt-full"] }

tokio-udp/src/frame.rs

@@ -1,9 +1,12 @@
 use super::UdpSocket;
 use bytes::{BufMut, BytesMut};
-use futures::{try_ready, Async, AsyncSink, Poll, Sink, StartSend, Stream};
+use core::task::{Context, Poll};
+use futures_core::{ready, Stream};
+use futures_sink::Sink;
 use log::trace;
 use std::io;
 use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
+use std::pin::Pin;
 use tokio_codec::{Decoder, Encoder};
 
 /// A unified `Stream` and `Sink` interface to an underlying `UdpSocket`, using
@@ -33,79 +36,98 @@ pub struct UdpFramed<C> {
     flushed: bool,
 }
 
-impl<C: Decoder> Stream for UdpFramed<C> {
-    type Item = (C::Item, SocketAddr);
-    type Error = C::Error;
+impl<C: Decoder + Unpin> Stream for UdpFramed<C> {
+    type Item = Result<(C::Item, SocketAddr), C::Error>;
 
-    fn poll(&mut self) -> Poll<Option<(Self::Item)>, Self::Error> {
-        self.rd.reserve(INITIAL_RD_CAPACITY);
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let pin = self.get_mut();
+
+        pin.rd.reserve(INITIAL_RD_CAPACITY);
 
         let (n, addr) = unsafe {
             // Read into the buffer without having to initialize the memory.
-            let (n, addr) = try_ready!(self.socket.poll_recv_from(self.rd.bytes_mut()));
-            self.rd.advance_mut(n);
+            let res = ready!(Pin::new(&mut pin.socket).poll_recv_from_priv(cx, pin.rd.bytes_mut()));
+            let (n, addr) = res?;
+            pin.rd.advance_mut(n);
             (n, addr)
         };
         trace!("received {} bytes, decoding", n);
-        let frame_res = self.codec.decode(&mut self.rd);
-        self.rd.clear();
+        let frame_res = pin.codec.decode(&mut pin.rd);
+        pin.rd.clear();
         let frame = frame_res?;
-        let result = frame.map(|frame| (frame, addr)); // frame -> (frame, addr)
+        let result = frame.map(|frame| Ok((frame, addr))); // frame -> (frame, addr)
+
         trace!("frame decoded from buffer");
-        Ok(Async::Ready(result))
+        Poll::Ready(result)
     }
 }
 
-impl<C: Encoder> Sink for UdpFramed<C> {
-    type SinkItem = (C::Item, SocketAddr);
-    type SinkError = C::Error;
-
-    fn start_send(&mut self, item: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
-        trace!("sending frame");
+impl<C: Encoder + Unpin> Sink<(C::Item, SocketAddr)> for UdpFramed<C> {
+    type Error = C::Error;
 
+    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
         if !self.flushed {
-            match self.poll_complete()? {
-                Async::Ready(()) => {}
-                Async::NotReady => return Ok(AsyncSink::NotReady(item)),
+            match self.poll_flush(cx)? {
+                Poll::Ready(()) => {}
+                Poll::Pending => return Poll::Pending,
             }
         }
 
-        let (frame, out_addr) = item;
-        self.codec.encode(frame, &mut self.wr)?;
-        self.out_addr = out_addr;
-        self.flushed = false;
-        trace!("frame encoded; length={}", self.wr.len());
-
-        Ok(AsyncSink::Ready)
+        Poll::Ready(Ok(()))
     }
 
-    fn poll_complete(&mut self) -> Poll<(), C::Error> {
+    fn start_send(self: Pin<&mut Self>, item: (C::Item, SocketAddr)) -> Result<(), Self::Error> {
+        trace!("sending frame");
+
+        let (frame, out_addr) = item;
+
+        let pin = self.get_mut();
+
+        pin.codec.encode(frame, &mut pin.wr)?;
+        pin.out_addr = out_addr;
+        pin.flushed = false;
+        trace!("frame encoded; length={}", pin.wr.len());
+
+        Ok(())
+    }
+
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
         if self.flushed {
-            return Ok(Async::Ready(()));
+            return Poll::Ready(Ok(()));
         }
 
         trace!("flushing frame; length={}", self.wr.len());
-        let n = try_ready!(self.socket.poll_send_to(&self.wr, &self.out_addr));
+
+        let Self {
+            ref mut socket,
+            ref mut out_addr,
+            ref mut wr,
+            ..
+        } = *self;
+
+        let n = ready!(socket.poll_send_to_priv(cx, &wr, &out_addr))?;
         trace!("written {}", n);
 
         let wrote_all = n == self.wr.len();
         self.wr.clear();
         self.flushed = true;
 
-        if wrote_all {
-            Ok(Async::Ready(()))
+        let res = if wrote_all {
+            Ok(())
         } else {
             Err(io::Error::new(
                 io::ErrorKind::Other,
                 "failed to write entire datagram to socket",
             )
             .into())
-        }
+        };
+
+        Poll::Ready(res)
     }
 
-    fn close(&mut self) -> Poll<(), C::Error> {
-        try_ready!(self.poll_complete());
-        Ok(().into())
+    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        ready!(self.poll_flush(cx))?;
+        Poll::Ready(Ok(()))
     }
 }
 
@@ -118,8 +140,8 @@ impl<C> UdpFramed<C> {
     /// See struct level documentation for more details.
     pub fn new(socket: UdpSocket, codec: C) -> UdpFramed<C> {
         UdpFramed {
-            socket: socket,
-            codec: codec,
+            socket,
+            codec,
             out_addr: SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0)),
             rd: BytesMut::with_capacity(INITIAL_RD_CAPACITY),
             wr: BytesMut::with_capacity(INITIAL_WR_CAPACITY),

tokio-udp/src/lib.rs

@@ -17,10 +17,9 @@
 //! Reading and writing to it can be done using futures, which return the
 //! [`Recv`], [`Send`], [`RecvFrom`] and [`SendTo`] structs respectively.
 
-// mod frame;
+mod frame;
 mod socket;
 pub mod split;
 
-// pub use self::frame::UdpFramed;
-
+pub use self::frame::UdpFramed;
 pub use self::socket::UdpSocket;

tokio-udp/tests/udp.rs

@@ -1,7 +1,11 @@
 #![feature(async_await)]
 #![warn(rust_2018_idioms)]
 
-use tokio_udp::UdpSocket;
+use bytes::{BufMut, BytesMut};
+use futures_util::{future::FutureExt, sink::SinkExt, stream::StreamExt, try_future::try_join};
+use std::io;
+use tokio_codec::{Decoder, Encoder};
+use tokio_udp::{UdpFramed, UdpSocket};
 
 #[tokio::test]
 async fn send_recv() -> std::io::Result<()> {
@@ -72,68 +76,73 @@ async fn reunite_error() -> std::io::Result<()> {
     Ok(())
 }
 
-// pub struct ByteCodec;
+pub struct ByteCodec;
 
-// impl Decoder for ByteCodec {
-//     type Item = Vec<u8>;
-//     type Error = io::Error;
+impl Decoder for ByteCodec {
+    type Item = Vec<u8>;
+    type Error = io::Error;
 
-//     fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Vec<u8>>, io::Error> {
-//         let len = buf.len();
-//         Ok(Some(buf.split_to(len).to_vec()))
-//     }
-// }
+    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Vec<u8>>, io::Error> {
+        let len = buf.len();
+        Ok(Some(buf.split_to(len).to_vec()))
+    }
+}
 
-// impl Encoder for ByteCodec {
-//     type Item = Vec<u8>;
-//     type Error = io::Error;
+impl Encoder for ByteCodec {
+    type Item = Vec<u8>;
+    type Error = io::Error;
 
-//     fn encode(&mut self, data: Vec<u8>, buf: &mut BytesMut) -> Result<(), io::Error> {
-//         buf.reserve(data.len());
-//         buf.put(data);
-//         Ok(())
-//     }
-// }
+    fn encode(&mut self, data: Vec<u8>, buf: &mut BytesMut) -> Result<(), io::Error> {
+        buf.reserve(data.len());
+        buf.put(data);
+        Ok(())
+    }
+}
 
-// #[test]
-// fn send_framed() {
-//     drop(env_logger::try_init());
+#[tokio::test]
+async fn send_framed() -> std::io::Result<()> {
+    drop(env_logger::try_init());
 
-//     let mut a_soc = t!(UdpSocket::bind(&t!("127.0.0.1:0".parse())));
-//     let mut b_soc = t!(UdpSocket::bind(&t!("127.0.0.1:0".parse())));
-//     let a_addr = t!(a_soc.local_addr());
-//     let b_addr = t!(b_soc.local_addr());
+    let mut a_soc = UdpSocket::bind(&"127.0.0.1:0".parse().unwrap())?;
+    let mut b_soc = UdpSocket::bind(&"127.0.0.1:0".parse().unwrap())?;
 
-//     {
-//         let a = UdpFramed::new(a_soc, ByteCodec);
-//         let b = UdpFramed::new(b_soc, ByteCodec);
+    let a_addr = a_soc.local_addr()?;
+    let b_addr = b_soc.local_addr()?;
 
-//         let msg = b"4567".to_vec();
+    // test sending & receiving bytes
+    {
+        let mut a = UdpFramed::new(a_soc, ByteCodec);
+        let mut b = UdpFramed::new(b_soc, ByteCodec);
 
-//         let send = a.send((msg.clone(), b_addr));
-//         let recv = b.into_future().map_err(|e| e.0);
-//         let (sendt, received) = t!(send.join(recv).wait());
+        let msg = b"4567".to_vec();
 
-//         let (data, addr) = received.0.unwrap();
-//         assert_eq!(msg, data);
-//         assert_eq!(a_addr, addr);
+        let send = a.send((msg.clone(), b_addr));
+        let recv = b.next().map(|e| e.unwrap());
+        let (_, received) = try_join(send, recv).await.unwrap();
 
-//         a_soc = sendt.into_inner();
-//         b_soc = received.1.into_inner();
-//     }
+        let (data, addr) = received;
+        assert_eq!(msg, data);
+        assert_eq!(a_addr, addr);
 
-//     {
-//         let a = UdpFramed::new(a_soc, ByteCodec);
-//         let b = UdpFramed::new(b_soc, ByteCodec);
+        a_soc = a.into_inner();
+        b_soc = b.into_inner();
+    }
 
-//         let msg = b"".to_vec();
+    // test sending & receiving an empty message
+    {
+        let mut a = UdpFramed::new(a_soc, ByteCodec);
+        let mut b = UdpFramed::new(b_soc, ByteCodec);
 
-//         let send = a.send((msg.clone(), b_addr));
-//         let recv = b.into_future().map_err(|e| e.0);
-//         let received = t!(send.join(recv).wait()).1;
+        let msg = b"".to_vec();
 
-//         let (data, addr) = received.0.unwrap();
-//         assert_eq!(msg, data);
-//         assert_eq!(a_addr, addr);
-//     }
-// }
+        let send = a.send((msg.clone(), b_addr));
+        let recv = b.next().map(|e| e.unwrap());
+        let (_, received) = try_join(send, recv).await.unwrap();
+
+        let (data, addr) = received;
+        assert_eq!(msg, data);
+        assert_eq!(a_addr, addr);
+    }
+
+    Ok(())
+}

tokio/examples/udp-codec.rs

@@ -10,15 +10,19 @@
 #![cfg(feature = "rt-full")]
 #![warn(rust_2018_idioms)]
 
-use tokio::io;
-use tokio::net::UdpSocket;
-use tokio::prelude::*;
-
 use std::env;
 use std::error::Error;
 use std::net::SocketAddr;
 use std::time::Duration;
 
+use bytes::Bytes;
+
+use futures::{FutureExt, SinkExt, StreamExt};
+use tokio::codec::BytesCodec;
+use tokio::future::FutureExt as TokioFutureExt;
+use tokio::io;
+use tokio::net::{UdpFramed, UdpSocket};
+
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn Error>> {
     let _ = env_logger::init();
@@ -27,10 +31,14 @@ async fn main() -> Result<(), Box<dyn Error>> {
     let addr = addr.parse::<SocketAddr>()?;
 
     // Bind both our sockets and then figure out what ports we got.
-    let mut a = UdpSocket::bind(&addr)?;
-    let mut b = UdpSocket::bind(&addr)?;
+    let a = UdpSocket::bind(&addr)?;
+    let b = UdpSocket::bind(&addr)?;
+
     let b_addr = b.local_addr()?;
 
+    let mut a = UdpFramed::new(a, BytesCodec::new());
+    let mut b = UdpFramed::new(b, BytesCodec::new());
+
     // Start off by sending a ping from a to b, afterwards we just print out
     // what they send us and continually send pings
     let a = ping(&mut a, b_addr);
@@ -48,39 +56,27 @@ async fn main() -> Result<(), Box<dyn Error>> {
     Ok(())
 }
 
-async fn ping(socket: &mut UdpSocket, b_addr: SocketAddr) -> Result<(), io::Error> {
-    socket.send_to(b"PING", &b_addr).await?;
+async fn ping(socket: &mut UdpFramed<BytesCodec>, b_addr: SocketAddr) -> Result<(), io::Error> {
+    socket.send((Bytes::from(&b"PING"[..]), b_addr)).await?;
 
     for _ in 0..4usize {
-        let mut buffer = [0u8; 255];
+        let (bytes, addr) = socket.next().map(|e| e.unwrap()).await?;
 
-        let (bytes_read, addr) = socket.recv_from(&mut buffer).await?;
+        println!("[a] recv: {}", String::from_utf8_lossy(&bytes));
 
-        println!(
-            "[a] recv: {}",
-            String::from_utf8_lossy(&buffer[..bytes_read])
-        );
-
-        socket.send_to(b"PING", &addr).await?;
+        socket.send((Bytes::from(&b"PING"[..]), addr)).await?;
     }
 
     Ok(())
 }
 
-async fn pong(socket: &mut UdpSocket) -> Result<(), io::Error> {
-    let mut buffer = [0u8; 255];
+async fn pong(socket: &mut UdpFramed<BytesCodec>) -> Result<(), io::Error> {
+    let timeout = Duration::from_millis(200);
 
-    while let Ok(Ok((bytes_read, addr))) = socket
-        .recv_from(&mut buffer)
-        .timeout(Duration::from_millis(200))
-        .await
-    {
-        println!(
-            "[b] recv: {}",
-            String::from_utf8_lossy(&buffer[..bytes_read])
-        );
+    while let Ok(Some(Ok((bytes, addr)))) = socket.next().timeout(timeout).await {
+        println!("[b] recv: {}", String::from_utf8_lossy(&bytes));
 
-        socket.send_to(b"PONG", &addr).await?;
+        socket.send((Bytes::from(&b"PONG"[..]), addr)).await?;
     }
 
     Ok(())

tokio/src/net.rs

@@ -58,10 +58,10 @@ pub mod udp {
     //! [`Send`]: struct.Send.html
     //! [`RecvFrom`]: struct.RecvFrom.html
     //! [`SendTo`]: struct.SendTo.html
-    pub use tokio_udp::{split, UdpSocket};
+    pub use tokio_udp::{split, UdpFramed, UdpSocket};
 }
 #[cfg(feature = "udp")]
-pub use self::udp::UdpSocket;
+pub use self::udp::{UdpFramed, UdpSocket};
 
 #[cfg(all(unix, feature = "uds"))]
 pub mod unix {