fix(codec): Fix streaming reponses w/ many status (#689)

Closes #681

tests/integration_tests/Cargo.toml

@@ -23,6 +23,7 @@ futures = "0.3"
 tower = { version = "0.4", features = [] }
 http-body = "0.4"
 http = "0.2"
+tracing-subscriber = "0.2"
 
 [build-dependencies]
 tonic-build = { path = "../../tonic-build" }

tests/integration_tests/build.rs

@@ -1,3 +1,4 @@
 fn main() {
     tonic_build::compile_protos("proto/test.proto").unwrap();
+    tonic_build::compile_protos("proto/stream.proto").unwrap();
 }

tests/integration_tests/proto/stream.proto

@@ -0,0 +1,10 @@
+syntax = "proto3";
+
+package stream;
+
+service TestStream {
+  rpc StreamCall(InputStream) returns (stream OutputStream);
+}
+
+message InputStream {}
+message OutputStream {}

tests/integration_tests/src/lib.rs

@@ -1,3 +1,4 @@
 pub mod pb {
     tonic::include_proto!("test");
+    tonic::include_proto!("stream");
 }

tests/integration_tests/tests/status.rs

@@ -1,6 +1,9 @@
 use bytes::Bytes;
 use futures_util::FutureExt;
-use integration_tests::pb::{test_client, test_server, Input, Output};
+use integration_tests::pb::{
+    test_client, test_server, test_stream_client, test_stream_server, Input, InputStream, Output,
+    OutputStream,
+};
 use std::time::Duration;
 use tokio::sync::oneshot;
 use tonic::metadata::{MetadataMap, MetadataValue};
@@ -117,3 +120,61 @@ async fn status_with_metadata() {
 
     jh.await.unwrap();
 }
+
+type Stream<T> = std::pin::Pin<
+    Box<dyn futures::Stream<Item = std::result::Result<T, Status>> + Send + Sync + 'static>,
+>;
+
+#[tokio::test]
+async fn status_from_server_stream() {
+    trace_init();
+
+    struct Svc;
+
+    #[tonic::async_trait]
+    impl test_stream_server::TestStream for Svc {
+        type StreamCallStream = Stream<OutputStream>;
+
+        async fn stream_call(
+            &self,
+            _: Request<InputStream>,
+        ) -> Result<Response<Self::StreamCallStream>, Status> {
+            let s = futures::stream::iter(vec![
+                Err::<OutputStream, _>(Status::unavailable("foo")),
+                Err::<OutputStream, _>(Status::unavailable("bar")),
+            ]);
+            Ok(Response::new(Box::pin(s) as Self::StreamCallStream))
+        }
+    }
+
+    let svc = test_stream_server::TestStreamServer::new(Svc);
+
+    tokio::spawn(async move {
+        Server::builder()
+            .add_service(svc)
+            .serve("127.0.0.1:1339".parse().unwrap())
+            .await
+            .unwrap();
+    });
+
+    tokio::time::sleep(Duration::from_millis(100)).await;
+
+    let mut client = test_stream_client::TestStreamClient::connect("http://127.0.0.1:1339")
+        .await
+        .unwrap();
+
+    let mut stream = client
+        .stream_call(InputStream {})
+        .await
+        .unwrap()
+        .into_inner();
+
+    assert_eq!(stream.message().await.unwrap_err().message(), "foo");
+    assert_eq!(stream.message().await.unwrap(), None);
+}
+
+fn trace_init() {
+    let _ = tracing_subscriber::FmtSubscriber::builder()
+        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
+        .try_init();
+}

tonic/src/codec/decode.rs

@@ -258,7 +258,11 @@ impl<T> Stream for Streaming<T> {
             match ready!(Pin::new(&mut self.body).poll_trailers(cx)) {
                 Ok(trailer) => {
                     if let Err(e) = crate::status::infer_grpc_status(trailer.as_ref(), status) {
-                        return Some(Err(e)).into();
+                        if let Some(e) = e {
+                            return Some(Err(e)).into();
+                        } else {
+                            return Poll::Ready(None);
+                        }
                     } else {
                         self.trailers = trailer.map(MetadataMap::from_headers);
                     }

tonic/src/codec/encode.rs

@@ -88,6 +88,7 @@ pub(crate) struct EncodeBody<S> {
     inner: S,
     error: Option<Status>,
     role: Role,
+    is_end_stream: bool,
 }
 
 impl<S> EncodeBody<S>
@@ -99,6 +100,7 @@ where
             inner,
             error: None,
             role: Role::Client,
+            is_end_stream: false,
         }
     }
 
@@ -107,6 +109,7 @@ where
             inner,
             error: None,
             role: Role::Server,
+            is_end_stream: false,
         }
     }
 }
@@ -119,7 +122,7 @@ where
     type Error = Status;
 
     fn is_end_stream(&self) -> bool {
-        false
+        self.is_end_stream
     }
 
     fn poll_data(
@@ -148,7 +151,13 @@ where
             Role::Client => Poll::Ready(Ok(None)),
             Role::Server => {
                 let self_proj = self.project();
+
+                if *self_proj.is_end_stream {
+                    return Poll::Ready(Ok(None));
+                }
+
                 let status = if let Some(status) = self_proj.error.take() {
+                    *self_proj.is_end_stream = true;
                     status
                 } else {
                     Status::new(Code::Ok, "")

tonic/src/status.rs

@@ -657,13 +657,13 @@ impl Error for Status {
 pub(crate) fn infer_grpc_status(
     trailers: Option<&HeaderMap>,
     status_code: http::StatusCode,
-) -> Result<(), Status> {
+) -> Result<(), Option<Status>> {
     if let Some(trailers) = trailers {
         if let Some(status) = Status::from_header_map(&trailers) {
             if status.code() == Code::Ok {
                 return Ok(());
             } else {
-                return Err(status);
+                return Err(status.into());
             }
         }
     }
@@ -678,6 +678,13 @@ pub(crate) fn infer_grpc_status(
         | http::StatusCode::BAD_GATEWAY
         | http::StatusCode::SERVICE_UNAVAILABLE
         | http::StatusCode::GATEWAY_TIMEOUT => Code::Unavailable,
+        // We got a 200 but no trailers, we can infer that this request is finished.
+        //
+        // This can happen when a streaming response sends two Status but
+        // gRPC requires that we end the stream after the first status.
+        //
+        // https://github.com/hyperium/tonic/issues/681
+        http::StatusCode::OK => return Err(None),
         _ => Code::Unknown,
     };
 
@@ -686,7 +693,7 @@ pub(crate) fn infer_grpc_status(
         status_code.as_u16(),
     );
     let status = Status::new(code, msg);
-    Err(status)
+    Err(status.into())
 }
 
 // ===== impl Code =====