From 8d55f98f6fbd6cda5093f0ace9eb31e6c06f3e8d Mon Sep 17 00:00:00 2001 From: John Doneth Date: Wed, 14 Aug 2019 14:18:21 -0400 Subject: [PATCH] udp: update `tokio_udp::UdpFramed` to std::future (#1370) --- tokio-udp/Cargo.toml | 4 ++ tokio-udp/src/frame.rs | 100 +++++++++++++++++++------------- tokio-udp/src/lib.rs | 5 +- tokio-udp/tests/udp.rs | 111 +++++++++++++++++++----------------- tokio/examples/udp-codec.rs | 52 ++++++++--------- tokio/src/net.rs | 4 +- 6 files changed, 153 insertions(+), 123 deletions(-) diff --git a/tokio-udp/Cargo.toml b/tokio-udp/Cargo.toml index b33b3c000..2871d6736 100644 --- a/tokio-udp/Cargo.toml +++ b/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"] } diff --git a/tokio-udp/src/frame.rs b/tokio-udp/src/frame.rs index 70111594b..015685218 100644 --- a/tokio-udp/src/frame.rs +++ b/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 { flushed: bool, } -impl Stream for UdpFramed { - type Item = (C::Item, SocketAddr); - type Error = C::Error; +impl Stream for UdpFramed { + type Item = Result<(C::Item, SocketAddr), C::Error>; - fn poll(&mut self) -> Poll, Self::Error> { - self.rd.reserve(INITIAL_RD_CAPACITY); + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + 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 Sink for UdpFramed { - type SinkItem = (C::Item, SocketAddr); - type SinkError = C::Error; - - fn start_send(&mut self, item: Self::SinkItem) -> StartSend { - trace!("sending frame"); +impl Sink<(C::Item, SocketAddr)> for UdpFramed { + type Error = C::Error; + fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 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> { 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> { + ready!(self.poll_flush(cx))?; + Poll::Ready(Ok(())) } } @@ -118,8 +140,8 @@ impl UdpFramed { /// See struct level documentation for more details. pub fn new(socket: UdpSocket, codec: C) -> UdpFramed { 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), diff --git a/tokio-udp/src/lib.rs b/tokio-udp/src/lib.rs index aa200bfb8..04ac9d4db 100644 --- a/tokio-udp/src/lib.rs +++ b/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; diff --git a/tokio-udp/tests/udp.rs b/tokio-udp/tests/udp.rs index 12b005166..2e48fbb40 100644 --- a/tokio-udp/tests/udp.rs +++ b/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; -// type Error = io::Error; +impl Decoder for ByteCodec { + type Item = Vec; + type Error = io::Error; -// fn decode(&mut self, buf: &mut BytesMut) -> Result>, io::Error> { -// let len = buf.len(); -// Ok(Some(buf.split_to(len).to_vec())) -// } -// } + fn decode(&mut self, buf: &mut BytesMut) -> Result>, io::Error> { + let len = buf.len(); + Ok(Some(buf.split_to(len).to_vec())) + } +} -// impl Encoder for ByteCodec { -// type Item = Vec; -// type Error = io::Error; +impl Encoder for ByteCodec { + type Item = Vec; + type Error = io::Error; -// fn encode(&mut self, data: Vec, buf: &mut BytesMut) -> Result<(), io::Error> { -// buf.reserve(data.len()); -// buf.put(data); -// Ok(()) -// } -// } + fn encode(&mut self, data: Vec, 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(()) +} diff --git a/tokio/examples/udp-codec.rs b/tokio/examples/udp-codec.rs index 2d206e344..a5bd3f642 100644 --- a/tokio/examples/udp-codec.rs +++ b/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> { let _ = env_logger::init(); @@ -27,10 +31,14 @@ async fn main() -> Result<(), Box> { let addr = addr.parse::()?; // 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> { 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, 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) -> 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(()) diff --git a/tokio/src/net.rs b/tokio/src/net.rs index 74ee97fa8..84675fdf4 100644 --- a/tokio/src/net.rs +++ b/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 {