feat(grpc): Remove trailers from SendStream and make Handle return Trailers .

Author: sauravzg

grpc/examples/inmemory.rs

@@ -37,6 +37,7 @@ use grpc::core::RecvMessage;
 use grpc::core::RequestHeaders;
 use grpc::core::SendMessage;
 use grpc::core::ServerResponseStreamItem;
+use grpc::core::Trailers;
 use grpc::credentials::InsecureChannelCredentials;
 use grpc::inmemory;
 use grpc::server;
@@ -84,7 +85,7 @@ impl Handle for Handler {
         _options: CallOptions,
         tx: &mut impl server::SendStream,
         mut rx: impl server::RecvStream + 'static,
-    ) {
+    ) -> Trailers {
         let method = headers.method_name().clone();
         let id = self.id.clone();
         // Send headers
@@ -108,16 +109,8 @@ impl Handle for Handler {
                 )
                 .await;
         }
-        // Send trailers
-        let _ = tx
-            .send(
-                ServerResponseStreamItem::Trailers(grpc::core::Trailers::new(grpc::Status::new(
-                    grpc::StatusCode::Ok,
-                    "OK",
-                ))),
-                server::SendOptions::default(),
-            )
-            .await;
+        // Return trailers
+        Trailers::new(grpc::Status::new(grpc::StatusCode::Ok, "OK"))
     }
 }
 

grpc/src/core/mod.rs

@@ -101,8 +101,7 @@ impl dyn RecvMessage + '_ {
     }
 }
 
-/// ResponseStreamItem represents an item in a response stream (either server
-/// sending or client receiving).
+/// ResponseStreamItem represents an item in a response stream from the client's view.
 ///
 /// A response stream must always contain items exactly as follows:
 ///
@@ -135,8 +134,6 @@ pub enum ServerResponseStreamItem<'a> {
     Headers(ResponseHeaders),
     /// Indicates a message on the stream.
     Message(&'a dyn SendMessage),
-    /// Indicates trailers were received on the stream and includes the trailers.
-    Trailers(Trailers),
 }
 
 /// Contains all information transmitted in the response headers of an RPC.

grpc/src/inmemory/mod.rs

@@ -85,6 +85,7 @@ struct InMemoryServerCall {
     headers: RequestHeaders,
     req_rx: mpsc::UnboundedReceiver<InMemoryRequestStreamItem>,
     resp_tx: mpsc::UnboundedSender<InMemoryResponseStreamItem>,
+    trailer_tx: oneshot::Sender<Trailers>,
 }
 
 enum InMemoryRequestStreamItem {
@@ -95,7 +96,6 @@ enum InMemoryRequestStreamItem {
 enum InMemoryResponseStreamItem {
     Headers(ResponseHeaders),
     Message(Box<dyn Buf + Send + Sync>),
-    Trailers(Trailers),
     StreamClosed,
 }
 
@@ -179,6 +179,7 @@ impl ServerListener for InMemoryListener {
                     headers: call.headers,
                     send: InMemoryServerSendStream { tx: call.resp_tx },
                     recv: InMemoryServerRecvStream { rx: call.req_rx },
+                    trailers_tx: call.trailer_tx,
                 })
             }
             _ = self.inner.close_notify.notified() => {
@@ -204,7 +205,6 @@ impl ServerSendStream for InMemoryServerSendStream {
                 let buf = m.encode().map_err(|_| ())?;
                 InMemoryResponseStreamItem::Message(buf)
             }
-            ServerResponseStreamItem::Trailers(t) => InMemoryResponseStreamItem::Trailers(t),
         };
 
         self.tx.send(inmemory_item).map_err(|_| ())
@@ -246,18 +246,23 @@ impl Invoke for InMemoryConnection {
     ) -> (Self::SendStream, Self::RecvStream) {
         let (req_tx, req_rx) = mpsc::unbounded_channel::<InMemoryRequestStreamItem>();
         let (resp_tx, resp_rx) = mpsc::unbounded_channel::<InMemoryResponseStreamItem>();
+        let (trailer_tx, trailer_rx) = oneshot::channel();
 
         let call = InMemoryServerCall {
             headers,
             req_rx,
             resp_tx,
+            trailer_tx,
         };
 
-        let _ = self.s.try_send(call);
+        let _ = self.s.send(call).await;
 
         (
             Box::new(InMemoryClientSendStream { tx: Some(req_tx) }),
-            Box::new(InMemoryClientRecvStream { rx: resp_rx }),
+            Box::new(InMemoryClientRecvStream {
+                rx: resp_rx,
+                trailer_rx: Some(trailer_rx),
+            }),
         )
     }
 }
