diff --git a/tonic-interop/src/server.rs b/tonic-interop/src/server.rs index e9e3f78..1858b8c 100644 --- a/tonic-interop/src/server.rs +++ b/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 { 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; diff --git a/tonic-interop/test.sh b/tonic-interop/test.sh index df47e3b..1f8ffd1 100755 --- a/tonic-interop/test.sh +++ b/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 diff --git a/tonic/src/body.rs b/tonic/src/body.rs index d24c6b2..f793efb 100644 --- a/tonic/src/body.rs +++ b/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>> { + Poll::Ready(None) + } + + fn poll_trailers( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>> { + Poll::Ready(Ok(None)) + } +} diff --git a/tonic/src/codec/decode.rs b/tonic/src/codec/decode.rs index 87f1784..2ec7c53 100644 --- a/tonic/src/codec/decode.rs +++ b/tonic/src/codec/decode.rs @@ -91,8 +91,12 @@ impl Streaming { impl Streaming { /// Fetch the next message from this stream. - pub async fn message(&mut self) -> Option> { - future::poll_fn(|cx| Pin::new(&mut *self).poll_next(cx)).await + pub async fn message(&mut self) -> Result, 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 Streaming { } // 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. diff --git a/tonic/src/server/grpc.rs b/tonic/src/server/grpc.rs index d53a86c..a34c0b4 100644 --- a/tonic/src/server/grpc.rs +++ b/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()) } } }