Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

util: Refactor UdpFramed & add UdpFramedRead/UdpFramedWrite #3086

Closed
wants to merge 5 commits into from
Closed

util: Refactor UdpFramed & add UdpFramedRead/UdpFramedWrite #3086

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
1155051
util: add read/write frames udp
leshow Nov 12, 2020
5b32418
util: add some docs udp
leshow Nov 12, 2020
9e1aec4
util: add some more tests udp
leshow Nov 12, 2020
04ddbff
util: remove write_buffer from framed
leshow Nov 12, 2020
dacfa9f
util: udp remove unnecessary Pin
leshow Dec 30, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 0 additions & 10 deletions tokio-util/src/codec/framed.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -208,16 +208,6 @@ impl<T, U> Framed<T, U> {
&mut self.inner.state.read.buffer
}

/// Returns a reference to the write buffer.
pub fn write_buffer(&self) -> &BytesMut {
&self.inner.state.write.buffer
}

/// Returns a mutable reference to the write buffer.
pub fn write_buffer_mut(&mut self) -> &mut BytesMut {
&mut self.inner.state.write.buffer
}

/// Consumes the `Framed`, returning its underlying I/O stream.
///
/// Note that care should be taken to not tamper with the underlying stream
Expand Down
3 changes: 1 addition & 2 deletions tokio-util/src/codec/framed_read.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,9 +48,8 @@ where
inner,
codec: decoder,
state: ReadFrame {
eof: false,
is_readable: false,
buffer: BytesMut::with_capacity(capacity),
..ReadFrame::default()
},
},
}
Expand Down
265 changes: 114 additions & 151 deletions tokio-util/src/udp/frame.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,158 +1,80 @@
use crate::codec::{Decoder, Encoder};
use crate::codec::{Decoder, Encoder, RWFrames, ReadFrame, WriteFrame};
use crate::udp::framed_impl::UdpFramedImpl;

use tokio::{io::ReadBuf, net::UdpSocket};
use pin_project_lite::pin_project;
use tokio::net::UdpSocket;
use tokio_stream::Stream;

use bytes::{BufMut, BytesMut};
use futures_core::ready;
use bytes::BytesMut;
use futures_sink::Sink;
use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
use std::pin::Pin;
use std::task::{Context, Poll};
use std::{io, mem::MaybeUninit};

