Merge pull request #208 from Tim-Zhang/fix-stream-recv-default-data

stream: Fix default data won't be received issue

Author: lifupan

example/async-stream-client.rs

@@ -41,9 +41,12 @@ async fn main() {
     let sc1 = sc.clone();
     let t5 = tokio::spawn(echo_null(sc1));
 
-    let t6 = tokio::spawn(echo_null_stream(sc));
+    let sc1 = sc.clone();
+    let t6 = tokio::spawn(echo_null_stream(sc1));
+
+    let t7 = tokio::spawn(echo_default_value(sc));
 
-    let _ = tokio::join!(t1, t2, t3, t4, t5, t6);
+    let _ = tokio::join!(t1, t2, t3, t4, t5, t6, t7);
 }
 
 fn default_ctx() -> Context {
@@ -185,3 +188,16 @@ async fn echo_null_stream(cli: streaming_ttrpc::StreamingClient) {
         .unwrap()
         .unwrap();
 }
+
+#[cfg(unix)]
+async fn echo_default_value(cli: streaming_ttrpc::StreamingClient) {
+    let mut stream = cli
+        .echo_default_value(default_ctx(), &Default::default()) // send default value to verify #208
+        .await
+        .unwrap();
+
+    let received = stream.recv().await.unwrap().unwrap();
+
+    assert_eq!(received.seq, 0);
+    assert_eq!(received.msg, "");
+}

example/async-stream-server.rs

@@ -13,7 +13,7 @@ use log::{info, LevelFilter};
 #[cfg(unix)]
 use protocols::r#async::{empty, streaming, streaming_ttrpc};
 #[cfg(unix)]
-use ttrpc::asynchronous::Server;
+use ttrpc::{asynchronous::Server, Error};
 
 #[cfg(unix)]
 use async_trait::async_trait;
@@ -134,6 +134,24 @@ impl streaming_ttrpc::Streaming for StreamingService {
         }
         Ok(())
     }
+
+    // It verifies PR #208
+    async fn echo_default_value(
+        &self,
+        _ctx: &::ttrpc::r#async::TtrpcContext,
+        e: streaming::EchoPayload,
+        s: ::ttrpc::r#async::ServerStreamSender<streaming::EchoPayload>,
+    ) -> ::ttrpc::Result<()> {
+        if e.seq != 0 || !e.msg.is_empty() {
+            return Err(Error::Others(
+                "Expect a request with empty payload to verify #208".to_string(),
+            ));
+        }
+
+        s.send(&e).await.unwrap();
+
+        Ok(())
+    }
 }
 
 #[cfg(windows)]

example/protocols/protos/streaming.proto

@@ -32,6 +32,7 @@ service Streaming {
 	rpc DivideStream(Sum) returns (stream Part);
 	rpc EchoNull(stream EchoPayload) returns (google.protobuf.Empty);
 	rpc EchoNullStream(stream EchoPayload) returns (stream google.protobuf.Empty);
+	rpc EchoDefaultValue(EchoPayload) returns (stream EchoPayload);
 }
 
 message EchoPayload {

src/asynchronous/client.rs

@@ -15,10 +15,10 @@ use nix::unistd::close;
 use tokio::{self, sync::mpsc, task};
 
 use crate::common::client_connect;
-use crate::error::{Error, Result};
+use crate::error::{get_rpc_status, Error, Result};
 use crate::proto::{
-    Code, Codec, GenMessage, Message, MessageHeader, Request, Response, FLAG_REMOTE_CLOSED,
-    FLAG_REMOTE_OPEN, MESSAGE_TYPE_DATA, MESSAGE_TYPE_RESPONSE,
+    Code, Codec, GenMessage, Message, MessageHeader, Request, Response, FLAG_NO_DATA,
+    FLAG_REMOTE_CLOSED, FLAG_REMOTE_OPEN, MESSAGE_TYPE_DATA, MESSAGE_TYPE_RESPONSE,
 };
 use crate::r#async::connection::*;
 use crate::r#async::shutdown;
@@ -117,13 +117,20 @@ impl Client {
         streaming_server: bool,
     ) -> Result<StreamInner> {
         let stream_id = self.next_stream_id.fetch_add(2, Ordering::Relaxed);
+        let is_req_payload_empty = req.payload.is_empty();
 
         let mut msg: GenMessage = Message::new_request(stream_id, req)?
             .try_into()
             .map_err(|e: protobuf::Error| Error::Others(e.to_string()))?;
 
         if streaming_client {
-            msg.header.add_flags(FLAG_REMOTE_OPEN);
+            if !is_req_payload_empty {
+                return Err(get_rpc_status(
+                    Code::INVALID_ARGUMENT,
+                    "Creating a ClientStream and sending payload at the same time is not allowed",
+                ));
+            }
+            msg.header.add_flags(FLAG_REMOTE_OPEN | FLAG_NO_DATA);
         } else {
             msg.header.add_flags(FLAG_REMOTE_CLOSED);
         }

src/asynchronous/server.rs

@@ -37,7 +37,7 @@ use crate::context;
 use crate::error::{get_status, Error, Result};
 use crate::proto::{
     check_oversize, Code, Codec, GenMessage, Message, MessageHeader, Request, Response, Status,
-    FLAG_NO_DATA, FLAG_REMOTE_CLOSED, FLAG_REMOTE_OPEN, MESSAGE_TYPE_DATA, MESSAGE_TYPE_REQUEST,
+    FLAG_NO_DATA, FLAG_REMOTE_CLOSED, MESSAGE_TYPE_DATA, MESSAGE_TYPE_REQUEST,
 };
 use crate::r#async::connection::*;
 use crate::r#async::shutdown;
@@ -606,8 +606,8 @@ impl HandlerContext {
         let stream_tx = tx.clone();
         self.streams.lock().unwrap().insert(stream_id, tx);
 
-        let _remote_close = (req_msg.header.flags & FLAG_REMOTE_CLOSED) == FLAG_REMOTE_CLOSED;
-        let _remote_open = (req_msg.header.flags & FLAG_REMOTE_OPEN) == FLAG_REMOTE_OPEN;
+        let no_data = (req_msg.header.flags & FLAG_NO_DATA) == FLAG_NO_DATA;
+
         let si = StreamInner::new(
             stream_id,
             self.tx.clone(),
@@ -627,7 +627,7 @@ impl HandlerContext {
 
         let task = spawn(async move { stream.handler(ctx, si).await });
 
-        if !req.payload.is_empty() {
+        if !no_data {
             // Fake the first data message.
             let msg = GenMessage {
                 header: MessageHeader::new_data(stream_id, req.payload.len() as u32),