test(h1): fix flaky streaming test

src/mock.rs

@@ -1,7 +1,8 @@
 use std::cmp;
 use std::io::{self, Read, Write};
 
-use futures::Poll;
+use bytes::Buf as BufTrait;
+use futures::{Async, Poll};
 use tokio_io::{AsyncRead, AsyncWrite};
 
 #[derive(Debug)]
@@ -11,10 +12,6 @@ pub struct Buf {
 }
 
 impl Buf {
-    pub fn new() -> Buf {
-        Buf::wrap(vec![])
-    }
-
     pub fn wrap(vec: Vec<u8>) -> Buf {
         Buf {
             vec: vec,
@@ -63,23 +60,27 @@ impl Read for Buf {
     }
 }
 
+const READ_VECS_CNT: usize = 64;
+
 #[derive(Debug)]
 pub struct AsyncIo<T> {
-    inner: T,
+    blocked: bool,
     bytes_until_block: usize,
     error: Option<io::Error>,
-    blocked: bool,
     flushed: bool,
+    inner: T,
+    max_read_vecs: usize,
 }
 
 impl<T> AsyncIo<T> {
     pub fn new(inner: T, bytes: usize) -> AsyncIo<T> {
         AsyncIo {
-            inner: inner,
+            blocked: false,
             bytes_until_block: bytes,
             error: None,
             flushed: false,
-            blocked: false,
+            inner: inner,
+            max_read_vecs: READ_VECS_CNT,
         }
     }
 
@@ -90,6 +91,22 @@ impl<T> AsyncIo<T> {
     pub fn error(&mut self, err: io::Error) {
         self.error = Some(err);
     }
+
+    pub fn max_read_vecs(&mut self, cnt: usize) {
+        assert!(cnt <= READ_VECS_CNT);
+        self.max_read_vecs = cnt;
+    }
+
+    #[cfg(feature = "tokio-proto")]
+    //TODO: fix proto::conn::tests to not use tokio-proto API,
+    //and then this cfg flag go away
+    pub fn flushed(&self) -> bool {
+        self.flushed
+    }
+
+    pub fn blocked(&self) -> bool {
+        self.blocked
+    }
 }
 
 impl AsyncIo<Buf> {
@@ -103,16 +120,17 @@ impl AsyncIo<Buf> {
     pub fn new_eof() -> AsyncIo<Buf> {
         AsyncIo::new(Buf::wrap(Vec::new().into()), 1)
     }
+}
 
-    #[cfg(feature = "tokio-proto")]
-    //TODO: fix proto::conn::tests to not use tokio-proto API,
-    //and then this cfg flag go away
-    pub fn flushed(&self) -> bool {
-        self.flushed
-    }
+impl<T: Read + Write> AsyncIo<T> {
+    fn write_no_vecs<B: BufTrait>(&mut self, buf: &mut B) -> Poll<usize, io::Error> {
+        if !buf.has_remaining() {
+            return Ok(Async::Ready(0));
+        }
 
-    pub fn blocked(&self) -> bool {
-        self.blocked
+        let n = try_nb!(self.write(buf.bytes()));
+        buf.advance(n);
+        Ok(Async::Ready(n))
     }
 }
 
@@ -170,12 +188,14 @@ impl<T: Read + Write> AsyncWrite for AsyncIo<T> {
         Ok(().into())
     }
 
-    fn write_buf<B: ::bytes::Buf>(&mut self, buf: &mut B) -> Poll<usize, io::Error> {
-        use futures::Async;
+    fn write_buf<B: BufTrait>(&mut self, buf: &mut B) -> Poll<usize, io::Error> {
+        if self.max_read_vecs == 0 {
+            return self.write_no_vecs(buf);
+        }
         let r = {
             static DUMMY: &[u8] = &[0];
-            let mut bufs = [From::from(DUMMY); 64];
-            let i = ::bytes::Buf::bytes_vec(&buf, &mut bufs);
+            let mut bufs = [From::from(DUMMY); READ_VECS_CNT];
+            let i = ::bytes::Buf::bytes_vec(&buf, &mut bufs[..self.max_read_vecs]);
             let mut n = 0;
             let mut ret = Ok(0);
             for iovec in &bufs[..i] {

src/proto/h1/io.rs

@@ -440,6 +440,9 @@ impl<T: Buf> Buf for BufDeque<T> {
                 return buf.bytes();
             }
         }
+        if let Some(ref buf) = self.bufs.front() {
+            return buf.bytes();
+        }
         &[]
     }
 
@@ -483,42 +486,54 @@ impl<T: Buf> Buf for BufDeque<T> {
     }
 }
 
-// TODO: Move tests to their own mod
 #[cfg(test)]
-use std::io::Read;
+mod tests {
+    use super::*;
+    use std::io::Read;
+    use mock::AsyncIo;
 
-#[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()))
+    #[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 iobuf_write_empty_slice() {
+        let mut mock = AsyncIo::new_buf(vec![], 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 parse_reads_until_blocked() {
+        // missing last line ending
+        let raw = "HTTP/1.1 200 OK\r\n";
+
+        let mock = AsyncIo::new_buf(raw, raw.len());
+        let mut buffered = Buffered::<_, Cursor<Vec<u8>>>::new(mock);
+        assert_eq!(buffered.parse::<::proto::ClientTransaction>().unwrap(), Async::NotReady);
+        assert!(buffered.io.blocked());
+    }
+
+    #[test]
+    fn write_buf_skips_empty_bufs() {
+        let mut mock = AsyncIo::new_buf(vec![], 1024);
+        mock.max_read_vecs(0); // disable vectored IO
+        let mut buffered = Buffered::<_, Cursor<Vec<u8>>>::new(mock);
+
+        buffered.buffer(Cursor::new(Vec::new()));
+        buffered.buffer(Cursor::new(b"hello".to_vec()));
+        buffered.flush().unwrap();
+        assert_eq!(buffered.io, b"hello");
     }
 }
-
-#[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());
-}