/// A unified [`Stream`] and [`Sink`] interface to an underlying `UdpSocket`, using
/// the `Encoder` and `Decoder` traits to encode and decode frames.
///
/// Raw UDP sockets work with datagrams, but higher-level code usually wants to
/// batch these into meaningful chunks, called "frames". This method layers
/// framing on top of this socket by using the `Encoder` and `Decoder` traits to
/// handle encoding and decoding of messages frames. Note that the incoming and
/// outgoing frame types may be distinct.
///
/// This function returns a *single* object that is both [`Stream`] and [`Sink`];
/// grouping this into a single object is often useful for layering things which
/// require both read and write access to the underlying object.
///
/// If you want to work more directly with the streams and sink, consider
/// calling [`split`] on the `UdpFramed` returned by this method, which will break
/// them into separate objects, allowing them to interact more easily.
///
/// [`Stream`]: tokio_stream::Stream
/// [`Sink`]: futures_sink::Sink
/// [`split`]: https://docs.rs/futures/0.3/futures/stream/trait.StreamExt.html#method.split
#[must_use = "sinks do nothing unless polled"]
#[cfg_attr(docsrs, doc(all(feature = "codec", feature = "udp")))]
#[derive(Debug)]
pub struct UdpFramed<C> {
socket: UdpSocket,
codec: C,
rd: BytesMut,
wr: BytesMut,
out_addr: SocketAddr,
flushed: bool,
is_readable: bool,
current_addr: Option<SocketAddr>,
use std::{
fmt, io,
net::{Ipv4Addr, SocketAddr, SocketAddrV4},
pin::Pin,
task::{Context, Poll},
};

pin_project! {
/// A unified [`Stream`] and [`Sink`] interface to an underlying `UdpSocket`, using
/// the `Encoder` and `Decoder` traits to encode and decode frames.
///
/// Raw UDP sockets work with datagrams, but higher-level code usually wants to
/// batch these into meaningful chunks, called "frames". This method layers
/// framing on top of this socket by using the `Encoder` and `Decoder` traits to
/// handle encoding and decoding of messages frames. Note that the incoming and
/// outgoing frame types may be distinct.
///
/// This function returns a *single* object that is both [`Stream`] and [`Sink`];
/// grouping this into a single object is often useful for layering things which
/// require both read and write access to the underlying object.
///
/// If you want to work more directly with the streams and sink, consider
/// calling [`split`] on the `UdpFramed` returned by this method, which will break
/// them into separate objects, allowing them to interact more easily.
///
/// [`Stream`]: tokio::stream::Stream
/// [`Sink`]: futures_sink::Sink
/// [`split`]: https://docs.rs/futures/0.3/futures/stream/trait.StreamExt.html#method.split.
#[cfg_attr(docsrs, doc(all(feature = "codec", feature = "udp")))]
pub struct UdpFramed<C> {
#[pin]
inner: UdpFramedImpl<C, RWFrames>,
}
}

const INITIAL_RD_CAPACITY: usize = 64 * 1024;
const INITIAL_WR_CAPACITY: usize = 8 * 1024;

impl<C: Decoder + Unpin> Stream for UdpFramed<C> {
impl<C> Stream for UdpFramed<C>
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removes Unpin bound-- is this okay?

where
C: Decoder,
{
type Item = Result<(C::Item, SocketAddr), C::Error>;

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);

loop {
// Are there still bytes left in the read buffer to decode?
if pin.is_readable {
if let Some(frame) = pin.codec.decode_eof(&mut pin.rd)? {
let current_addr = pin
.current_addr
.expect("will always be set before this line is called");

return Poll::Ready(Some(Ok((frame, current_addr))));
}

// if this line has been reached then decode has returned `None`.
pin.is_readable = false;
pin.rd.clear();
}

// We're out of data. Try and fetch more data to decode
let addr = unsafe {
// Convert `&mut [MaybeUnit<u8>]` to `&mut [u8]` because we will be
// writing to it via `poll_recv_from` and therefore initializing the memory.
let buf = &mut *(pin.rd.chunk_mut() as *mut _ as *mut [MaybeUninit<u8>]);
let mut read = ReadBuf::uninit(buf);
let ptr = read.filled().as_ptr();
let res = ready!(Pin::new(&mut pin.socket).poll_recv_from(cx, &mut read));

assert_eq!(ptr, read.filled().as_ptr());
let addr = res?;
pin.rd.advance_mut(read.filled().len());
addr
};

pin.current_addr = Some(addr);
pin.is_readable = true;
}
self.project().inner.poll_next(cx)
}
}

impl<I, C: Encoder<I> + Unpin> Sink<(I, SocketAddr)> for UdpFramed<C> {
// This impl just defers to the underlying FramedImpl
impl<I, C> Sink<(I, SocketAddr)> for UdpFramed<C>
where
C: Encoder<I>,
C::Error: From<io::Error>,
{
type Error = C::Error;

fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
if !self.flushed {
match self.poll_flush(cx)? {
Poll::Ready(()) => {}
Poll::Pending => return Poll::Pending,
}
}

Poll::Ready(Ok(()))
self.project().inner.poll_ready(cx)
}

fn start_send(self: Pin<&mut Self>, item: (I, SocketAddr)) -> Result<(), Self::Error> {
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;

Ok(())
self.project().inner.start_send(item)
}

fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
if self.flushed {
return Poll::Ready(Ok(()));
}

let Self {
ref mut socket,
ref mut out_addr,
ref mut wr,
..
} = *self;

let n = ready!(socket.poll_send_to(cx, &wr, *out_addr))?;

let wrote_all = n == self.wr.len();
self.wr.clear();
self.flushed = true;

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 poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.project().inner.poll_flush(cx)
}

fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
ready!(self.poll_flush(cx))?;
Poll::Ready(Ok(()))
self.project().inner.poll_close(cx)
}
}

Expand All @@ -162,14 +84,26 @@ impl<C> UdpFramed<C> {
/// See struct level documentation for more details.
pub fn new(socket: UdpSocket, codec: C) -> UdpFramed<C> {
Self {
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),
flushed: true,
is_readable: false,
current_addr: None,
inner: UdpFramedImpl {
inner: socket,
codec,
state: RWFrames {
read: ReadFrame {
buffer: BytesMut::with_capacity(
crate::udp::framed_impl::INITIAL_RD_CAPACITY,
),
..ReadFrame::default()
},
write: WriteFrame {
buffer: BytesMut::with_capacity(
crate::udp::framed_impl::INITIAL_WR_CAPACITY,
),
},
},
out_addr: SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0)),
flushed: true,
current_addr: None,
},
}
}

Expand All @@ -181,7 +115,7 @@ impl<C> UdpFramed<C> {
/// coming in as it may corrupt the stream of frames otherwise being worked
/// with.
pub fn get_ref(&self) -> &UdpSocket {
&self.socket
&self.inner.inner
}

/// Returns a mutable reference to the underlying I/O stream wrapped by
Expand All @@ -193,12 +127,7 @@ impl<C> UdpFramed<C> {
/// coming in as it may corrupt the stream of frames otherwise being worked
/// with.
pub fn get_mut(&mut self) -> &mut UdpSocket {
&mut self.socket
}

/// Consumes the `Framed`, returning its underlying I/O stream.
pub fn into_inner(self) -> UdpSocket {
self.socket
&mut self.inner.inner
}

/// Returns a reference to the underlying codec wrapped by
Expand All @@ -207,7 +136,7 @@ impl<C> UdpFramed<C> {
/// Note that care should be taken to not tamper with the underlying codec
/// as it may corrupt the stream of frames otherwise being worked with.
pub fn codec(&self) -> &C {
&self.codec
&self.inner.codec
}

/// Returns a mutable reference to the underlying codec wrapped by
Expand All @@ -216,16 +145,50 @@ impl<C> UdpFramed<C> {
/// Note that care should be taken to not tamper with the underlying codec
/// as it may corrupt the stream of frames otherwise being worked with.
pub fn codec_mut(&mut self) -> &mut C {
&mut self.codec
&mut self.inner.codec
}

/// Returns a reference to the read buffer.
pub fn read_buffer(&self) -> &BytesMut {
&self.rd
&self.inner.state.read.buffer
}

/// Returns a mutable reference to the read buffer.
pub fn read_buffer_mut(&mut self) -> &mut BytesMut {
&mut self.rd
&mut self.inner.state.read.buffer
}

/// Returns a reference to the write buffer.
pub fn write_buffer(&self) -> &BytesMut {
&self.inner.state.write.buffer
}

/// Returns a mutable reference to the write buffer.
pub fn write_buffer_mut(&mut self) -> &mut BytesMut {
&mut self.inner.state.write.buffer
}

/// Consumes the `Framed`, returning its underlying I/O stream.
pub fn into_inner(self) -> UdpSocket {
self.inner.inner
}
}

impl<C> fmt::Debug for UdpFramed<C>
where
C: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("UdpFramed")
.field("io", self.get_ref())
.field("codec", self.codec())
.field("current_addr", &self.inner.current_addr)
.field("flushed", &self.inner.flushed)
.field("out_addr", &self.inner.out_addr)
.field("is_readable", &self.inner.state.read.is_readable)
.field("eof", &self.inner.state.read.eof)
.field("read_buffer", &self.read_buffer())
.field("write_buffer", &self.write_buffer())
.finish()
}
}
Loading