@@ -299,6 +304,7 @@ impl Drop for InMemoryClientSendStream {
 
 pub struct InMemoryClientRecvStream {
     rx: mpsc::UnboundedReceiver<InMemoryResponseStreamItem>,
+    trailer_rx: Option<oneshot::Receiver<Trailers>>,
 }
 
 impl ClientRecvStream for InMemoryClientRecvStream {
@@ -309,8 +315,14 @@ impl ClientRecvStream for InMemoryClientRecvStream {
                 msg.decode(&mut buf).unwrap();
                 ClientResponseStreamItem::Message(())
             }
-            Some(InMemoryResponseStreamItem::Trailers(t)) => ClientResponseStreamItem::Trailers(t),
-            _ => ClientResponseStreamItem::StreamClosed,
+            _ => {
+                if let Some(trailer_rx) = self.trailer_rx.take() {
+                    if let Ok(trailers) = trailer_rx.await {
+                        return ClientResponseStreamItem::Trailers(trailers);
+                    }
+                }
+                ClientResponseStreamItem::StreamClosed
+            }
         }
     }
 }

grpc/src/server/mod.rs

@@ -29,6 +29,8 @@ use crate::client::CallOptions;
 use crate::core::RecvMessage;
 use crate::core::RequestHeaders;
 use crate::core::ServerResponseStreamItem;
+use crate::core::Trailers;
+use tokio::sync::oneshot;
 
 pub struct Server {
     handler: Option<Arc<dyn DynHandle>>,
@@ -38,6 +40,7 @@ pub struct Call<SS, RS> {
     pub headers: RequestHeaders,
     pub send: SS,
     pub recv: RS,
+    pub trailers_tx: oneshot::Sender<Trailers>,
 }
 
 #[trait_variant::make(Send)]
@@ -64,11 +67,14 @@ impl Server {
             let mut send: Box<dyn DynSendStream> = Box::new(call.send);
             let recv = BoxedRecvStream(Box::new(call.recv));
             let options = CallOptions::default();
-            self.handler
+            let trailers_tx = call.trailers_tx;
+            let trailers = self
+                .handler
                 .as_ref()
                 .unwrap()
                 .dyn_handle(call.headers, options, &mut *send, recv)
                 .await;
+            let _ = trailers_tx.send(trailers);
         }
     }
 }
@@ -92,7 +98,7 @@ pub trait Handle: Send + Sync {
         options: CallOptions,
         tx: &mut impl SendStream,
         rx: impl RecvStream + 'static,
-    );
+    ) -> Trailers;
 }
 
 #[async_trait]
@@ -103,7 +109,7 @@ trait DynHandle: Send + Sync {
         options: CallOptions,
         tx: &mut dyn DynSendStream,
         rx: BoxedRecvStream,
-    );
+    ) -> Trailers;
 }
 
 #[async_trait]
@@ -114,7 +120,7 @@ impl<T: Handle> DynHandle for T {
         options: CallOptions,
         mut tx: &mut dyn DynSendStream,
         rx: BoxedRecvStream,
-    ) {
+    ) -> Trailers {
         self.handle(headers, options, &mut tx, rx).await
     }
 }
@@ -146,10 +152,9 @@ impl RecvStream for BoxedRecvStream {
 pub trait SendStream {
     /// Sends the next item on the stream. Returns `Ok(())` on success, or
     /// `Err(())` on failure. `Err(())` is a terminal state.
-    /// Sending is also considered complete after successfully sending
-    /// `ServerResponseStreamItem::Trailers`.
-    /// Calling this method after an error or after sending trailers should be
-    /// avoided and is unspecified.
+    /// Sending is considered complete when the `handle` method returns
+    /// `Trailers`.
+    /// Calling this method after an error should be avoided and is unspecified.
     ///
     /// # Cancel safety
     ///