refactor(proto): move more h1-specific pieces into h1 module
This commit is contained in:
Sean McArthur
1254
src/proto/h1/conn.rs
Normal file
1254
src/proto/h1/conn.rs
Normal file
File diff suppressed because it is too large
Load Diff
@@ -5,7 +5,8 @@ use std::io;
|
||||
|
||||
use futures::{Async, Poll};
|
||||
use bytes::Bytes;
|
||||
use proto::io::MemRead;
|
||||
|
||||
use super::io::MemRead;
|
||||
|
||||
use self::Kind::{Length, Chunked, Eof};
|
||||
|
||||
@@ -320,7 +321,7 @@ mod tests {
|
||||
use std::io::Write;
|
||||
use super::Decoder;
|
||||
use super::ChunkedState;
|
||||
use proto::io::MemRead;
|
||||
use super::super::io::MemRead;
|
||||
use futures::{Async, Poll};
|
||||
use bytes::{BytesMut, Bytes};
|
||||
use mock::AsyncIo;
|
||||
|
||||
470
src/proto/h1/dispatch.rs
Normal file
470
src/proto/h1/dispatch.rs
Normal file
@@ -0,0 +1,470 @@
|
||||
use std::io;
|
||||
|
||||
use futures::{Async, AsyncSink, Future, Poll, Stream};
|
||||
use futures::sync::{mpsc, oneshot};
|
||||
use tokio_io::{AsyncRead, AsyncWrite};
|
||||
use tokio_service::Service;
|
||||
|
||||
use proto::{Body, Conn, KeepAlive, Http1Transaction, MessageHead, RequestHead, ResponseHead};
|
||||
use ::StatusCode;
|
||||
|
||||
pub struct Dispatcher<D, Bs, I, B, T, K> {
|
||||
conn: Conn<I, B, T, K>,
|
||||
dispatch: D,
|
||||
body_tx: Option<::proto::body::ChunkSender>,
|
||||
body_rx: Option<Bs>,
|
||||
is_closing: bool,
|
||||
}
|
||||
|
||||
pub trait Dispatch {
|
||||
type PollItem;
|
||||
type PollBody;
|
||||
type RecvItem;
|
||||
fn poll_msg(&mut self) -> Poll<Option<(Self::PollItem, Option<Self::PollBody>)>, ::Error>;
|
||||
fn recv_msg(&mut self, msg: ::Result<(Self::RecvItem, Option<Body>)>) -> ::Result<()>;
|
||||
fn poll_ready(&mut self) -> Poll<(), ()>;
|
||||
fn should_poll(&self) -> bool;
|
||||
}
|
||||
|
||||
pub struct Server<S: Service> {
|
||||
in_flight: Option<S::Future>,
|
||||
service: S,
|
||||
}
|
||||
|
||||
pub struct Client<B> {
|
||||
callback: Option<oneshot::Sender<::Result<::Response>>>,
|
||||
rx: ClientRx<B>,
|
||||
}
|
||||
|
||||
pub enum ClientMsg<B> {
|
||||
Request(RequestHead, Option<B>, oneshot::Sender<::Result<::Response>>),
|
||||
Close,
|
||||
}
|
||||
|
||||
type ClientRx<B> = mpsc::Receiver<ClientMsg<B>>;
|
||||
|
||||
impl<D, Bs, I, B, T, K> Dispatcher<D, Bs, I, B, T, K>
|
||||
where
|
||||
D: Dispatch<PollItem=MessageHead<T::Outgoing>, PollBody=Bs, RecvItem=MessageHead<T::Incoming>>,
|
||||
I: AsyncRead + AsyncWrite,
|
||||
B: AsRef<[u8]>,
|
||||
T: Http1Transaction,
|
||||
K: KeepAlive,
|
||||
Bs: Stream<Item=B, Error=::Error>,
|
||||
{
|
||||
pub fn new(dispatch: D, conn: Conn<I, B, T, K>) -> Self {
|
||||
Dispatcher {
|
||||
conn: conn,
|
||||
dispatch: dispatch,
|
||||
body_tx: None,
|
||||
body_rx: None,
|
||||
is_closing: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn disable_keep_alive(&mut self) {
|
||||
self.conn.disable_keep_alive()
|
||||
}
|
||||
|
||||
fn poll2(&mut self) -> Poll<(), ::Error> {
|
||||
self.poll_read()?;
|
||||
self.poll_write()?;
|
||||
self.poll_flush()?;
|
||||
|
||||
if self.is_done() {
|
||||
try_ready!(self.conn.shutdown());
|
||||
self.conn.take_error()?;
|
||||
trace!("Dispatch::poll done");
|
||||
Ok(Async::Ready(()))
|
||||
} else {
|
||||
Ok(Async::NotReady)
|
||||
}
|
||||
}
|
||||
|
||||
fn poll_read(&mut self) -> Poll<(), ::Error> {
|
||||
loop {
|
||||
if self.is_closing {
|
||||
return Ok(Async::Ready(()));
|
||||
} else if self.conn.can_read_head() {
|
||||
try_ready!(self.poll_read_head());
|
||||
} else if let Some(mut body) = self.body_tx.take() {
|
||||
if self.conn.can_read_body() {
|
||||
match body.poll_ready() {
|
||||
Ok(Async::Ready(())) => (),
|
||||
Ok(Async::NotReady) => {
|
||||
self.body_tx = Some(body);
|
||||
return Ok(Async::NotReady);
|
||||
},
|
||||
Err(_canceled) => {
|
||||
// user doesn't care about the body
|
||||
// so we should stop reading
|
||||
trace!("body receiver dropped before eof, closing");
|
||||
self.conn.close_read();
|
||||
return Ok(Async::Ready(()));
|
||||
}
|
||||
}
|
||||
match self.conn.read_body() {
|
||||
Ok(Async::Ready(Some(chunk))) => {
|
||||
match body.start_send(Ok(chunk)) {
|
||||
Ok(AsyncSink::Ready) => {
|
||||
self.body_tx = Some(body);
|
||||
},
|
||||
Ok(AsyncSink::NotReady(_chunk)) => {
|
||||
unreachable!("mpsc poll_ready was ready, start_send was not");
|
||||
}
|
||||
Err(_canceled) => {
|
||||
if self.conn.can_read_body() {
|
||||
trace!("body receiver dropped before eof, closing");
|
||||
self.conn.close_read();
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
},
|
||||
Ok(Async::Ready(None)) => {
|
||||
// just drop, the body will close automatically
|
||||
},
|
||||
Ok(Async::NotReady) => {
|
||||
self.body_tx = Some(body);
|
||||
return Ok(Async::NotReady);
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = body.start_send(Err(::Error::Io(e)));
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// just drop, the body will close automatically
|
||||
}
|
||||
} else {
|
||||
return self.conn.read_keep_alive().map(Async::Ready);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn poll_read_head(&mut self) -> Poll<(), ::Error> {
|
||||
// can dispatch receive, or does it still care about, an incoming message?
|
||||
match self.dispatch.poll_ready() {
|
||||
Ok(Async::Ready(())) => (),
|
||||
Ok(Async::NotReady) => unreachable!("dispatch not ready when conn is"),
|
||||
Err(()) => {
|
||||
trace!("dispatch no longer receiving messages");
|
||||
self.close();
|
||||
return Ok(Async::Ready(()));
|
||||
}
|
||||
}
|
||||
// dispatch is ready for a message, try to read one
|
||||
match self.conn.read_head() {
|
||||
Ok(Async::Ready(Some((head, has_body)))) => {
|
||||
let body = if has_body {
|
||||
let (mut tx, rx) = ::proto::body::channel();
|
||||
let _ = tx.poll_ready(); // register this task if rx is dropped
|
||||
self.body_tx = Some(tx);
|
||||
Some(rx)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
self.dispatch.recv_msg(Ok((head, body)))?;
|
||||
Ok(Async::Ready(()))
|
||||
},
|
||||
Ok(Async::Ready(None)) => {
|
||||
// read eof, conn will start to shutdown automatically
|
||||
Ok(Async::Ready(()))
|
||||
}
|
||||
Ok(Async::NotReady) => Ok(Async::NotReady),
|
||||
Err(err) => {
|
||||
debug!("read_head error: {}", err);
|
||||
self.dispatch.recv_msg(Err(err))?;
|
||||
// if here, the dispatcher gave the user the error
|
||||
// somewhere else. we still need to shutdown, but
|
||||
// not as a second error.
|
||||
Ok(Async::Ready(()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn poll_write(&mut self) -> Poll<(), ::Error> {
|
||||
loop {
|
||||
if self.is_closing {
|
||||
return Ok(Async::Ready(()));
|
||||
} else if self.body_rx.is_none() && self.dispatch.should_poll() {
|
||||
if let Some((head, body)) = try_ready!(self.dispatch.poll_msg()) {
|
||||
self.conn.write_head(head, body.is_some());
|
||||
self.body_rx = body;
|
||||
} else {
|
||||
self.close();
|
||||
return Ok(Async::Ready(()));
|
||||
}
|
||||
} else if !self.conn.can_buffer_body() {
|
||||
try_ready!(self.poll_flush());
|
||||
} else if let Some(mut body) = self.body_rx.take() {
|
||||
let chunk = match body.poll()? {
|
||||
Async::Ready(Some(chunk)) => {
|
||||
self.body_rx = Some(body);
|
||||
chunk
|
||||
},
|
||||
Async::Ready(None) => {
|
||||
if self.conn.can_write_body() {
|
||||
self.conn.write_body(None)?;
|
||||
}
|
||||
continue;
|
||||
},
|
||||
Async::NotReady => {
|
||||
self.body_rx = Some(body);
|
||||
return Ok(Async::NotReady);
|
||||
}
|
||||
};
|
||||
|
||||
if self.conn.can_write_body() {
|
||||
assert!(self.conn.write_body(Some(chunk))?.is_ready());
|
||||
// This allows when chunk is `None`, or `Some([])`.
|
||||
} else if chunk.as_ref().len() == 0 {
|
||||
// ok
|
||||
} else {
|
||||
warn!("unexpected chunk when body cannot write");
|
||||
}
|
||||
} else {
|
||||
return Ok(Async::NotReady);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn poll_flush(&mut self) -> Poll<(), ::Error> {
|
||||
self.conn.flush().map_err(|err| {
|
||||
debug!("error writing: {}", err);
|
||||
err.into()
|
||||
})
|
||||
}
|
||||
|
||||
fn close(&mut self) {
|
||||
self.is_closing = true;
|
||||
self.conn.close_read();
|
||||
self.conn.close_write();
|
||||
}
|
||||
|
||||
fn is_done(&self) -> bool {
|
||||
if self.is_closing {
|
||||
return true;
|
||||
}
|
||||
|
||||
let read_done = self.conn.is_read_closed();
|
||||
|
||||
if !T::should_read_first() && read_done {
|
||||
// a client that cannot read may was well be done.
|
||||
true
|
||||
} else {
|
||||
let write_done = self.conn.is_write_closed() ||
|
||||
(!self.dispatch.should_poll() && self.body_rx.is_none());
|
||||
read_done && write_done
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
impl<D, Bs, I, B, T, K> Future for Dispatcher<D, Bs, I, B, T, K>
|
||||
where
|
||||
D: Dispatch<PollItem=MessageHead<T::Outgoing>, PollBody=Bs, RecvItem=MessageHead<T::Incoming>>,
|
||||
I: AsyncRead + AsyncWrite,
|
||||
B: AsRef<[u8]>,
|
||||
T: Http1Transaction,
|
||||
K: KeepAlive,
|
||||
Bs: Stream<Item=B, Error=::Error>,
|
||||
{
|
||||
type Item = ();
|
||||
type Error = ::Error;
|
||||
|
||||
#[inline]
|
||||
fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
|
||||
trace!("Dispatcher::poll");
|
||||
self.poll2().or_else(|e| {
|
||||
// An error means we're shutting down either way.
|
||||
// We just try to give the error to the user,
|
||||
// and close the connection with an Ok. If we
|
||||
// cannot give it to the user, then return the Err.
|
||||
self.dispatch.recv_msg(Err(e)).map(Async::Ready)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ===== impl Server =====
|
||||
|
||||
impl<S> Server<S> where S: Service {
|
||||
pub fn new(service: S) -> Server<S> {
|
||||
Server {
|
||||
in_flight: None,
|
||||
service: service,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, Bs> Dispatch for Server<S>
|
||||
where
|
||||
S: Service<Request=::Request, Response=::Response<Bs>, Error=::Error>,
|
||||
Bs: Stream<Error=::Error>,
|
||||
Bs::Item: AsRef<[u8]>,
|
||||
{
|
||||
type PollItem = MessageHead<StatusCode>;
|
||||
type PollBody = Bs;
|
||||
type RecvItem = RequestHead;
|
||||
|
||||
fn poll_msg(&mut self) -> Poll<Option<(Self::PollItem, Option<Self::PollBody>)>, ::Error> {
|
||||
if let Some(mut fut) = self.in_flight.take() {
|
||||
let resp = match fut.poll()? {
|
||||
Async::Ready(res) => res,
|
||||
Async::NotReady => {
|
||||
self.in_flight = Some(fut);
|
||||
return Ok(Async::NotReady);
|
||||
}
|
||||
};
|
||||
let (head, body) = ::proto::response::split(resp);
|
||||
Ok(Async::Ready(Some((head.into(), body))))
|
||||
} else {
|
||||
unreachable!("poll_msg shouldn't be called if no inflight");
|
||||
}
|
||||
}
|
||||
|
||||
fn recv_msg(&mut self, msg: ::Result<(Self::RecvItem, Option<Body>)>) -> ::Result<()> {
|
||||
let (msg, body) = msg?;
|
||||
let req = ::proto::request::from_wire(None, msg, body);
|
||||
self.in_flight = Some(self.service.call(req));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn poll_ready(&mut self) -> Poll<(), ()> {
|
||||
if self.in_flight.is_some() {
|
||||
Ok(Async::NotReady)
|
||||
} else {
|
||||
Ok(Async::Ready(()))
|
||||
}
|
||||
}
|
||||
|
||||
fn should_poll(&self) -> bool {
|
||||
self.in_flight.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
// ===== impl Client =====
|
||||
|
||||
|
||||
impl<B> Client<B> {
|
||||
pub fn new(rx: ClientRx<B>) -> Client<B> {
|
||||
Client {
|
||||
callback: None,
|
||||
rx: rx,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<B> Dispatch for Client<B>
|
||||
where
|
||||
B: Stream<Error=::Error>,
|
||||
B::Item: AsRef<[u8]>,
|
||||
{
|
||||
type PollItem = RequestHead;
|
||||
type PollBody = B;
|
||||
type RecvItem = ResponseHead;
|
||||
|
||||
fn poll_msg(&mut self) -> Poll<Option<(Self::PollItem, Option<Self::PollBody>)>, ::Error> {
|
||||
match self.rx.poll() {
|
||||
Ok(Async::Ready(Some(ClientMsg::Request(head, body, mut cb)))) => {
|
||||
// check that future hasn't been canceled already
|
||||
match cb.poll_cancel().expect("poll_cancel cannot error") {
|
||||
Async::Ready(()) => {
|
||||
trace!("request canceled");
|
||||
Ok(Async::Ready(None))
|
||||
},
|
||||
Async::NotReady => {
|
||||
self.callback = Some(cb);
|
||||
Ok(Async::Ready(Some((head, body))))
|
||||
}
|
||||
}
|
||||
},
|
||||
Ok(Async::Ready(Some(ClientMsg::Close))) |
|
||||
Ok(Async::Ready(None)) => {
|
||||
trace!("client tx closed");
|
||||
// user has dropped sender handle
|
||||
Ok(Async::Ready(None))
|
||||
},
|
||||
Ok(Async::NotReady) => return Ok(Async::NotReady),
|
||||
Err(()) => unreachable!("mpsc receiver cannot error"),
|
||||
}
|
||||
}
|
||||
|
||||
fn recv_msg(&mut self, msg: ::Result<(Self::RecvItem, Option<Body>)>) -> ::Result<()> {
|
||||
match msg {
|
||||
Ok((msg, body)) => {
|
||||
if let Some(cb) = self.callback.take() {
|
||||
let res = ::proto::response::from_wire(msg, body);
|
||||
let _ = cb.send(Ok(res));
|
||||
Ok(())
|
||||
} else {
|
||||
Err(::Error::Io(io::Error::new(io::ErrorKind::InvalidData, "response received without matching request")))
|
||||
}
|
||||
},
|
||||
Err(err) => {
|
||||
if let Some(cb) = self.callback.take() {
|
||||
let _ = cb.send(Err(err));
|
||||
Ok(())
|
||||
} else if let Ok(Async::Ready(Some(ClientMsg::Request(_, _, cb)))) = self.rx.poll() {
|
||||
let _ = cb.send(Err(err));
|
||||
Ok(())
|
||||
} else {
|
||||
Err(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn poll_ready(&mut self) -> Poll<(), ()> {
|
||||
match self.callback {
|
||||
Some(ref mut cb) => match cb.poll_cancel() {
|
||||
Ok(Async::Ready(())) => {
|
||||
trace!("callback receiver has dropped");
|
||||
Err(())
|
||||
},
|
||||
Ok(Async::NotReady) => Ok(Async::Ready(())),
|
||||
Err(_) => unreachable!("oneshot poll_cancel cannot error"),
|
||||
},
|
||||
None => Err(()),
|
||||
}
|
||||
}
|
||||
|
||||
fn should_poll(&self) -> bool {
|
||||
self.callback.is_none()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use futures::Sink;
|
||||
|
||||
use super::*;
|
||||
use mock::AsyncIo;
|
||||
use proto::ClientTransaction;
|
||||
|
||||
#[test]
|
||||
fn client_read_response_before_writing_request() {
|
||||
extern crate pretty_env_logger;
|
||||
let _ = pretty_env_logger::try_init();
|
||||
::futures::lazy(|| {
|
||||
let io = AsyncIo::new_buf(b"HTTP/1.1 200 OK\r\n\r\n".to_vec(), 100);
|
||||
let (mut tx, rx) = mpsc::channel(0);
|
||||
let conn = Conn::<_, ::Chunk, ClientTransaction>::new(io, Default::default());
|
||||
let mut dispatcher = Dispatcher::new(Client::new(rx), conn);
|
||||
|
||||
let req = RequestHead {
|
||||
version: ::HttpVersion::Http11,
|
||||
subject: ::proto::RequestLine::default(),
|
||||
headers: Default::default(),
|
||||
};
|
||||
let (res_tx, res_rx) = oneshot::channel();
|
||||
tx.start_send(ClientMsg::Request(req, None::<::Body>, res_tx)).unwrap();
|
||||
|
||||
dispatcher.poll().expect("dispatcher poll 1");
|
||||
dispatcher.poll().expect("dispatcher poll 2");
|
||||
let _res = res_rx.wait()
|
||||
.expect("callback poll")
|
||||
.expect("callback response");
|
||||
Ok::<(), ()>(())
|
||||
}).wait().unwrap();
|
||||
}
|
||||
}
|
||||
@@ -261,7 +261,7 @@ impl Buf for CrLf {
|
||||
mod tests {
|
||||
use bytes::{BufMut};
|
||||
|
||||
use proto::io::Cursor;
|
||||
use super::super::io::Cursor;
|
||||
use super::Encoder;
|
||||
|
||||
#[test]
|
||||
|
||||
510
src/proto/h1/io.rs
Normal file
510
src/proto/h1/io.rs
Normal file
@@ -0,0 +1,510 @@
|
||||
use std::collections::VecDeque;
|
||||
use std::fmt;
|
||||
use std::io;
|
||||
|
||||
use bytes::{Buf, BufMut, Bytes, BytesMut};
|
||||
use futures::{Async, Poll};
|
||||
use iovec::IoVec;
|
||||
use tokio_io::{AsyncRead, AsyncWrite};
|
||||
|
||||
use proto::{Http1Transaction, MessageHead};
|
||||
|
||||
const INIT_BUFFER_SIZE: usize = 8192;
|
||||
pub const DEFAULT_MAX_BUFFER_SIZE: usize = 8192 + 4096 * 100;
|
||||
const MAX_BUF_LIST_BUFFERS: usize = 16;
|
||||
|
||||
pub struct Buffered<T, B> {
|
||||
flush_pipeline: bool,
|
||||
io: T,
|
||||
max_buf_size: usize,
|
||||
read_blocked: bool,
|
||||
read_buf: BytesMut,
|
||||
write_buf: WriteBuf<B>,
|
||||
}
|
||||
|
||||
impl<T, B> fmt::Debug for Buffered<T, B>
|
||||
where
|
||||
B: Buf,
|
||||
{
|
||||
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
||||
f.debug_struct("Buffered")
|
||||
.field("read_buf", &self.read_buf)
|
||||
.field("write_buf", &self.write_buf)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl<T, B> Buffered<T, B>
|
||||
where
|
||||
T: AsyncRead + AsyncWrite,
|
||||
B: Buf,
|
||||
{
|
||||
pub fn new(io: T) -> Buffered<T, B> {
|
||||
Buffered {
|
||||
flush_pipeline: false,
|
||||
io: io,
|
||||
max_buf_size: DEFAULT_MAX_BUFFER_SIZE,
|
||||
read_buf: BytesMut::with_capacity(0),
|
||||
write_buf: WriteBuf::new(),
|
||||
read_blocked: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set_flush_pipeline(&mut self, enabled: bool) {
|
||||
self.flush_pipeline = enabled;
|
||||
self.write_buf.set_strategy(if enabled {
|
||||
Strategy::Flatten
|
||||
} else {
|
||||
Strategy::Queue
|
||||
});
|
||||
}
|
||||
|
||||
pub fn set_max_buf_size(&mut self, max: usize) {
|
||||
self.max_buf_size = max;
|
||||
self.write_buf.max_buf_size = max;
|
||||
}
|
||||
|
||||
pub fn read_buf(&self) -> &[u8] {
|
||||
self.read_buf.as_ref()
|
||||
}
|
||||
|
||||
pub fn write_buf_mut(&mut self) -> &mut Vec<u8> {
|
||||
let buf = self.write_buf.head_mut();
|
||||
buf.maybe_reset();
|
||||
&mut buf.bytes
|
||||
}
|
||||
|
||||
pub fn buffer(&mut self, buf: B) {
|
||||
self.write_buf.buffer(buf)
|
||||
}
|
||||
|
||||
pub fn can_buffer(&self) -> bool {
|
||||
self.flush_pipeline || self.write_buf.can_buffer()
|
||||
}
|
||||
|
||||
pub fn consume_leading_lines(&mut self) {
|
||||
if !self.read_buf.is_empty() {
|
||||
let mut i = 0;
|
||||
while i < self.read_buf.len() {
|
||||
match self.read_buf[i] {
|
||||
b'\r' | b'\n' => i += 1,
|
||||
_ => break,
|
||||
}
|
||||
}
|
||||
self.read_buf.split_to(i);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn parse<S: Http1Transaction>(&mut self) -> Poll<MessageHead<S::Incoming>, ::Error> {
|
||||
loop {
|
||||
match try!(S::parse(&mut self.read_buf)) {
|
||||
Some((head, len)) => {
|
||||
debug!("parsed {} headers ({} bytes)", head.headers.len(), len);
|
||||
return Ok(Async::Ready(head))
|
||||
},
|
||||
None => {
|
||||
if self.read_buf.capacity() >= self.max_buf_size {
|
||||
debug!("max_buf_size ({}) reached, closing", self.max_buf_size);
|
||||
return Err(::Error::TooLarge);
|
||||
}
|
||||
},
|
||||
}
|
||||
match try_ready!(self.read_from_io()) {
|
||||
0 => {
|
||||
trace!("parse eof");
|
||||
return Err(::Error::Incomplete);
|
||||
}
|
||||
_ => {},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn read_from_io(&mut self) -> Poll<usize, io::Error> {
|
||||
use bytes::BufMut;
|
||||
self.read_blocked = false;
|
||||
if self.read_buf.remaining_mut() < INIT_BUFFER_SIZE {
|
||||
self.read_buf.reserve(INIT_BUFFER_SIZE);
|
||||
}
|
||||
self.io.read_buf(&mut self.read_buf).map(|ok| {
|
||||
match ok {
|
||||
Async::Ready(n) => {
|
||||
debug!("read {} bytes", n);
|
||||
Async::Ready(n)
|
||||
},
|
||||
Async::NotReady => {
|
||||
self.read_blocked = true;
|
||||
Async::NotReady
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
pub fn io_mut(&mut self) -> &mut T {
|
||||
&mut self.io
|
||||
}
|
||||
|
||||
pub fn is_read_blocked(&self) -> bool {
|
||||
self.read_blocked
|
||||
}
|
||||
|
||||
pub fn flush(&mut self) -> Poll<(), io::Error> {
|
||||
if self.flush_pipeline && !self.read_buf.is_empty() {
|
||||
//Ok(())
|
||||
} else if self.write_buf.remaining() == 0 {
|
||||
try_nb!(self.io.flush());
|
||||
} else {
|
||||
loop {
|
||||
let n = try_ready!(self.io.write_buf(&mut self.write_buf));
|
||||
debug!("flushed {} bytes", n);
|
||||
if self.write_buf.remaining() == 0 {
|
||||
break;
|
||||
}
|
||||
}
|
||||
try_nb!(self.io.flush())
|
||||
}
|
||||
Ok(Async::Ready(()))
|
||||
}
|
||||
}
|
||||
|
||||
pub trait MemRead {
|
||||
fn read_mem(&mut self, len: usize) -> Poll<Bytes, io::Error>;
|
||||
}
|
||||
|
||||
impl<T, B> MemRead for Buffered<T, B>
|
||||
where
|
||||
T: AsyncRead + AsyncWrite,
|
||||
B: Buf,
|
||||
{
|
||||
fn read_mem(&mut self, len: usize) -> Poll<Bytes, io::Error> {
|
||||
trace!("Buffered.read_mem read_buf={}, wanted={}", self.read_buf.len(), len);
|
||||
if !self.read_buf.is_empty() {
|
||||
let n = ::std::cmp::min(len, self.read_buf.len());
|
||||
trace!("Buffered.read_mem read_buf is not empty, slicing {}", n);
|
||||
Ok(Async::Ready(self.read_buf.split_to(n).freeze()))
|
||||
} else {
|
||||
let n = try_ready!(self.read_from_io());
|
||||
Ok(Async::Ready(self.read_buf.split_to(::std::cmp::min(len, n)).freeze()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Cursor<T> {
|
||||
bytes: T,
|
||||
pos: usize,
|
||||
}
|
||||
|
||||
impl<T: AsRef<[u8]>> Cursor<T> {
|
||||
pub fn new(bytes: T) -> Cursor<T> {
|
||||
Cursor {
|
||||
bytes: bytes,
|
||||
pos: 0,
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub fn buf(&self) -> &[u8] {
|
||||
&self.bytes.as_ref()[self.pos..]
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub fn consume(&mut self, num: usize) {
|
||||
self.pos += num;
|
||||
}
|
||||
}
|
||||
|
||||
impl Cursor<Vec<u8>> {
|
||||
fn maybe_reset(&mut self) {
|
||||
if self.pos != 0 && self.remaining() == 0 {
|
||||
self.pos = 0;
|
||||
unsafe {
|
||||
self.bytes.set_len(0);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: AsRef<[u8]>> fmt::Debug for Cursor<T> {
|
||||
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
||||
f.debug_struct("Cursor")
|
||||
.field("pos", &self.pos)
|
||||
.field("len", &self.bytes.as_ref().len())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: AsRef<[u8]>> Buf for Cursor<T> {
|
||||
#[inline]
|
||||
fn remaining(&self) -> usize {
|
||||
self.bytes.as_ref().len() - self.pos
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn bytes(&self) -> &[u8] {
|
||||
self.buf()
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn advance(&mut self, cnt: usize) {
|
||||
self.consume(cnt)
|
||||
}
|
||||
}
|
||||
|
||||
// an internal buffer to collect writes before flushes
|
||||
struct WriteBuf<B> {
|
||||
buf: BufDeque<B>,
|
||||
max_buf_size: usize,
|
||||
strategy: Strategy,
|
||||
}
|
||||
|
||||
impl<B> WriteBuf<B> {
|
||||
fn new() -> WriteBuf<B> {
|
||||
WriteBuf {
|
||||
buf: BufDeque::new(),
|
||||
max_buf_size: DEFAULT_MAX_BUFFER_SIZE,
|
||||
strategy: Strategy::Queue,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
impl<B> WriteBuf<B>
|
||||
where
|
||||
B: Buf,
|
||||
{
|
||||
fn set_strategy(&mut self, strategy: Strategy) {
|
||||
self.strategy = strategy;
|
||||
}
|
||||
|
||||
fn buffer(&mut self, buf: B) {
|
||||
match self.strategy {
|
||||
Strategy::Flatten => {
|
||||
let head = self.head_mut();
|
||||
head.maybe_reset();
|
||||
head.bytes.put(buf);
|
||||
},
|
||||
Strategy::Queue => {
|
||||
self.buf.bufs.push_back(VecOrBuf::Buf(buf));
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn can_buffer(&self) -> bool {
|
||||
match self.strategy {
|
||||
Strategy::Flatten => {
|
||||
self.remaining() < self.max_buf_size
|
||||
},
|
||||
Strategy::Queue => {
|
||||
// for now, the simplest of heuristics
|
||||
self.buf.bufs.len() < MAX_BUF_LIST_BUFFERS
|
||||
&& self.remaining() < self.max_buf_size
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn head_mut(&mut self) -> &mut Cursor<Vec<u8>> {
|
||||
// this dance is brought to you, The Borrow Checker!
|
||||
|
||||
let reuse_back = if let Some(&VecOrBuf::Vec(_)) = self.buf.bufs.back() {
|
||||
true
|
||||
} else {
|
||||
false
|
||||
};
|
||||
|
||||
if !reuse_back {
|
||||
let head_buf = Cursor::new(Vec::with_capacity(INIT_BUFFER_SIZE));
|
||||
self.buf.bufs.push_back(VecOrBuf::Vec(head_buf));
|
||||
}
|
||||
if let Some(&mut VecOrBuf::Vec(ref mut v)) = self.buf.bufs.back_mut() {
|
||||
v
|
||||
} else {
|
||||
unreachable!("head_buf just pushed on back");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<B: Buf> fmt::Debug for WriteBuf<B> {
|
||||
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
||||
f.debug_struct("WriteBuf")
|
||||
.field("remaining", &self.remaining())
|
||||
.field("strategy", &self.strategy)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl<B: Buf> Buf for WriteBuf<B> {
|
||||
#[inline]
|
||||
fn remaining(&self) -> usize {
|
||||
self.buf.remaining()
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn bytes(&self) -> &[u8] {
|
||||
self.buf.bytes()
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn advance(&mut self, cnt: usize) {
|
||||
self.buf.advance(cnt)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn bytes_vec<'t>(&'t self, dst: &mut [&'t IoVec]) -> usize {
|
||||
self.buf.bytes_vec(dst)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum Strategy {
|
||||
Flatten,
|
||||
Queue,
|
||||
}
|
||||
|
||||
enum VecOrBuf<B> {
|
||||
Vec(Cursor<Vec<u8>>),
|
||||
Buf(B),
|
||||
}
|
||||
|
||||
impl<B: Buf> Buf for VecOrBuf<B> {
|
||||
#[inline]
|
||||
fn remaining(&self) -> usize {
|
||||
match *self {
|
||||
VecOrBuf::Vec(ref v) => v.remaining(),
|
||||
VecOrBuf::Buf(ref b) => b.remaining(),
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn bytes(&self) -> &[u8] {
|
||||
match *self {
|
||||
VecOrBuf::Vec(ref v) => v.bytes(),
|
||||
VecOrBuf::Buf(ref b) => b.bytes(),
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn advance(&mut self, cnt: usize) {
|
||||
match *self {
|
||||
VecOrBuf::Vec(ref mut v) => v.advance(cnt),
|
||||
VecOrBuf::Buf(ref mut b) => b.advance(cnt),
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn bytes_vec<'t>(&'t self, dst: &mut [&'t IoVec]) -> usize {
|
||||
match *self {
|
||||
VecOrBuf::Vec(ref v) => v.bytes_vec(dst),
|
||||
VecOrBuf::Buf(ref b) => b.bytes_vec(dst),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct BufDeque<T> {
|
||||
bufs: VecDeque<VecOrBuf<T>>,
|
||||
}
|
||||
|
||||
|
||||
impl<T> BufDeque<T> {
|
||||
fn new() -> BufDeque<T> {
|
||||
BufDeque {
|
||||
bufs: VecDeque::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Buf> Buf for BufDeque<T> {
|
||||
#[inline]
|
||||
fn remaining(&self) -> usize {
|
||||
self.bufs.iter()
|
||||
.map(|buf| buf.remaining())
|
||||
.sum()
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn bytes(&self) -> &[u8] {
|
||||
if let Some(buf) = self.bufs.front() {
|
||||
buf.bytes()
|
||||
} else {
|
||||
&[]
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn advance(&mut self, mut cnt: usize) {
|
||||
let mut maybe_reclaim = None;
|
||||
while cnt > 0 {
|
||||
{
|
||||
let front = &mut self.bufs[0];
|
||||
let rem = front.remaining();
|
||||
if rem > cnt {
|
||||
front.advance(cnt);
|
||||
return;
|
||||
} else {
|
||||
front.advance(rem);
|
||||
cnt -= rem;
|
||||
}
|
||||
}
|
||||
maybe_reclaim = self.bufs.pop_front();
|
||||
}
|
||||
|
||||
if let Some(VecOrBuf::Vec(v)) = maybe_reclaim {
|
||||
trace!("reclaiming write buf Vec");
|
||||
self.bufs.push_back(VecOrBuf::Vec(v));
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn bytes_vec<'t>(&'t self, dst: &mut [&'t IoVec]) -> usize {
|
||||
if dst.is_empty() {
|
||||
return 0;
|
||||
}
|
||||
let mut vecs = 0;
|
||||
for buf in &self.bufs {
|
||||
vecs += buf.bytes_vec(&mut dst[vecs..]);
|
||||
if vecs == dst.len() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
vecs
|
||||
}
|
||||
}
|
||||
|
||||
// TODO: Move tests to their own mod
|
||||
#[cfg(test)]
|
||||
use std::io::Read;
|
||||
|
||||
#[cfg(test)]
|
||||
impl<T: Read> MemRead for ::mock::AsyncIo<T> {
|
||||
fn read_mem(&mut self, len: usize) -> Poll<Bytes, io::Error> {
|
||||
let mut v = vec![0; len];
|
||||
let n = try_nb!(self.read(v.as_mut_slice()));
|
||||
Ok(Async::Ready(BytesMut::from(&v[..n]).freeze()))
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_iobuf_write_empty_slice() {
|
||||
use mock::{AsyncIo, Buf as MockBuf};
|
||||
|
||||
let mut mock = AsyncIo::new(MockBuf::new(), 256);
|
||||
mock.error(io::Error::new(io::ErrorKind::Other, "logic error"));
|
||||
|
||||
let mut io_buf = Buffered::<_, Cursor<Vec<u8>>>::new(mock);
|
||||
|
||||
// underlying io will return the logic error upon write,
|
||||
// so we are testing that the io_buf does not trigger a write
|
||||
// when there is nothing to flush
|
||||
io_buf.flush().expect("should short-circuit flush");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_reads_until_blocked() {
|
||||
use mock::{AsyncIo, Buf as MockBuf};
|
||||
// missing last line ending
|
||||
let raw = "HTTP/1.1 200 OK\r\n";
|
||||
|
||||
let mock = AsyncIo::new(MockBuf::wrap(raw.into()), raw.len());
|
||||
let mut buffered = Buffered::<_, Cursor<Vec<u8>>>::new(mock);
|
||||
assert_eq!(buffered.parse::<::proto::ClientTransaction>().unwrap(), Async::NotReady);
|
||||
assert!(buffered.io.blocked());
|
||||
}
|
||||
@@ -1,8 +1,12 @@
|
||||
pub use self::conn::{Conn, KeepAlive, KA};
|
||||
pub use self::decode::Decoder;
|
||||
pub use self::encode::{EncodedBuf, Encoder};
|
||||
|
||||
mod conn;
|
||||
mod date;
|
||||
mod decode;
|
||||
pub mod dispatch;
|
||||
mod encode;
|
||||
pub mod parse;
|
||||
mod io;
|
||||
pub mod role;
|
||||
|
||||
|
||||
Reference in New Issue
Block a user