fix(http2): send trailers if Payload includes them

src/proto/h2/mod.rs

@@ -30,6 +30,7 @@ where
     S: Payload,
 {
     body_tx: SendStream<SendBuf<S::Data>>,
+    data_done: bool,
     stream: S,
 }
 
@@ -40,9 +41,23 @@ where
     fn new(stream: S, tx: SendStream<SendBuf<S::Data>>) -> PipeToSendStream<S> {
         PipeToSendStream {
             body_tx: tx,
+            data_done: false,
             stream: stream,
         }
     }
+
+    fn on_err(&mut self, err: S::Error) -> ::Error {
+        let err = ::Error::new_user_body(err);
+        trace!("send body user stream error: {}", err);
+        self.body_tx.send_reset(Reason::INTERNAL_ERROR);
+        err
+    }
+
+    fn send_eos_frame(&mut self) -> ::Result<()> {
+        trace!("send body eos");
+        self.body_tx.send_data(SendBuf(None), true)
+            .map_err(::Error::new_body_write)
+    }
 }
 
 impl<S> Future for PipeToSendStream<S>
@@ -54,6 +69,7 @@ where
 
     fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
         loop {
+            if !self.data_done {
                 // TODO: make use of flow control on SendStream
                 // If you're looking at this and thinking of trying to fix this TODO,
                 // you may want to look at:
@@ -68,8 +84,8 @@ where
                 // - else:
                 // -   try reserve a smallish amount of capacity
                 // -   call self.body_tx.poll_capacity(), return if NotReady
-            match self.stream.poll_data() {
-                Ok(Async::Ready(Some(chunk))) => {
+                match try_ready!(self.stream.poll_data().map_err(|e| self.on_err(e))) {
+                    Some(chunk) => {
                         let is_eos = self.stream.is_end_stream();
                         trace!(
                             "send body chunk: {} bytes, eos={}",
@@ -85,18 +101,27 @@ where
                             return Ok(Async::Ready(()))
                         }
                     },
-                Ok(Async::Ready(None)) => {
-                    trace!("send body eos");
-                    self.body_tx.send_data(SendBuf(None), true)
+                    None => {
+                        let is_eos = self.stream.is_end_stream();
+                        if is_eos {
+                            return self.send_eos_frame().map(Async::Ready);
+                        } else {
+                            self.data_done = true;
+                            // loop again to poll_trailers
+                        }
+                    },
+                }
+            } else {
+                match try_ready!(self.stream.poll_trailers().map_err(|e| self.on_err(e))) {
+                    Some(trailers) => {
+                        self.body_tx.send_trailers(trailers)
                             .map_err(::Error::new_body_write)?;
                         return Ok(Async::Ready(()));
                     },
-                Ok(Async::NotReady) => return Ok(Async::NotReady),
-                Err(err) => {
-                    let err = ::Error::new_user_body(err);
-                    trace!("send body user stream error: {}", err);
-                    self.body_tx.send_reset(Reason::INTERNAL_ERROR);
-                    return Err(err);
+                    None => {
+                        // There were no trailers, so send an empty DATA frame...
+                        return self.send_eos_frame().map(Async::Ready);
+                    },
                 }
             }
         }