From b2af3d4b281d35468514c8dd9a28bc13aaaeb3c6 Mon Sep 17 00:00:00 2001
From: Sean McArthur <sean@seanmonstar.com>
Date: Tue, 20 Aug 2019 16:20:49 -0700
Subject: [PATCH] Change trailers() to return Result<Option> instead of
 Option<Result>

---
 examples/client.rs               |  4 ++--
 src/share.rs                     | 10 +++++---
 tests/h2-tests/tests/hammer.rs   | 41 +++++++-------------------------
 tests/h2-tests/tests/trailers.rs |  2 +-
 4 files changed, 18 insertions(+), 39 deletions(-)

diff --git a/examples/client.rs b/examples/client.rs
index 9df87f2..3bbd163 100644
--- a/examples/client.rs
+++ b/examples/client.rs
@@ -46,8 +46,8 @@ pub async fn main() -> Result<(), Box<dyn Error>> {
         println!("GOT CHUNK = {:?}", chunk?);
     }
 
-    if let Some(trailers) = body.trailers().await {
-        println!("GOT TRAILERS: {:?}", trailers?);
+    if let Some(trailers) = body.trailers().await? {
+        println!("GOT TRAILERS: {:?}", trailers);
     }
 
     Ok(())
diff --git a/src/share.rs b/src/share.rs
index 5ac723f..509e9e0 100644
--- a/src/share.rs
+++ b/src/share.rs
@@ -422,7 +422,7 @@ impl RecvStream {
     }
 
     /// Get optional trailers for this stream.
-    pub async fn trailers(&mut self) -> Option<Result<HeaderMap, crate::Error>> {
+    pub async fn trailers(&mut self) -> Result<Option<HeaderMap>, crate::Error> {
         futures::future::poll_fn(move |cx| self.poll_trailers(cx)).await
     }
 
@@ -435,8 +435,12 @@ impl RecvStream {
     pub fn poll_trailers(
         &mut self,
         cx: &mut Context,
-    ) -> Poll<Option<Result<HeaderMap, crate::Error>>> {
-        self.inner.inner.poll_trailers(cx).map_err_(Into::into)
+    ) -> Poll<Result<Option<HeaderMap>, crate::Error>> {
+        match ready!(self.inner.inner.poll_trailers(cx)) {
+            Some(Ok(map)) => Poll::Ready(Ok(Some(map))),
+            Some(Err(e)) => Poll::Ready(Err(e.into())),
+            None => Poll::Ready(Ok(None)),
+        }
     }
 
     /// Returns the stream ID of this stream.
diff --git a/tests/h2-tests/tests/hammer.rs b/tests/h2-tests/tests/hammer.rs
index d3a6052..73045d6 100644
--- a/tests/h2-tests/tests/hammer.rs
+++ b/tests/h2-tests/tests/hammer.rs
@@ -1,10 +1,7 @@
 #![feature(async_await)]
 
-use futures::{ready, FutureExt, StreamExt, TryFutureExt};
+use futures::{FutureExt, StreamExt, TryFutureExt};
 use h2_support::prelude::*;
-use std::future::Future;
-use std::pin::Pin;
-use std::task::{Context, Poll};
 
 use std::io;
 use std::{
@@ -89,31 +86,6 @@ where
     Ok(())
 }
 
-struct Process {
-    body: RecvStream,
-    trailers: bool,
-}
-
-impl Future for Process {
-    type Output = Result<(), h2::Error>;
-
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        loop {
-            if self.trailers {
-                ready!(self.body.poll_trailers(cx));
-                return Poll::Ready(Ok(()));
-            } else {
-                match ready!(Pin::new(&mut self.body).poll_next(cx)) {
-                    None => {
-                        self.trailers = true;
-                    }
-                    _ => {}
-                }
-            }
-        }
-    }
-}
-
 #[test]
 fn hammer_client_concurrency() {
     // This reproduces issue #326.
@@ -150,11 +122,14 @@ fn hammer_client_concurrency() {
 
                 response
                     .and_then(|response| {
-                        let (_, body) = response.into_parts();
+                        let mut body = response.into_body();
 
-                        Process {
-                            body,
-                            trailers: false,
+                        async move {
+                            while let Some(res) = body.data().await {
+                                res?;
+                            }
+                            body.trailers().await?;
+                            Ok(())
                         }
                     })
                     .map_err(|e| {
diff --git a/tests/h2-tests/tests/trailers.rs b/tests/h2-tests/tests/trailers.rs
index 4793534..f65a5e9 100644
--- a/tests/h2-tests/tests/trailers.rs
+++ b/tests/h2-tests/tests/trailers.rs
@@ -100,7 +100,7 @@ async fn send_trailers_immediately() {
     let chunk = h2.run(body.next()).await;
     assert!(chunk.is_none());
 
-    let trailers = h2.run(poll_fn(|cx| body.poll_trailers(cx))).await;
+    let trailers = h2.run(poll_fn(|cx| body.poll_trailers(cx))).await.unwrap();
     assert!(trailers.is_none());
 
     h2.await.unwrap();