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> { 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> { + pub async fn trailers(&mut self) -> Result, 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>> { - self.inner.inner.poll_trailers(cx).map_err_(Into::into) + ) -> Poll, 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 { - 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();