diff --git a/src/proto/streams/recv.rs b/src/proto/streams/recv.rs index a2c2977..b5f1158 100644 --- a/src/proto/streams/recv.rs +++ b/src/proto/streams/recv.rs @@ -780,6 +780,16 @@ impl Recv { // Frame is trailer stream.pending_recv.push_front(&mut self.buffer, event); + // Notify the recv task. This is done just in case + // `poll_trailers` was called. + // + // It is very likely that `notify_recv` will just be a no-op (as + // the task will be None), so this isn't really much of a + // performance concern. It also means we don't have to track + // state to see if `poll_trailers` was called before `poll_data` + // returned `None`. + stream.notify_recv(); + // No more data frames Ok(None.into()) }, @@ -793,11 +803,11 @@ impl Recv { ) -> Poll, proto::Error> { match stream.pending_recv.pop_front(&mut self.buffer) { Some(Event::Trailers(trailers)) => Ok(Some(trailers).into()), - Some(_) => { - // TODO: This is a user error. `poll_trailers` was called before - // the entire set of data frames have been consumed. What should - // we do? - panic!("poll_trailers called before data has been consumed"); + Some(event) => { + // Frame is not trailers.. not ready to poll trailers yet. + stream.pending_recv.push_front(&mut self.buffer, event); + + Ok(Async::NotReady) }, None => self.schedule_recv(stream), } diff --git a/src/share.rs b/src/share.rs index b85657c..081abcb 100644 --- a/src/share.rs +++ b/src/share.rs @@ -149,8 +149,6 @@ impl RecvStream { } /// Returns received trailers. - /// - /// This function **must** not be called until `Body::poll` returns `None`. pub fn poll_trailers(&mut self) -> Poll, ::Error> { self.inner.inner.poll_trailers().map_err(Into::into) }