Update tokio-udp to use std-future (#1199)

This commit is contained in:
Yin Guanhao
2019-06-26 14:41:36 -04:00
committed by Lucio Franco
parent 0784dc2767
commit 6316aa1d0b
16 changed files with 347 additions and 605 deletions
+21 -16
View File
@@ -10,23 +10,28 @@
//!
//! The main struct for UDP is the [`UdpSocket`], which represents a UDP socket.
//! Reading and writing to it can be done using futures, which return the
//! [`RecvDgram`] and [`SendDgram`] structs respectively.
//!
//! For convenience it's also possible to convert raw datagrams into higher-level
//! frames.
//!
//! [`UdpSocket`]: struct.UdpSocket.html
//! [`RecvDgram`]: struct.RecvDgram.html
//! [`SendDgram`]: struct.SendDgram.html
//! [`UdpFramed`]: struct.UdpFramed.html
//! [`framed`]: struct.UdpSocket.html#method.framed
//! [`Recv`], [`Send`], [`RecvFrom`] and [`SendTo`] structs respectively.
mod frame;
mod recv_dgram;
mod send_dgram;
macro_rules! ready {
($e:expr) => {
match $e {
::std::task::Poll::Ready(t) => t,
::std::task::Poll::Pending => return ::std::task::Poll::Pending,
}
};
}
// mod frame;
mod recv;
mod recv_from;
mod send;
mod send_to;
mod socket;
pub use self::frame::UdpFramed;
pub use self::recv_dgram::RecvDgram;
pub use self::send_dgram::SendDgram;
// pub use self::frame::UdpFramed;
pub use self::recv::Recv;
pub use self::recv_from::RecvFrom;
pub use self::send::Send;
pub use self::send_to::SendTo;
pub use self::socket::UdpSocket;
+30
View File
@@ -0,0 +1,30 @@
use super::UdpSocket;
use std::future::Future;
use std::io;
use std::pin::Pin;
use std::task::{Context, Poll};
/// A future that receives a datagram from the connected address.
///
/// This `struct` is created by [`recv`](super::UdpSocket::recv).
#[must_use = "futures do nothing unless polled"]
#[derive(Debug)]
pub struct Recv<'a, 'b> {
socket: &'a mut UdpSocket,
buf: &'b mut [u8],
}
impl<'a, 'b> Recv<'a, 'b> {
pub(super) fn new(socket: &'a mut UdpSocket, buf: &'b mut [u8]) -> Self {
Self { socket, buf }
}
}
impl<'a, 'b> Future for Recv<'a, 'b> {
type Output = io::Result<usize>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let Recv { socket, buf } = self.get_mut();
Pin::new(&mut **socket).poll_recv(cx, buf)
}
}
-103
View File
@@ -1,103 +0,0 @@
use super::socket::UdpSocket;
use futures::{try_ready, Async, Future, Poll};
use std::io;
use std::net::SocketAddr;
/// A future used to receive a datagram from a UDP socket.
///
/// This is created by the `UdpSocket::recv_dgram` method.
#[must_use = "futures do nothing unless polled"]
#[derive(Debug)]
pub struct RecvDgram<T> {
/// None means future was completed
state: Option<RecvDgramInner<T>>,
}
/// A struct is used to represent the full info of RecvDgram.
#[derive(Debug)]
struct RecvDgramInner<T> {
/// Rx socket
socket: UdpSocket,
/// The received data will be put in the buffer
buffer: T,
}
/// Components of a `RecvDgram` future, returned from `into_parts`.
#[derive(Debug)]
pub struct Parts<T> {
/// The socket
pub socket: UdpSocket,
/// The buffer
pub buffer: T,
_priv: (),
}
impl<T> RecvDgram<T> {
/// Create a new future to receive UDP Datagram
pub(crate) fn new(socket: UdpSocket, buffer: T) -> RecvDgram<T> {
let inner = RecvDgramInner {
socket: socket,
buffer: buffer,
};
RecvDgram { state: Some(inner) }
}
/// Consume the `RecvDgram`, returning the socket and buffer.
///
/// # Panics
///
/// If called after the future has completed.
///
/// # Examples
///
/// ```
/// use tokio_udp::UdpSocket;
///
/// let socket = UdpSocket::bind(&([127, 0, 0, 1], 0).into()).unwrap();
/// let mut buffer = vec![0; 4096];
///
/// let future = socket.recv_dgram(buffer);
///
/// // ... polling `future` ... giving up (e.g. after timeout)
///
/// let parts = future.into_parts();
///
/// let socket = parts.socket; // extract the socket
/// let buffer = parts.buffer; // extract the buffer
/// ```
pub fn into_parts(mut self) -> Parts<T> {
let state = self
.state
.take()
.expect("into_parts called after completion");
Parts {
socket: state.socket,
buffer: state.buffer,
_priv: (),
}
}
}
impl<T> Future for RecvDgram<T>
where
T: AsMut<[u8]>,
{
type Item = (UdpSocket, T, usize, SocketAddr);
type Error = io::Error;
fn poll(&mut self) -> Poll<Self::Item, io::Error> {
let (n, addr) = {
let ref mut inner = self
.state
.as_mut()
.expect("RecvDgram polled after completion");
try_ready!(inner.socket.poll_recv_from(inner.buffer.as_mut()))
};
let inner = self.state.take().unwrap();
Ok(Async::Ready((inner.socket, inner.buffer, n, addr)))
}
}
+31
View File
@@ -0,0 +1,31 @@
use super::UdpSocket;
use std::future::Future;
use std::io;
use std::net::SocketAddr;
use std::pin::Pin;
use std::task::{Context, Poll};
/// A future that receives a datagram.
///
/// This `struct` is created by [`recv_from`](super::UdpSocket::recv_from).
#[must_use = "futures do nothing unless polled"]
#[derive(Debug)]
pub struct RecvFrom<'a, 'b> {
socket: &'a mut UdpSocket,
buf: &'b mut [u8],
}
impl<'a, 'b> RecvFrom<'a, 'b> {
pub(super) fn new(socket: &'a mut UdpSocket, buf: &'b mut [u8]) -> Self {
Self { socket, buf }
}
}
impl<'a, 'b> Future for RecvFrom<'a, 'b> {
type Output = io::Result<(usize, SocketAddr)>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let RecvFrom { socket, buf } = self.get_mut();
Pin::new(&mut **socket).poll_recv_from(cx, buf)
}
}
+30
View File
@@ -0,0 +1,30 @@
use super::UdpSocket;
use std::future::Future;
use std::io;
use std::pin::Pin;
use std::task::{Context, Poll};
/// A future that sends a datagram to the connected address.
///
/// This `struct` is created by [`send`](super::UdpSocket::send).
#[must_use = "futures do nothing unless polled"]
#[derive(Debug)]
pub struct Send<'a, 'b> {
socket: &'a mut UdpSocket,
buf: &'b [u8],
}
impl<'a, 'b> Send<'a, 'b> {
pub(super) fn new(socket: &'a mut UdpSocket, buf: &'b [u8]) -> Self {
Self { socket, buf }
}
}
impl<'a, 'b> Future for Send<'a, 'b> {
type Output = io::Result<usize>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let Send { socket, buf } = self.get_mut();
Pin::new(&mut **socket).poll_send(cx, buf)
}
}
-70
View File
@@ -1,70 +0,0 @@
use super::socket::UdpSocket;
use futures::{try_ready, Async, Future, Poll};
use std::io;
use std::net::SocketAddr;
/// A future used to write the entire contents of some data to a UDP socket.
///
/// This is created by the `UdpSocket::send_dgram` method.
#[must_use = "futures do nothing unless polled"]
#[derive(Debug)]
pub struct SendDgram<T> {
/// None means future was completed
state: Option<SendDgramInner<T>>,
}
/// A struct is used to represent the full info of SendDgram.
#[derive(Debug)]
struct SendDgramInner<T> {
/// Tx socket
socket: UdpSocket,
/// The whole buffer will be sent
buffer: T,
/// Destination addr
addr: SocketAddr,
}
impl<T> SendDgram<T> {
/// Create a new future to send UDP Datagram
pub(crate) fn new(socket: UdpSocket, buffer: T, addr: SocketAddr) -> SendDgram<T> {
let inner = SendDgramInner {
socket: socket,
buffer: buffer,
addr: addr,
};
SendDgram { state: Some(inner) }
}
}
fn incomplete_write(reason: &str) -> io::Error {
io::Error::new(io::ErrorKind::Other, reason)
}
impl<T> Future for SendDgram<T>
where
T: AsRef<[u8]>,
{
type Item = (UdpSocket, T);
type Error = io::Error;
fn poll(&mut self) -> Poll<(UdpSocket, T), io::Error> {
{
let ref mut inner = self
.state
.as_mut()
.expect("SendDgram polled after completion");
let n = try_ready!(inner
.socket
.poll_send_to(inner.buffer.as_ref(), &inner.addr));
if n != inner.buffer.as_ref().len() {
return Err(incomplete_write(
"failed to send entire message \
in datagram",
));
}
}
let inner = self.state.take().unwrap();
Ok(Async::Ready((inner.socket, inner.buffer)))
}
}
+40
View File
@@ -0,0 +1,40 @@
use super::UdpSocket;
use std::future::Future;
use std::io;
use std::net::SocketAddr;
use std::pin::Pin;
use std::task::{Context, Poll};
/// A future that sends a datagram to a given address.
///
/// This `struct` is created by [`send_to`](super::UdpSocket::send_to).
#[must_use = "futures do nothing unless polled"]
#[derive(Debug)]
pub struct SendTo<'a, 'b> {
socket: &'a mut UdpSocket,
buf: &'b [u8],
target: &'b SocketAddr,
}
impl<'a, 'b> SendTo<'a, 'b> {
pub(super) fn new(socket: &'a mut UdpSocket, buf: &'b [u8], target: &'b SocketAddr) -> Self {
Self {
socket,
buf,
target,
}
}
}
impl<'a, 'b> Future for SendTo<'a, 'b> {
type Output = io::Result<usize>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let SendTo {
socket,
buf,
target,
} = self.get_mut();
Pin::new(&mut **socket).poll_send_to(cx, buf, target)
}
}
+95 -130
View File
@@ -1,10 +1,11 @@
use super::{RecvDgram, SendDgram};
use futures::{try_ready, Async, Poll};
use super::{Recv, RecvFrom, Send, SendTo};
use mio;
use std::convert::TryFrom;
use std::fmt;
use std::io;
use std::net::{self, Ipv4Addr, Ipv6Addr, SocketAddr};
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio_reactor::{Handle, PollEvented};
/// An I/O object representing a UDP socket.
@@ -53,13 +54,15 @@ impl UdpSocket {
self.io.get_ref().connect(*addr)
}
#[deprecated(since = "0.1.2", note = "use poll_send instead")]
#[doc(hidden)]
pub fn send(&mut self, buf: &[u8]) -> io::Result<usize> {
match self.poll_send(buf)? {
Async::Ready(n) => Ok(n),
Async::NotReady => Err(io::ErrorKind::WouldBlock.into()),
}
/// Returns a future that sends data on the socket to the remote address to which it is connected.
/// On success, the future will resolve to the number of bytes written.
///
/// The [`connect`] method will connect this socket to a remote address. The future
/// will resolve to an error if the socket is not connected.
///
/// [`connect`]: #method.connect
pub fn send<'a, 'b>(&'a mut self, buf: &'b [u8]) -> Send<'a, 'b> {
Send::new(self, buf)
}
/// Sends data on the socket to the remote address to which it is connected.
@@ -71,35 +74,41 @@ impl UdpSocket {
///
/// # Return
///
/// On success, returns `Ok(Async::Ready(num_bytes_written))`.
/// On success, returns `Poll::Ready(Ok(num_bytes_written))`.
///
/// If the socket is not ready for writing, the method returns
/// `Ok(Async::NotReady)` and arranges for the current task to receive a
/// `Poll::Pending` and arranges for the current task to receive a
/// notification when the socket becomes writable.
///
/// # Panics
///
/// This function will panic if called from outside of a task context.
pub fn poll_send(&mut self, buf: &[u8]) -> Poll<usize, io::Error> {
try_ready!(self.io.poll_write_ready());
pub fn poll_send(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
ready!(self.io.poll_write_ready(cx))?;
match self.io.get_ref().send(buf) {
Ok(n) => Ok(n.into()),
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_write_ready()?;
Ok(Async::NotReady)
self.io.clear_write_ready(cx)?;
Poll::Pending
}
Err(e) => Err(e),
x => Poll::Ready(x),
}
}
#[deprecated(since = "0.1.2", note = "use poll_recv instead")]
#[doc(hidden)]
pub fn recv(&mut self, buf: &mut [u8]) -> io::Result<usize> {
match self.poll_recv(buf)? {
Async::Ready(n) => Ok(n),
Async::NotReady => Err(io::ErrorKind::WouldBlock.into()),
}
/// Returns a future that receives a single datagram message on the socket from
/// the remote address to which it is connected. On success, the future will resolve
/// to the number of bytes read.
///
/// The function must be called with valid byte array `buf` of sufficient size to
/// hold the message bytes. If a message is too long to fit in the supplied buffer,
/// excess bytes may be discarded.
///
/// The [`connect`] method will connect this socket to a remote address. The future
/// will fail if the socket is not connected.
///
/// [`connect`]: #method.connect
pub fn recv<'a, 'b>(&'a mut self, buf: &'b mut [u8]) -> Recv<'a, 'b> {
Recv::new(self, buf)
}
/// Receives a single datagram message on the socket from the remote address to
@@ -116,35 +125,34 @@ impl UdpSocket {
///
/// # Return
///
/// On success, returns `Ok(Async::Ready(num_bytes_read))`.
/// On success, returns `Poll::Ready(Ok(num_bytes_read))`.
///
/// If no data is available for reading, the method returns
/// `Ok(Async::NotReady)` and arranges for the current task to receive a
/// `Poll::Pending` and arranges for the current task to receive a
/// notification when the socket becomes receivable or is closed.
///
/// # Panics
///
/// This function will panic if called from outside of a task context.
pub fn poll_recv(&mut self, buf: &mut [u8]) -> Poll<usize, io::Error> {
try_ready!(self.io.poll_read_ready(mio::Ready::readable()));
pub fn poll_recv(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
ready!(self.io.poll_read_ready(cx, mio::Ready::readable()))?;
match self.io.get_ref().recv(buf) {
Ok(n) => Ok(n.into()),
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_read_ready(mio::Ready::readable())?;
Ok(Async::NotReady)
self.io.clear_read_ready(cx, mio::Ready::readable())?;
Poll::Pending
}
Err(e) => Err(e),
x => Poll::Ready(x),
}
}
#[deprecated(since = "0.1.2", note = "use poll_send_to instead")]
#[doc(hidden)]
pub fn send_to(&mut self, buf: &[u8], target: &SocketAddr) -> io::Result<usize> {
match self.poll_send_to(buf, target)? {
Async::Ready(n) => Ok(n),
Async::NotReady => Err(io::ErrorKind::WouldBlock.into()),
}
/// Returns a future that sends data on the socket to the given address.
/// On success, the future will resolve to the number of bytes written.
///
/// The future will resolve to an error if the IP version of the socket does
/// not match that of `target`.
pub fn send_to<'a, 'b>(&'a mut self, buf: &'b [u8], target: &'b SocketAddr) -> SendTo<'a, 'b> {
SendTo::new(self, buf, target)
}
/// Sends data on the socket to the given address. On success, returns the
@@ -155,133 +163,90 @@ impl UdpSocket {
///
/// # Return
///
/// On success, returns `Ok(Async::Ready(num_bytes_written))`.
/// On success, returns `Poll::Ready(Ok(num_bytes_written))`.
///
/// If the socket is not ready for writing, the method returns
/// `Ok(Async::NotReady)` and arranges for the current task to receive a
/// `Poll::Pending` and arranges for the current task to receive a
/// notification when the socket becomes writable.
///
/// # Panics
///
/// This function will panic if called from outside of a task context.
pub fn poll_send_to(&mut self, buf: &[u8], target: &SocketAddr) -> Poll<usize, io::Error> {
try_ready!(self.io.poll_write_ready());
pub fn poll_send_to(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
target: &SocketAddr,
) -> Poll<io::Result<usize>> {
ready!(self.io.poll_write_ready(cx))?;
match self.io.get_ref().send_to(buf, target) {
Ok(n) => Ok(n.into()),
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_write_ready()?;
Ok(Async::NotReady)
self.io.clear_write_ready(cx)?;
Poll::Pending
}
Err(e) => Err(e),
x => Poll::Ready(x),
}
}
/// Creates a future that will write the entire contents of the buffer
/// `buf` provided as a datagram to this socket.
/// Returns a future that receives a single datagram on the socket. On success,
/// the future resolves to the number of bytes read and the origin.
///
/// The returned future will return after data has been written to the
/// outbound socket. The future will resolve to the stream as well as the
/// buffer (for reuse if needed).
///
/// Any error which happens during writing will cause both the stream and
/// the buffer to get destroyed. Note that failure to write the entire
/// buffer is considered an error for the purposes of sending a datagram.
///
/// The `buf` parameter here only requires the `AsRef<[u8]>` trait, which
/// should be broadly applicable to accepting data which can be converted
/// to a slice.
pub fn send_dgram<T>(self, buf: T, addr: &SocketAddr) -> SendDgram<T>
where
T: AsRef<[u8]>,
{
SendDgram::new(self, buf, *addr)
}
#[deprecated(since = "0.1.2", note = "use poll_recv_from instead")]
#[doc(hidden)]
pub fn recv_from(&mut self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
match self.poll_recv_from(buf)? {
Async::Ready(ret) => Ok(ret),
Async::NotReady => Err(io::ErrorKind::WouldBlock.into()),
}
/// The function must be called with valid byte array `buf` of sufficient size
/// to hold the message bytes. If a message is too long to fit in the supplied
/// buffer, excess bytes may be discarded.
pub fn recv_from<'a, 'b>(&'a mut self, buf: &'b mut [u8]) -> RecvFrom<'a, 'b> {
RecvFrom::new(self, buf)
}
/// Receives data from the socket. On success, returns the number of bytes
/// read and the address from whence the data came.
///
/// # Panics
///
/// This function will panic if called outside the context of a future's
/// task.
pub fn poll_recv_from(&mut self, buf: &mut [u8]) -> Poll<(usize, SocketAddr), io::Error> {
try_ready!(self.io.poll_read_ready(mio::Ready::readable()));
pub fn poll_recv_from(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<Result<(usize, SocketAddr), io::Error>> {
ready!(self.io.poll_read_ready(cx, mio::Ready::readable()))?;
match self.io.get_ref().recv_from(buf) {
Ok(n) => Ok(n.into()),
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
self.io.clear_read_ready(mio::Ready::readable())?;
Ok(Async::NotReady)
self.io.clear_read_ready(cx, mio::Ready::readable())?;
Poll::Pending
}
Err(e) => Err(e),
x => Poll::Ready(x),
}
}
/// Creates a future that receive a datagram to be written to the buffer
/// provided.
///
/// The returned future will return after a datagram has been received on
/// this socket. The future will resolve to the socket, the buffer, the
/// amount of data read, and the address the data was received from.
///
/// An error during reading will cause the socket and buffer to get
/// destroyed.
///
/// The `buf` parameter here only requires the `AsMut<[u8]>` trait, which
/// should be broadly applicable to accepting data which can be converted
/// to a slice.
pub fn recv_dgram<T>(self, buf: T) -> RecvDgram<T>
where
T: AsMut<[u8]>,
{
RecvDgram::new(self, buf)
}
/// Check the UDP socket's read readiness state.
///
/// The mask argument allows specifying what readiness to notify on. This
/// can be any value, including platform specific readiness, **except**
/// `writable`.
///
/// If the socket is not ready for receiving then `Async::NotReady` is
/// If the socket is not ready for receiving then `Poll::Pending` is
/// returned and the current task is notified once a new event is received.
///
/// The socket will remain in a read-ready state until calls to `poll_recv`
/// return `NotReady`.
/// return `Poll::Pending`.
///
/// # Panics
///
/// This function panics if:
///
/// * `ready` includes writable.
/// * called from outside of a task context.
pub fn poll_read_ready(&self, mask: mio::Ready) -> Poll<mio::Ready, io::Error> {
self.io.poll_read_ready(mask)
pub fn poll_read_ready(
&self,
cx: &mut Context<'_>,
mask: mio::Ready,
) -> Poll<Result<mio::Ready, io::Error>> {
self.io.poll_read_ready(cx, mask)
}
/// Check the UDP socket's write readiness state.
///
/// If the socket is not ready for sending then `Async::NotReady` is
/// If the socket is not ready for sending then `Poll::Pending` is
/// returned and the current task is notified once a new event is received.
///
/// The I/O resource will remain in a write-ready state until calls to
/// `poll_send` return `NotReady`.
///
/// # Panics
///
/// This function panics if called from outside of a task context.
pub fn poll_write_ready(&self) -> Poll<mio::Ready, io::Error> {
self.io.poll_write_ready()
/// `poll_send` return `Poll::Pending`.
pub fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll<Result<mio::Ready, io::Error>> {
self.io.poll_write_ready(cx)
}
/// Gets the value of the `SO_BROADCAST` option for this socket.