Respect flow control when flushing writes (#26)

Previously, writing a body did not respect HTTP/2.0 flow control. This
resulted in unbounded buffering. This patch updates the library to respect
flow control when writing the streaming body.

This also adds tests 🎉

Author: carllerche

Diff for: .travis.yml

@@ -13,8 +13,12 @@ matrix:
     - rust: nightly
 
 script:
+  # Run lib and doc tests
   - cargo test
 
+  # Run integration tests
+  - cargo test -p tests
+
 deploy:
   provider:  pages
   skip_cleanup: true

Diff for: Cargo.toml

@@ -8,6 +8,12 @@ authors = [
 description = "Provides an HTTP/2.0 Tower Service"
 publish = false
 
+[workspace]
+
+members = [
+    "tests",
+]
+
 [dependencies]
 bytes = "0.4"
 futures = "0.1"

Diff for: src/buf.rs

@@ -0,0 +1,39 @@
+use bytes::Buf;
+
+pub struct SendBuf<T>
+{
+    inner: Option<T>,
+}
+
+impl<T: Buf> SendBuf<T> {
+    pub fn new(buf: T) -> SendBuf<T> {
+        SendBuf { inner: Some(buf) }
+    }
+
+    pub fn none() -> SendBuf<T> {
+        SendBuf { inner: None }
+    }
+}
+
+impl<T: Buf> Buf for SendBuf<T> {
+    fn remaining(&self) -> usize {
+        match self.inner {
+            Some(ref v) => v.remaining(),
+            None => 0,
+        }
+    }
+
+    fn bytes(&self) -> &[u8] {
+        match self.inner {
+            Some(ref v) => v.bytes(),
+            None => &[],
+        }
+    }
+
+    fn advance(&mut self, cnt: usize) {
+        match self.inner {
+            Some(ref mut v) => v.advance(cnt),
+            None => {}
+        }
+    }
+}

Diff for: src/client/background.rs

@@ -1,6 +1,8 @@
 use Body;
+use buf::SendBuf;
 use flush::Flush;
 
+use bytes::IntoBuf;
 use futures::{Future, Poll};
 use h2::client::Connection;
 use tokio_io::{AsyncRead, AsyncWrite};
@@ -18,7 +20,7 @@ where S: Body,
 enum Task<T, S>
 where S: Body,
 {
-    Connection(Connection<T, S::Data>),
+    Connection(Connection<T, SendBuf<<S::Data as IntoBuf>::Buf>>),
     Flush(Flush<S>),
 }
 
@@ -27,7 +29,7 @@ where S: Body,
 impl<T, S> Background<T, S>
 where S: Body,
 {
-    pub(crate) fn connection(connection: Connection<T, S::Data>) -> Self {
+    pub(crate) fn connection(connection: Connection<T, SendBuf<<S::Data as IntoBuf>::Buf>>) -> Self {
         let task = Task::Connection(connection);
         Background { task }
     }

Diff for: src/client/connection.rs

@@ -1,5 +1,6 @@
 use {Body, RecvBody};
 use super::Background;
+use buf::SendBuf;
 use flush::Flush;
 
 use bytes::IntoBuf;
@@ -18,7 +19,7 @@ use std::marker::PhantomData;
 pub struct Connection<T, E, S>
 where S: Body,
 {
-    client: SendRequest<S::Data>,
+    client: SendRequest<SendBuf<<S::Data as IntoBuf>::Buf>>,
     executor: E,
     _p: PhantomData<(T, S)>,
 }
@@ -27,7 +28,7 @@ where S: Body,
 pub struct Handshake<T, E, S>
 where S: Body,
 {
-    inner: h2::client::Handshake<T, S::Data>,
+    inner: h2::client::Handshake<T, SendBuf<<S::Data as IntoBuf>::Buf>>,
     executor: E,
 }
 
@@ -80,7 +81,9 @@ where S: Body,
       T: AsyncRead + AsyncWrite,
 {
     /// Builds Connection on an H2 client connection.
-    pub fn new(client: SendRequest<S::Data>, executor: E) -> Self {
+    pub(crate) fn new(client: SendRequest<SendBuf<<S::Data as IntoBuf>::Buf>>, executor: E)
+        -> Self
+    {
         let _p = PhantomData;
 
         Connection {
@@ -195,7 +198,7 @@ where T: AsyncRead + AsyncWrite,
       S: Body,
 {
     /// Start an HTTP/2.0 handshake with the provided builder
-    pub fn new(io: T, executor: E, builder: &Builder) -> Self {
+    pub(crate) fn new(io: T, executor: E, builder: &Builder) -> Self {
         let inner = builder.handshake(io);
 
         Handshake {
@@ -254,9 +257,9 @@ impl From<h2::Reason> for Error {
 impl fmt::Display for Error {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self.kind {
-            Kind::Inner(ref h2) => 
+            Kind::Inner(ref h2) =>
                 write!(f, "Error caused by underlying HTTP/2 error: {}", h2),
-            Kind::Spawn => 
+            Kind::Spawn =>
                 write!(f, "Error spawning background task"),
         }
     }
@@ -292,12 +295,12 @@ impl fmt::Display for HandshakeError {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match *self {
             HandshakeError::Proto(ref h2) =>
-                write!(f, 
+                write!(f,
                     "An error occurred while attempting to perform the HTTP/2 \
-                    handshake: {}", 
+                    handshake: {}",
                     h2),
             HandshakeError::Execute =>
-                write!(f, 
+                write!(f,
                     "An error occurred while attempting to execute a worker \
                      task."),
         }
@@ -321,5 +324,4 @@ impl error::Error for HandshakeError {
                 "error attempting to execute a worker task",
         }
     }
-
 }

Diff for: src/flush.rs

@@ -1,5 +1,7 @@
 use Body;
+use buf::SendBuf;
 
+use bytes::IntoBuf;
 use futures::{Future, Poll, Async};
 use h2::{self, SendStream};
 use http::HeaderMap;
@@ -8,7 +10,7 @@ use http::HeaderMap;
 pub(crate) struct Flush<S>
 where S: Body,
 {
-    h2: SendStream<S::Data>,
+    h2: SendStream<SendBuf<<S::Data as IntoBuf>::Buf>>,
     body: S,
     state: FlushState,
 }
@@ -29,7 +31,9 @@ enum DataOrTrailers<B> {
 impl<S> Flush<S>
 where S: Body,
 {
-    pub fn new(src: S, dst: SendStream<S::Data>) -> Self {
+    pub fn new(src: S, dst: SendStream<SendBuf<<S::Data as IntoBuf>::Buf>>)
+        -> Self
+    {
         Flush {
             h2: dst,
             body: src,
@@ -39,41 +43,71 @@ where S: Body,
 
     /// Try to flush the body.
     fn poll_complete(&mut self) -> Poll<(), h2::Error> {
-        let mut first = try_ready!(self.poll_body());
+        use self::DataOrTrailers::*;
 
         loop {
-            if let Some(DataOrTrailers::Data(buf)) = first {
-                let second = self.poll_body()?;
-                let eos = if let Async::Ready(None) = second {
-                    true
-                } else {
-                    false
-                };
-                self.h2.send_data(buf, eos)?;
-                if eos {
+            match try_ready!(self.poll_body()) {
+                Some(Data(buf)) => {
+                    let eos = self.body.is_end_stream();
+
+                    self.h2.send_data(SendBuf::new(buf.into_buf()), eos)?;
+
+                    if eos {
+                        self.state = FlushState::Done;
+                        return Ok(Async::Ready(()));
+                    }
+                }
+                Some(Trailers(trailers)) => {
+                    self.h2.send_trailers(trailers)?;
+                    return Ok(Async::Ready(()));
+                }
+                None => {
+                    // If this is hit, then an EOS was not reached via the other
+                    // paths. So, we must send an empty data frame with EOS.
+                    self.h2.send_data(SendBuf::none(), true)?;
+
                     return Ok(Async::Ready(()));
-                } else if let Async::Ready(item) = second {
-                    first = item;
-                } else {
-                    return Ok(Async::NotReady);
                 }
-            } else if let Some(DataOrTrailers::Trailers(trailers)) = first {
-                self.h2.send_trailers(trailers)?;
-                return Ok(Async::Ready(()));
-            } else {
-                return Ok(Async::Ready(()));
             }
         }
     }
 
     /// Get the next message to write, either a data frame or trailers.
-    fn poll_body(&mut self) -> Poll<Option<DataOrTrailers<S::Data>>, h2::Error> {
+    fn poll_body(&mut self)
+        -> Poll<Option<DataOrTrailers<S::Data>>, h2::Error>
+    {
         loop {
             match self.state {
                 FlushState::Data => {
+                    // Before trying to poll the next chunk, we have to see if
+                    // the h2 connection has capacity. We do this by requesting
+                    // a single byte (since we don't know how big the next chunk
+                    // will be.
+                    self.h2.reserve_capacity(1);
+
+                    if self.h2.capacity() == 0 {
+                        // TODO: The loop should not be needed once
+                        // carllerche/h2#270 is fixed.
+                        loop {
+                            match try_ready!(self.h2.poll_capacity()) {
+                                Some(0) => {}
+                                Some(_) => break,
+                                None => {
+                                    debug!("connection closed early");
+                                    // The error shouldn't really matter at this
+                                    // point as the peer has disconnected, the
+                                    // error will be discarded anyway.
+                                    return Err(h2::Reason::INTERNAL_ERROR.into());
+                                }
+                            }
+                        }
+                    }
+
                     if let Some(data) = try_ready!(self.body.poll_data()) {
                         return Ok(Async::Ready(Some(DataOrTrailers::Data(data))));
                     } else {
+                        // Release all capacity back to the connection
+                        self.h2.reserve_capacity(0);
                         self.state = FlushState::Trailers;
                     }
                 }

Diff for: src/lib.rs

@@ -14,6 +14,7 @@ pub mod client;
 pub mod server;
 
 mod body;
+mod buf;
 mod flush;
 mod recv_body;
 mod service;

Diff for: src/server/mod.rs

@@ -1,7 +1,9 @@
 use {flush, Body, RecvBody};
+use buf::SendBuf;
 
 use tower_service::{NewService, Service};
 
+use bytes::IntoBuf;
 use futures::{Async, Future, Poll, Stream};
 use futures::future::{Executor, Either, Join, MapErr};
 use h2::{self, Reason};
@@ -47,17 +49,17 @@ where T: AsyncRead + AsyncWrite,
 {
     /// Establish the HTTP/2.0 connection and get a service to process inbound
     /// requests.
-    Init(Init<T, B::Data, S::Future, S::InitError>),
+    Init(Init<T, SendBuf<<B::Data as IntoBuf>::Buf>, S::Future, S::InitError>),
 
     /// Both the HTTP/2.0 connection and the service are ready.
     Ready {
-        connection: Accept<T, B::Data>,
+        connection: Accept<T, SendBuf<<B::Data as IntoBuf>::Buf>>,
         service: S::Service,
     },
 
     /// The service has closed, so poll until connection is closed.
     GoAway {
-        connection: Accept<T, B::Data>,
+        connection: Accept<T, SendBuf<<B::Data as IntoBuf>::Buf>>,
         error: Error<S>,
     },
 
@@ -84,7 +86,7 @@ enum BackgroundState<T, B>
 where B: Body,
 {
     Respond {
-        respond: SendResponse<B::Data>,
+        respond: SendResponse<SendBuf<<B::Data as IntoBuf>::Buf>>,
         response: T,
     },
     Flush(flush::Flush<B>),
@@ -379,7 +381,9 @@ impl<T, B> Background<T, B>
 where T: Future,
       B: Body,
 {
-    fn new(respond: SendResponse<B::Data>, response: T) -> Self {
+    fn new(respond: SendResponse<SendBuf<<B::Data as IntoBuf>::Buf>>, response: T)
+        -> Self
+    {
         Background {
             state: BackgroundState::Respond {
                 respond,

Diff for: tests/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "tests"
+version = "0.1.0"
+authors = ["Carl Lerche <[email protected]>"]
+publish = false
+
+[dependencies]
+bytes = "0.4.7"
+futures = "0.1.21"
+h2 = "0.1.4"
+h2-support = { git = "https://github.com/carllerche/h2" }
+http = "0.1.5"
+tokio = "0.1.5"
+tower-h2 = { path = ".." }
+tower-service = { git = "https://github.com/tower-rs/tower" }
+tower-util = { git = "https://github.com/tower-rs/tower" }

Diff for: tests/README.md

@@ -0,0 +1,5 @@
+# tower-h2 integration tests
+
+This is a separate crate because it uses a git dependency on h2-support (from
+the h2 crate). When deploying to crates.io, there cannot be any git
+dependencies. This applies to dev dependencies as well.