Fix server streaming inital error response headers

tonic-interop/src/server.rs

@@ -1,6 +1,6 @@
 use crate::pb::{self, *};
 use async_stream::try_stream;
-use futures_util::TryStreamExt;
+use futures_util::{TryStreamExt, StreamExt, stream};
 use std::pin::Pin;
 use std::time::{Duration, Instant};
 use tonic::{Code, Request, Response, Status};
@@ -109,24 +109,37 @@ impl pb::server::TestService for TestService {
     ) -> Result<Self::FullDuplexCallStream> {
         let mut stream = req.into_inner();
 
-        let stream = try_stream! {
-            while let Some(msg) = stream.try_next().await? {
-                if let Some(echo_status) = msg.response_status {
-                    let status = Status::new(Code::from_i32(echo_status.code), echo_status.message);
-                     Err(status)?;
-                }
-
-                for param in msg.response_parameters {
-                    let deadline = Instant::now() + Duration::from_micros(param.interval_us as u64);
-                    tokio::timer::delay(deadline).await;
-
-                    let payload = crate::server_payload(param.size as usize);
-                    yield StreamingOutputCallResponse { payload: Some(payload) };
-                }
+        if let Some(first_msg) = stream.message().await? {
+            if let Some(echo_status) = first_msg.response_status {
+                let status = Status::new(Code::from_i32(echo_status.code), echo_status.message);
+                return Err(status);
             }
-        };
 
-        Ok(Response::new(Box::pin(stream) as Self::FullDuplexCallStream))
+            let single_message = stream::iter(vec![Ok(first_msg)]);
+            let mut stream = single_message.chain(stream);
+
+            let stream = try_stream! {
+                while let Some(msg) = stream.try_next().await? {
+                    if let Some(echo_status) = msg.response_status {
+                        let status = Status::new(Code::from_i32(echo_status.code), echo_status.message);
+                        Err(status)?;
+                    }
+
+                    for param in msg.response_parameters {
+                        let deadline = Instant::now() + Duration::from_micros(param.interval_us as u64);
+                        tokio::timer::delay(deadline).await;
+
+                        let payload = crate::server_payload(param.size as usize);
+                        yield StreamingOutputCallResponse { payload: Some(payload) };
+                    }
+                }
+            };
+
+            Ok(Response::new(Box::pin(stream) as Self::FullDuplexCallStream))
+        } else {
+            let stream = stream::empty();
+            Ok(Response::new(Box::pin(stream) as Self::FullDuplexCallStream))
+        }
     }
 
     type HalfDuplexCallStream = Stream<StreamingOutputCallResponse>;

tonic-interop/test.sh

@@ -43,6 +43,6 @@ trap 'echo ":; killing test server"; kill ${SERVER_PID};' EXIT
 
 cargo run -p tonic-interop --bin client -- \
 --test_case=empty_unary,large_unary,client_streaming,server_streaming,ping_pong,\
-empty_stream
-# status_code_and_message,special_status_message,unimplemented_method,\
+empty_stream,status_code_and_message
+# special_status_message,unimplemented_method,\
 # unimplemented_service,custom_metadata $ARG

tonic/src/body.rs

@@ -110,6 +110,13 @@ impl BoxBody {
             inner: Box::pin(MapBody(inner)),
         }
     }
+
+    /// Create a new `BoxBody` that is empty.
+    pub fn empty() -> Self {
+        BoxBody {
+            inner: Box::pin(EmptyBody::default()),
+        }
+    }
 }
 
 impl HttpBody for BoxBody {
@@ -185,3 +192,31 @@ impl fmt::Debug for BoxBody {
         f.debug_struct("BoxBody").finish()
     }
 }
+
+#[derive(Debug, Default)]
+struct EmptyBody {
+    _p: (),
+}
+
+impl HttpBody for EmptyBody {
+    type Data = BytesBuf;
+    type Error = Status;
+
+    fn is_end_stream(&self) -> bool {
+        true
+    }
+
+    fn poll_data(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<Self::Data, Self::Error>>> {
+        Poll::Ready(None)
+    }
+
+    fn poll_trailers(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<Option<http::HeaderMap>, Self::Error>> {
+        Poll::Ready(Ok(None))
+    }
+}

tonic/src/codec/decode.rs

@@ -91,8 +91,12 @@ impl<T> Streaming<T> {
 
 impl<T> Streaming<T> {
     /// Fetch the next message from this stream.
-    pub async fn message(&mut self) -> Option<Result<T, Status>> {
-        future::poll_fn(|cx| Pin::new(&mut *self).poll_next(cx)).await
+    pub async fn message(&mut self) -> Result<Option<T>, Status> {
+        match future::poll_fn(|cx| Pin::new(&mut *self).poll_next(cx)).await {
+            Some(Ok(m)) => Ok(Some(m)),
+            Some(Err(e)) => Err(e),
+            None => Ok(None),
+        }
     }
 
     /// Fetch the trailing metadata.
@@ -108,9 +112,7 @@ impl<T> Streaming<T> {
         }
 
         // To fetch the trailers we must clear the body and drop it.
-        while let Some(res) = self.message().await {
-            res?;
-        }
+        while let Some(res) = self.message().await? {}
 
         // Since we call poll_trailers internally on poll_next we need to
         // check if it got cached again.

tonic/src/server/grpc.rs

@@ -186,8 +186,6 @@ where
                 http::Response::from_parts(parts, BoxBody::new(body))
             }
             Err(status) => {
-                let status = stream::once(future::err(status));
-                let body = encode_server(self.codec.encoder(), status);
                 let (mut parts, _body) = Response::new(()).into_http().into_parts();
 
                 parts.headers.insert(
@@ -195,7 +193,9 @@ where
                     http::header::HeaderValue::from_static(T::CONTENT_TYPE),
                 );
 
-                http::Response::from_parts(parts, BoxBody::new(body))
+                status.add_header(&mut parts.headers).unwrap();
+
+                http::Response::from_parts(parts, BoxBody::empty())
             }
         }
     }