Add trailing metadata for unary/client-streaming

tonic-interop/src/client.rs

@@ -402,22 +402,21 @@ pub async fn custom_metadata(client: &mut Client, assertions: &mut Vec<TestAsser
     req_stream.metadata_mut().insert(key1, value1.clone());
     req_stream.metadata_mut().insert_bin(key2, value2.clone());
 
-    // let response = client
+    let response = client
-    //     .unary_call(req_unary)
+        .unary_call(req_unary)
-    //     .await
+        .await
-    //     .expect("call should pass.");
+        .expect("call should pass.");
-
-    // assertions.push(test_assert!(
-    //     "metadata string must match in unary",
-    //     response.metadata().get(key1) == Some(&value1),
-    //     format!("result={:?}", response.metadata().get(key1))
-    // ));
-    // assertions.push(test_assert!(
-    //     "metadata bin must match in unary",
-    //     response.metadata().get_bin(key2) == Some(&value2),
-    //     format!("result={:?}", response.metadata().get_bin(key1))
-    // ));
 
+    assertions.push(test_assert!(
+        "metadata string must match in unary",
+        response.metadata().get(key1) == Some(&value1),
+        format!("result={:?}", response.metadata().get(key1))
+    ));
+    assertions.push(test_assert!(
+        "metadata bin must match in unary",
+        response.metadata().get_bin(key2) == Some(&value2),
+        format!("result={:?}", response.metadata().get_bin(key1))
+    ));
 
     let response = client
         .full_duplex_call(req_stream)
@@ -432,8 +431,6 @@ pub async fn custom_metadata(client: &mut Client, assertions: &mut Vec<TestAsser
 
     let mut stream = response.into_inner();
 
-    // while let Some(_) = stream.next().await {}
-
     let trailers = stream.trailers().await.unwrap().unwrap();
 
     assertions.push(test_assert!(

tonic/src/client/grpc.rs

@@ -77,7 +77,7 @@ impl<T> Grpc<T> {
         M1: Send,
         M2: Send + Unpin + 'static,
     {
-        let (parts, body) = self.streaming(request, path, codec).await?.into_parts();
+        let (mut parts, body) = self.streaming(request, path, codec).await?.into_parts();
 
         futures_util::pin_mut!(body);
 
@@ -86,6 +86,10 @@ impl<T> Grpc<T> {
             .await?
             .ok_or(Status::new(Code::Internal, "Missing response message."))?;
 
+        if let Some(trailers) = body.trailers().await? {
+            parts.merge(trailers);
+        }
+
         Ok(Response::from_parts(parts, message))
     }
 

tonic/src/codec/decode.rs

@@ -91,10 +91,10 @@ impl<T> Streaming<T> {
         future::poll_fn(|cx| Pin::new(&mut *self).poll_next(cx)).await
     }
 
-    pub async fn trailers(mut self) -> Result<Option<MetadataMap>, Status> {
+    pub async fn trailers(&mut self) -> Result<Option<MetadataMap>, Status> {
         // Shortcut to see if we already pulled the trailers in the stream step
         // we need to do that so that the stream can error on trailing grpc-status
-        if let Some(trailers) = self.trailers {
+        if let Some(trailers) = self.trailers.take() {
             return Ok(Some(trailers));
         }
 
@@ -105,11 +105,10 @@ impl<T> Streaming<T> {
 
         // Since we call poll_trailers internally on poll_next we need to
         // check if it got cached again.
-        if let Some(trailers) = self.trailers {
+        if let Some(trailers) = self.trailers.take() {
             return Ok(Some(trailers));
         }
 
-
         // Trailers were not caught during poll_next and thus lets poll for
         // them manually.
         let map =

tonic/src/metadata/map.rs

@@ -1181,6 +1181,10 @@ impl MetadataMap {
     {
         key.remove(self)
     }
+
+    pub(crate) fn merge(&mut self, other: MetadataMap) {
+        self.headers.extend(other.headers);
+    }
 }
 
 // ===== impl Iter =====

tonic/src/server/grpc.rs

@@ -136,7 +136,14 @@ where
             .await?
             .ok_or(Status::new(Code::Internal, "Missing request message."))?;
 
-        Ok(Request::from_http_parts(parts, message))
+
+        let mut req = Request::from_http_parts(parts, message);
+
+        if let Some(trailers) = stream.trailers().await? {
+            req.metadata_mut().merge(trailers);
+        }
+
+        Ok(req)
     }
 
     fn map_request_streaming<B>(
@@ -170,8 +177,6 @@ where
 
                 let body = encode_server(self.codec.encoder(), body.into_stream());
 
-                // FIXME: try to return impl Trait?
-                // let body = Box::pin(body) as BoxStream<BytesBuf>;
                 http::Response::from_parts(parts, BoxBody::new(body))
             }
             Err(status) => {