diff --git a/src/client/mod.rs b/src/client/mod.rs index ea2615af..3fcd9e7a 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -36,6 +36,8 @@ mod dns; mod pool; #[cfg(feature = "compat")] pub mod compat; +#[cfg(test)] +mod tests; /// A Client to make outgoing HTTP requests. pub struct Client { diff --git a/src/client/tests.rs b/src/client/tests.rs new file mode 100644 index 00000000..6aecb655 --- /dev/null +++ b/src/client/tests.rs @@ -0,0 +1,47 @@ +extern crate pretty_env_logger; + +use futures::Async; +use futures::future::poll_fn; +use tokio::reactor::Core; + +use mock::MockConnector; +use super::*; + +#[test] +fn retryable_request() { + let _ = pretty_env_logger::try_init(); + let mut core = Core::new().unwrap(); + + let mut connector = MockConnector::new(); + + let sock1 = connector.mock("http://mock.local/a"); + let sock2 = connector.mock("http://mock.local/b"); + + let client = Client::configure() + .connector(connector) + .build(&core.handle()); + + + { + let res1 = client.get("http://mock.local/a".parse().unwrap()); + let srv1 = poll_fn(|| { + try_ready!(sock1.read(&mut [0u8; 512])); + try_ready!(sock1.write(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n")); + Ok(Async::Ready(())) + }); + core.run(res1.join(srv1)).expect("res1"); + } + drop(sock1); + + let res2 = client.get("http://mock.local/b".parse().unwrap()) + .map(|res| { + assert_eq!(res.status().as_u16(), 222); + }); + let srv2 = poll_fn(|| { + try_ready!(sock2.read(&mut [0u8; 512])); + try_ready!(sock2.write(b"HTTP/1.1 222 OK\r\nContent-Length: 0\r\n\r\n")); + Ok(Async::Ready(())) + }); + + core.run(res2.join(srv2)).expect("res2"); +} diff --git a/src/mock.rs b/src/mock.rs index 1eb8125d..2e88b10d 100644 --- a/src/mock.rs +++ b/src/mock.rs @@ -1,26 +1,33 @@ +use std::cell::RefCell; +use std::collections::HashMap; use std::cmp; use std::io::{self, Read, Write}; +use std::rc::Rc; -use bytes::Buf as BufTrait; +use bytes::Buf; use futures::{Async, Poll}; +use futures::task::{self, Task}; use tokio_io::{AsyncRead, AsyncWrite}; +use tokio_service::Service; + +use ::Uri; #[derive(Debug)] -pub struct Buf { +pub struct MockCursor { vec: Vec, pos: usize, } -impl Buf { - pub fn wrap(vec: Vec) -> Buf { - Buf { +impl MockCursor { + pub fn wrap(vec: Vec) -> MockCursor { + MockCursor { vec: vec, pos: 0, } } } -impl ::std::ops::Deref for Buf { +impl ::std::ops::Deref for MockCursor { type Target = [u8]; fn deref(&self) -> &[u8] { @@ -28,19 +35,19 @@ impl ::std::ops::Deref for Buf { } } -impl AsRef<[u8]> for Buf { +impl AsRef<[u8]> for MockCursor { fn as_ref(&self) -> &[u8] { &self.vec } } -impl> PartialEq for Buf { +impl> PartialEq for MockCursor { fn eq(&self, other: &S) -> bool { self.vec == other.as_ref() } } -impl Write for Buf { +impl Write for MockCursor { fn write(&mut self, data: &[u8]) -> io::Result { self.vec.extend(data); Ok(data.len()) @@ -51,7 +58,7 @@ impl Write for Buf { } } -impl Read for Buf { +impl Read for MockCursor { fn read(&mut self, buf: &mut [u8]) -> io::Result { (&self.vec[self.pos..]).read(buf).map(|n| { self.pos += n; @@ -71,6 +78,8 @@ pub struct AsyncIo { inner: T, max_read_vecs: usize, num_writes: usize, + park_tasks: bool, + task: Option, } impl AsyncIo { @@ -83,11 +92,17 @@ impl AsyncIo { inner: inner, max_read_vecs: READ_VECS_CNT, num_writes: 0, + park_tasks: false, + task: None, } } pub fn block_in(&mut self, bytes: usize) { self.bytes_until_block = bytes; + + if let Some(task) = self.task.take() { + task.notify(); + } } pub fn error(&mut self, err: io::Error) { @@ -99,6 +114,10 @@ impl AsyncIo { self.max_read_vecs = cnt; } + pub fn park_tasks(&mut self, enabled: bool) { + self.park_tasks = enabled; + } + #[cfg(feature = "tokio-proto")] //TODO: fix proto::conn::tests to not use tokio-proto API, //and then this cfg flag go away @@ -113,23 +132,39 @@ impl AsyncIo { pub fn num_writes(&self) -> usize { self.num_writes } + + fn would_block(&mut self) -> io::Error { + self.blocked = true; + if self.park_tasks { + self.task = Some(task::current()); + } + io::ErrorKind::WouldBlock.into() + } + } -impl AsyncIo { - pub fn new_buf>>(buf: T, bytes: usize) -> AsyncIo { - AsyncIo::new(Buf::wrap(buf.into()), bytes) +impl AsyncIo { + pub fn new_buf>>(buf: T, bytes: usize) -> AsyncIo { + AsyncIo::new(MockCursor::wrap(buf.into()), bytes) } #[cfg(feature = "tokio-proto")] //TODO: fix proto::conn::tests to not use tokio-proto API, //and then this cfg flag go away - pub fn new_eof() -> AsyncIo { - AsyncIo::new(Buf::wrap(Vec::new().into()), 1) + pub fn new_eof() -> AsyncIo { + AsyncIo::new(MockCursor::wrap(Vec::new().into()), 1) + } + + fn close(&mut self) { + self.block_in(1); + assert_eq!(self.inner.vec.len(), self.inner.pos); + self.inner.vec.truncate(0); + self.inner.pos = 0; } } impl AsyncIo { - fn write_no_vecs(&mut self, buf: &mut B) -> Poll { + fn write_no_vecs(&mut self, buf: &mut B) -> Poll { if !buf.has_remaining() { return Ok(Async::Ready(0)); } @@ -153,8 +188,7 @@ impl Read for AsyncIo { if let Some(err) = self.error.take() { Err(err) } else if self.bytes_until_block == 0 { - self.blocked = true; - Err(io::ErrorKind::WouldBlock.into()) + Err(self.would_block()) } else { let n = cmp::min(self.bytes_until_block, buf.len()); let n = try!(self.inner.read(&mut buf[..n])); @@ -168,11 +202,13 @@ impl Write for AsyncIo { fn write(&mut self, data: &[u8]) -> io::Result { self.num_writes += 1; if let Some(err) = self.error.take() { + trace!("AsyncIo::write error"); Err(err) } else if self.bytes_until_block == 0 { - Err(io::ErrorKind::WouldBlock.into()) + trace!("AsyncIo::write would block"); + Err(self.would_block()) } else { - trace!("AsyncIo::write() block_in = {}, data.len() = {}", self.bytes_until_block, data.len()); + trace!("AsyncIo::write; {} bytes", data.len()); self.flushed = false; let n = cmp::min(self.bytes_until_block, data.len()); let n = try!(self.inner.write(&data[..n])); @@ -195,14 +231,14 @@ impl AsyncWrite for AsyncIo { Ok(().into()) } - fn write_buf(&mut self, buf: &mut B) -> Poll { + fn write_buf(&mut self, buf: &mut B) -> Poll { if self.max_read_vecs == 0 { return self.write_no_vecs(buf); } let r = { static DUMMY: &[u8] = &[0]; let mut bufs = [From::from(DUMMY); READ_VECS_CNT]; - let i = ::bytes::Buf::bytes_vec(&buf, &mut bufs[..self.max_read_vecs]); + let i = Buf::bytes_vec(&buf, &mut bufs[..self.max_read_vecs]); let mut n = 0; let mut ret = Ok(0); // each call to write() will increase our count, but we assume @@ -231,7 +267,7 @@ impl AsyncWrite for AsyncIo { }; match r { Ok(n) => { - ::bytes::Buf::advance(buf, n); + Buf::advance(buf, n); Ok(Async::Ready(n)) } Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { @@ -242,10 +278,153 @@ impl AsyncWrite for AsyncIo { } } -impl ::std::ops::Deref for AsyncIo { +impl ::std::ops::Deref for AsyncIo { type Target = [u8]; fn deref(&self) -> &[u8] { &self.inner } } + +pub struct Duplex { + inner: Rc>, +} + +struct DuplexInner { + handle_read_task: Option, + read: AsyncIo, + write: AsyncIo, +} + +impl Read for Duplex { + fn read(&mut self, buf: &mut [u8]) -> io::Result { + self.inner.borrow_mut().read.read(buf) + } +} + +impl Write for Duplex { + fn write(&mut self, buf: &[u8]) -> io::Result { + if let Some(task) = self.inner.borrow_mut().handle_read_task.take() { + trace!("waking DuplexHandle read"); + task.notify(); + } + self.inner.borrow_mut().write.write(buf) + } + + fn flush(&mut self) -> io::Result<()> { + self.inner.borrow_mut().write.flush() + } +} + +impl AsyncRead for Duplex { + +} + +impl AsyncWrite for Duplex { + fn shutdown(&mut self) -> Poll<(), io::Error> { + Ok(().into()) + } + + fn write_buf(&mut self, buf: &mut B) -> Poll { + if let Some(task) = self.inner.borrow_mut().handle_read_task.take() { + task.notify(); + } + self.inner.borrow_mut().write.write_buf(buf) + } +} + +pub struct DuplexHandle { + inner: Rc>, +} + +impl DuplexHandle { + pub fn read(&self, buf: &mut [u8]) -> Poll { + let mut inner = self.inner.borrow_mut(); + assert!(buf.len() >= inner.write.inner.len()); + if inner.write.inner.is_empty() { + trace!("DuplexHandle read parking"); + inner.handle_read_task = Some(task::current()); + return Ok(Async::NotReady); + } + inner.write.inner.vec.truncate(0); + Ok(Async::Ready(inner.write.inner.len())) + } + + pub fn write(&self, bytes: &[u8]) -> Poll { + let mut inner = self.inner.borrow_mut(); + assert!(inner.read.inner.vec.is_empty()); + assert_eq!(inner.read.inner.pos, 0); + inner + .read + .inner + .vec + .extend(bytes); + inner.read.block_in(bytes.len()); + Ok(Async::Ready(bytes.len())) + } +} + +impl Drop for DuplexHandle { + fn drop(&mut self) { + trace!("mock duplex handle drop"); + let mut inner = self.inner.borrow_mut(); + inner.read.close(); + inner.write.close(); + } +} + +pub struct MockConnector { + mocks: RefCell>>, +} + +impl MockConnector { + pub fn new() -> MockConnector { + MockConnector { + mocks: RefCell::new(HashMap::new()), + } + } + + pub fn mock(&mut self, key: &str) -> DuplexHandle { + let key = key.to_owned(); + let mut inner = DuplexInner { + handle_read_task: None, + read: AsyncIo::new_buf(Vec::new(), 0), + write: AsyncIo::new_buf(Vec::new(), ::std::usize::MAX), + }; + + inner.read.park_tasks(true); + inner.write.park_tasks(true); + + let inner = Rc::new(RefCell::new(inner)); + + let duplex = Duplex { + inner: inner.clone(), + }; + let handle = DuplexHandle { + inner: inner, + }; + + self.mocks.borrow_mut().entry(key) + .or_insert(Vec::new()) + .push(duplex); + + handle + } +} + +impl Service for MockConnector { + type Request = Uri; + type Response = Duplex; + type Error = io::Error; + type Future = ::futures::future::FutureResult; + + fn call(&self, uri: Uri) -> Self::Future { + use futures::future; + trace!("mock connect: {:?}", uri.as_ref()); + let mock = self.mocks.borrow_mut() + .get_mut(uri.as_ref()) + .expect(&format!("unknown mocks uri: {:?}", uri.as_ref())) + .remove(0); + future::ok(mock) + } +} diff --git a/tests/client.rs b/tests/client.rs index bbdb1332..76a90c06 100644 --- a/tests/client.rs +++ b/tests/client.rs @@ -584,72 +584,6 @@ fn client_keep_alive() { core.run(res.join(rx).map(|r| r.0)).unwrap(); } -#[test] -fn client_keep_alive_connreset() { - use std::sync::mpsc; - extern crate pretty_env_logger; - let _ = pretty_env_logger::try_init(); - - let server = TcpListener::bind("127.0.0.1:0").unwrap(); - let addr = server.local_addr().unwrap(); - let mut core = Core::new().unwrap(); - let handle = core.handle(); - - let client = client(&handle); - - let (tx1, rx1) = oneshot::channel(); - let (tx2, rx2) = mpsc::channel(); - thread::spawn(move || { - let mut sock = server.accept().unwrap().0; - sock.set_read_timeout(Some(Duration::from_secs(5))).unwrap(); - sock.set_write_timeout(Some(Duration::from_secs(5))).unwrap(); - let mut buf = [0; 4096]; - sock.read(&mut buf).expect("read 1"); - sock.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n").expect("write 1"); - - // Wait for client to indicate it is done processing the first request - // This is what seem to trigger the race condition -- without it client notices - // connection is closed while processing the first request. - let _ = rx2.recv(); - let _r = sock.shutdown(std::net::Shutdown::Both); - - // Let client know it can try to reuse the connection - let _ = tx1.send(()); - - // use sock2 so that sock isn't dropped yet - let mut sock2 = server.accept().unwrap().0; - sock2.set_read_timeout(Some(Duration::from_secs(5))).unwrap(); - sock2.set_write_timeout(Some(Duration::from_secs(5))).unwrap(); - let mut buf = [0; 4096]; - sock2.read(&mut buf).expect("read 2"); - sock2.write_all(b"HTTP/1.1 222 OK\r\nContent-Length: 0\r\n\r\n").expect("write 2"); - }); - - - let res = client.get(format!("http://{}/a", addr).parse().unwrap()); - core.run(res).unwrap(); - - let _ = tx2.send(()); - - let rx = rx1.map_err(|_| hyper::Error::Io(io::Error::new(io::ErrorKind::Other, "thread panicked"))); - core.run(rx).unwrap(); - - let t = Timeout::new(Duration::from_millis(100), &handle).unwrap(); - let res = client.get(format!("http://{}/b", addr).parse().unwrap()) - .map(|res| { - assert_eq!(res.status().as_u16(), 222); - }); - - let fut = res.select2(t).then(|result| match result { - Ok(Either::A((resp, _))) => Ok(resp), - Err(Either::A((err, _))) => Err(err), - Ok(Either::B(_)) | - Err(Either::B(_)) => Err(hyper::Error::Timeout), - }); - - core.run(fut).expect("req 2"); -} - #[test] fn client_keep_alive_extra_body() { let _ = pretty_env_logger::try_init();