diff --git a/examples/README.md b/examples/README.md index 688984b9f..3f4734c61 100644 --- a/examples/README.md +++ b/examples/README.md @@ -39,8 +39,8 @@ A high level description of each example is: showcasing running on multiple cores, working with futures and spawning tasks, and finally framing a TCP connection to discrete request/response objects. -* `udp-codec` - an example of using the `UdpCodec` trait along with a small - ping-pong protocol happening locally. +* `udp-codec` - an example of using the `Encoder`/`Decoder` traits for UDP + along with a small ping-pong protocol happening locally. * `compress` - an echo-like server where instead of echoing back everything read it echos back a gzip-compressed version of everything read! All compression occurs on a CPU pool to offload work from the event loop. diff --git a/examples/connect.rs b/examples/connect.rs index a4160449f..275864d12 100644 --- a/examples/connect.rs +++ b/examples/connect.rs @@ -76,17 +76,55 @@ fn main() { }).wait().unwrap(); } +mod codec { + use std::io; + use bytes::{BufMut, BytesMut}; + use tokio_io::codec::{Encoder, Decoder}; + /// A simple `Codec` implementation that just ships bytes around. + /// + /// This type is used for "framing" a TCP/UDP stream of bytes but it's really + /// just a convenient method for us to work with streams/sinks for now. + /// This'll just take any data read and interpret it as a "frame" and + /// conversely just shove data into the output location without looking at + /// it. + pub struct Bytes; + + impl Decoder for Bytes { + type Item = BytesMut; + type Error = io::Error; + + fn decode(&mut self, buf: &mut BytesMut) -> io::Result> { + if buf.len() > 0 { + let len = buf.len(); + Ok(Some(buf.split_to(len))) + } else { + Ok(None) + } + } + } + + impl Encoder for Bytes { + type Item = Vec; + type Error = io::Error; + + fn encode(&mut self, data: Vec, buf: &mut BytesMut) -> io::Result<()> { + buf.put(&data[..]); + Ok(()) + } + } +} + mod tcp { use std::io; use std::net::SocketAddr; - use bytes::{BufMut, BytesMut}; + use bytes::BytesMut; use futures::{Future, Stream}; use futures::future::Executor; use futures_cpupool::CpuPool; use tokio::net::TcpStream; use tokio_io::AsyncRead; - use tokio_io::codec::{Encoder, Decoder}; + use codec::Bytes; pub fn connect(addr: &SocketAddr, pool: &CpuPool, @@ -122,43 +160,6 @@ mod tcp { stream }).flatten_stream()) } - - /// A simple `Codec` implementation that just ships bytes around. - /// - /// This type is used for "framing" a TCP stream of bytes but it's really - /// just a convenient method for us to work with streams/sinks for now. - /// This'll just take any data read and interpret it as a "frame" and - /// conversely just shove data into the output location without looking at - /// it. - struct Bytes; - - impl Decoder for Bytes { - type Item = BytesMut; - type Error = io::Error; - - fn decode(&mut self, buf: &mut BytesMut) -> io::Result> { - if buf.len() > 0 { - let len = buf.len(); - Ok(Some(buf.split_to(len))) - } else { - Ok(None) - } - } - - fn decode_eof(&mut self, buf: &mut BytesMut) -> io::Result> { - self.decode(buf) - } - } - - impl Encoder for Bytes { - type Item = Vec; - type Error = io::Error; - - fn encode(&mut self, data: Vec, buf: &mut BytesMut) -> io::Result<()> { - buf.put(&data[..]); - Ok(()) - } - } } mod udp { @@ -169,7 +170,8 @@ mod udp { use futures::{Future, Stream}; use futures::future::Executor; use futures_cpupool::CpuPool; - use tokio::net::{UdpCodec, UdpSocket}; + use tokio::net::UdpSocket; + use codec::Bytes; pub fn connect(&addr: &SocketAddr, pool: &CpuPool, @@ -186,7 +188,7 @@ mod udp { let udp = UdpSocket::bind(&addr_to_bind) .expect("failed to bind socket"); - // Like above with TCP we use an instance of `UdpCodec` to transform + // Like above with TCP we use an instance of `Bytes` codec to transform // this UDP socket into a framed sink/stream which operates over // discrete values. In this case we're working with *pairs* of socket // addresses and byte buffers. @@ -195,7 +197,7 @@ mod udp { // All bytes from `stdin` will go to the `addr` specified in our // argument list. Like with TCP this is spawned concurrently pool.execute(stdin.map(move |chunk| { - (addr, chunk) + (chunk, addr) }).forward(sink).then(|result| { if let Err(e) = result { panic!("failed to write to socket: {}", e) @@ -205,7 +207,7 @@ mod udp { // With UDP we could receive data from any source, so filter out // anything coming from a different address - Box::new(stream.filter_map(move |(src, chunk)| { + Box::new(stream.filter_map(move |(chunk, src)| { if src == addr { Some(chunk.into()) } else { @@ -213,23 +215,6 @@ mod udp { } })) } - - struct Bytes; - - impl UdpCodec for Bytes { - type In = (SocketAddr, Vec); - type Out = (SocketAddr, Vec); - type Error = io::Error; - - fn decode(&mut self, addr: &SocketAddr, buf: &[u8]) -> io::Result { - Ok((*addr, buf.to_vec())) - } - - fn encode(&mut self, (addr, buf): Self::Out, into: &mut Vec) -> io::Result { - into.extend(buf); - Ok(addr) - } - } } // Our helper method which will read data from stdin and send it along the diff --git a/examples/udp-codec.rs b/examples/udp-codec.rs index c874ebd78..a2429e491 100644 --- a/examples/udp-codec.rs +++ b/examples/udp-codec.rs @@ -1,40 +1,24 @@ -//! This is a basic example of leveraging `UdpCodec` to create a simple UDP +//! This is a basic example of leveraging `BytesCodec` to create a simple UDP //! client and server which speak a custom protocol. //! -//! Here we're using the a custom codec to convert a UDP socket to a stream of +//! Here we're using the codec from tokio-io to convert a UDP socket to a stream of //! client messages. These messages are then processed and returned back as a //! new message with a new destination. Overall, we then use this to construct a //! "ping pong" pair where two sockets are sending messages back and forth. extern crate tokio; +extern crate tokio_io; extern crate env_logger; extern crate futures; extern crate futures_cpupool; -use std::io; use std::net::SocketAddr; use futures::{Future, Stream, Sink}; use futures::future::Executor; use futures_cpupool::CpuPool; -use tokio::net::{UdpSocket, UdpCodec}; - -pub struct LineCodec; - -impl UdpCodec for LineCodec { - type In = (SocketAddr, Vec); - type Out = (SocketAddr, Vec); - type Error = io::Error; - - fn decode(&mut self, addr: &SocketAddr, buf: &[u8]) -> io::Result { - Ok((*addr, buf.to_vec())) - } - - fn encode(&mut self, (addr, buf): Self::Out, into: &mut Vec) -> io::Result { - into.extend(buf); - Ok(addr) - } -} +use tokio::net::UdpSocket; +use tokio_io::codec::BytesCodec; fn main() { drop(env_logger::init()); @@ -50,27 +34,27 @@ fn main() { // We're parsing each socket with the `LineCodec` defined above, and then we // `split` each codec into the sink/stream halves. - let (a_sink, a_stream) = a.framed(LineCodec).split(); - let (b_sink, b_stream) = b.framed(LineCodec).split(); + let (a_sink, a_stream) = a.framed(BytesCodec::new()).split(); + let (b_sink, b_stream) = b.framed(BytesCodec::new()).split(); // Start off by sending a ping from a to b, afterwards we just print out // what they send us and continually send pings // let pings = stream::iter((0..5).map(Ok)); - let a = a_sink.send((b_addr, b"PING".to_vec())).and_then(|a_sink| { + let a = a_sink.send(("PING".into(), b_addr)).and_then(|a_sink| { let mut i = 0; - let a_stream = a_stream.take(4).map(move |(addr, msg)| { + let a_stream = a_stream.take(4).map(move |(msg, addr)| { i += 1; println!("[a] recv: {}", String::from_utf8_lossy(&msg)); - (addr, format!("PING {}", i).into_bytes()) + (format!("PING {}", i).into(), addr) }); a_sink.send_all(a_stream) }); // The second client we have will receive the pings from `a` and then send // back pongs. - let b_stream = b_stream.map(|(addr, msg)| { + let b_stream = b_stream.map(|(msg, addr)| { println!("[b] recv: {}", String::from_utf8_lossy(&msg)); - (addr, b"PONG".to_vec()) + ("PONG".into(), addr) }); let b = b_sink.send_all(b_stream); diff --git a/src/net/mod.rs b/src/net/mod.rs index d8b2f8589..d4a55eb22 100644 --- a/src/net/mod.rs +++ b/src/net/mod.rs @@ -29,18 +29,17 @@ //! //! For convience it's also possible to convert raw datagrams into higher-level //! frames. This done with [`UdpFramed`], created by calling [`framed`] on a -//! [`UdpSocket`], and using the [`UdpCodec`]. +//! [`UdpSocket`]. //! //! [`UdpSocket`]: struct.UdpSocket.html //! [`RecvDgram`]: struct.RecvDgram.html //! [`SendDgram`]: struct.SendDgram.html //! [`UdpFramed`]: struct.UdpFramed.html //! [`framed`]: struct.UdpSocket.html#method.framed -//! [`UdpCodec`]: trait.UdpCodec.html mod tcp; mod udp; pub use self::tcp::{TcpStream, ConnectFuture}; pub use self::tcp::{TcpListener, Incoming}; -pub use self::udp::{UdpSocket, UdpCodec, UdpFramed, SendDgram, RecvDgram}; +pub use self::udp::{UdpSocket, UdpFramed, SendDgram, RecvDgram}; diff --git a/src/net/udp/frame.rs b/src/net/udp/frame.rs index fe8d4ef0c..f84d0842b 100644 --- a/src/net/udp/frame.rs +++ b/src/net/udp/frame.rs @@ -5,63 +5,11 @@ use futures::{Async, Poll, Stream, Sink, StartSend, AsyncSink}; use net::UdpSocket; -/// Encoding of datagrams into frames via buffers. -/// -/// This trait is used when constructing an instance of `UdpFramed` and provides -/// the `In` and `Out` types which are decoded and encoded from the socket, -/// respectively. -/// -/// Because UDP is a connectionless protocol, the `decode` method receives the -/// address where data came from and the `encode` method is also responsible for -/// determining the remote host to which the datagram should be sent -/// -/// The trait itself is implemented on a type that can track state for decoding -/// or encoding, which is particularly useful for streaming parsers. In many -/// cases, though, this type will simply be a unit struct (e.g. `struct -/// MyCodec`). -pub trait UdpCodec { - /// The type of decoded frames. - type In; - - /// The type of frames to be encoded. - type Out; - - /// The type of unrecoverable frame encoding/decoding errors. - /// - /// If an individual message is ill-formed but can be ignored without - /// interfering with the processing of future messages, it may be more - /// useful to report the failure as an `Item`. - /// - /// Note that implementors of this trait can simply indicate `type Error = - /// io::Error` to use I/O errors as this type. - type Error: From; - - /// Attempts to decode a frame from the provided buffer of bytes. - /// - /// This method is called by `UdpFramed` on a single datagram which has been - /// read from a socket. The `buf` argument contains the data that was - /// received from the remote address, and `src` is the address the data came - /// from. Note that typically this method should require the entire contents - /// of `buf` to be valid or otherwise return an error with trailing data. - /// - /// Finally, if the bytes in the buffer are malformed then an error is - /// returned indicating why. This informs `Framed` that the stream is now - /// corrupt and should be terminated. - fn decode(&mut self, src: &SocketAddr, buf: &[u8]) -> Result; - - /// Encodes a frame into the buffer provided. - /// - /// This method will encode `msg` into the byte buffer provided by `buf`. - /// The `buf` provided is an internal buffer of the `Framed` instance and - /// will be written out when possible. - /// - /// The encode method also determines the destination to which the buffer - /// should be directed, which will be returned as a `SocketAddr`. - fn encode(&mut self, msg: Self::Out, buf: &mut Vec) -> Result; -} +use tokio_io::codec::{Decoder, Encoder}; +use bytes::{BytesMut, BufMut}; /// A unified `Stream` and `Sink` interface to an underlying `UdpSocket`, using -/// the `UdpCodec` trait to encode and decode frames. +/// the `Encoder` and `Decoder` traits to encode and decode frames. /// /// You can acquire a `UdpFramed` instance by using the `UdpSocket::framed` /// adapter. @@ -70,30 +18,40 @@ pub trait UdpCodec { pub struct UdpFramed { socket: UdpSocket, codec: C, - rd: Vec, - wr: Vec, + rd: BytesMut, + wr: BytesMut, out_addr: SocketAddr, flushed: bool, } -impl Stream for UdpFramed { - type Item = C::In; +impl Stream for UdpFramed { + type Item = (C::Item, SocketAddr); type Error = C::Error; - fn poll(&mut self) -> Poll, C::Error> { - let (n, addr) = try_nb!(self.socket.recv_from(&mut self.rd)); + fn poll(&mut self) -> Poll, Self::Error> { + self.rd.reserve(INITIAL_RD_CAPACITY); + + let (n, addr) = unsafe { + // Read into the buffer without having to initialize the memory. + let (n, addr) = try_nb!(self.socket.recv_from(self.rd.bytes_mut())); + self.rd.advance_mut(n); + (n, addr) + }; trace!("received {} bytes, decoding", n); - let frame = self.codec.decode(&addr, &self.rd[..n])?; + let frame_res = self.codec.decode(&mut self.rd); + self.rd.clear(); + let frame = frame_res?; + let result = frame.map(|frame| (frame, addr)); // frame -> (frame, addr) trace!("frame decoded from buffer"); - Ok(Async::Ready(Some(frame))) + Ok(Async::Ready(result)) } } -impl Sink for UdpFramed { - type SinkItem = C::Out; +impl Sink for UdpFramed { + type SinkItem = (C::Item, SocketAddr); type SinkError = C::Error; - fn start_send(&mut self, item: C::Out) -> StartSend { + fn start_send(&mut self, item: Self::SinkItem) -> StartSend { trace!("sending frame"); if !self.flushed { @@ -103,7 +61,9 @@ impl Sink for UdpFramed { } } - self.out_addr = self.codec.encode(item, &mut self.wr)?; + 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()); @@ -137,13 +97,16 @@ impl Sink for UdpFramed { } } -pub fn new(socket: UdpSocket, codec: C) -> UdpFramed { +const INITIAL_RD_CAPACITY: usize = 64 * 1024; +const INITIAL_WR_CAPACITY: usize = 8 * 1024; + +pub fn new(socket: UdpSocket, codec: C) -> UdpFramed { UdpFramed { socket: socket, codec: codec, out_addr: SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0)), - rd: vec![0; 64 * 1024], - wr: Vec::with_capacity(8 * 1024), + rd: BytesMut::with_capacity(INITIAL_RD_CAPACITY), + wr: BytesMut::with_capacity(INITIAL_WR_CAPACITY), flushed: true, } } diff --git a/src/net/udp/mod.rs b/src/net/udp/mod.rs index f8715a82b..1d022abfe 100644 --- a/src/net/udp/mod.rs +++ b/src/net/udp/mod.rs @@ -13,7 +13,8 @@ pub struct UdpSocket { } mod frame; -pub use self::frame::{UdpFramed, UdpCodec}; +pub use self::frame::UdpFramed; +use tokio_io::codec::{Decoder, Encoder}; impl UdpSocket { /// This function will create a new UDP socket and attempt to bind it to @@ -44,14 +45,14 @@ impl UdpSocket { } /// Provides a `Stream` and `Sink` interface for reading and writing to this - /// `UdpSocket` object, using the provided `UdpCodec` to read and write the - /// raw data. + /// `UdpSocket` object, using the provided codec that must implement + /// `Encoder` and `Decoder` traits to read and write the raw data. /// /// 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 `UdpCodec` - /// trait to handle encoding and decoding of messages frames. Note that - /// the incoming and outgoing frame types may be distinct. + /// 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 @@ -62,7 +63,7 @@ impl UdpSocket { /// calling `split` on the `UdpFramed` returned by this method, which will /// break them into separate objects, allowing them to interact more /// easily. - pub fn framed(self, codec: C) -> UdpFramed { + pub fn framed(self, codec: C) -> UdpFramed { frame::new(self, codec) } diff --git a/tests/udp.rs b/tests/udp.rs index f0a47d37c..37c986986 100644 --- a/tests/udp.rs +++ b/tests/udp.rs @@ -2,12 +2,17 @@ extern crate futures; extern crate tokio; #[macro_use] extern crate tokio_io; +extern crate bytes; +extern crate env_logger; use std::io; use std::net::SocketAddr; use futures::{Future, Poll, Stream, Sink}; -use tokio::net::{UdpSocket, UdpCodec}; + +use tokio::net::UdpSocket; +use tokio_io::codec::{Encoder, Decoder}; +use bytes::{BytesMut, BufMut}; macro_rules! t { ($e:expr) => (match $e { @@ -187,58 +192,68 @@ fn send_dgrams() { } } -#[derive(Debug, Clone)] -struct Codec { - data: &'static [u8], - from: SocketAddr, - to: SocketAddr, -} +pub struct ByteCodec; -impl UdpCodec for Codec { - type In = (); - type Out = &'static [u8]; +impl Decoder for ByteCodec { + type Item = Vec; type Error = io::Error; - fn decode(&mut self, src: &SocketAddr, buf: &[u8]) -> io::Result { - assert_eq!(src, &self.from); - assert_eq!(buf, self.data); - Ok(()) + fn decode(&mut self, buf: &mut BytesMut) -> Result>, io::Error> { + let len = buf.len(); + Ok(Some(buf.split_to(len).to_vec())) } +} - fn encode(&mut self, msg: Self::Out, buf: &mut Vec) -> io::Result { - assert_eq!(msg, self.data); - buf.extend_from_slice(msg); - Ok(self.to) +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(()) } } #[test] fn send_framed() { + drop(env_logger::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 a = a_soc.framed(Codec { data: &b"4567"[..], from: a_addr, to: b_addr}); - let b = b_soc.framed(Codec { data: &b"4567"[..], from: a_addr, to: b_addr}); + let a = a_soc.framed(ByteCodec); + let b = b_soc.framed(ByteCodec); - let send = a.send(&b"4567"[..]); + let msg = b"4567".to_vec(); + + 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()); - assert_eq!(received.0, Some(())); + + let (data, addr) = received.0.unwrap(); + assert_eq!(msg, data); + assert_eq!(a_addr, addr); a_soc = sendt.into_inner(); b_soc = received.1.into_inner(); } { - let a = a_soc.framed(Codec { data: &b""[..], from: a_addr, to: b_addr}); - let b = b_soc.framed(Codec { data: &b""[..], from: a_addr, to: b_addr}); + let a = a_soc.framed(ByteCodec); + let b = b_soc.framed(ByteCodec); - let send = a.send(&b""[..]); + let msg = b"".to_vec(); + + 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; - assert_eq!(received.0, Some(())); + + let (data, addr) = received.0.unwrap(); + assert_eq!(msg, data); + assert_eq!(a_addr, addr); } }