Merge pull request #289 from Tim-Zhang/fix-client-stream-receiver

stream/client: Fix early closing of socket caused by Client's lifetime

Author: teawater

src/asynchronous/client.rs

@@ -55,6 +55,7 @@ impl Client {
         };
 
         let conn = Connection::new(stream, delegate);
+        // Long-running receiver task
         tokio::spawn(async move { conn.run().await });
 
         Client {
@@ -86,9 +87,7 @@ impl Client {
             .map_err(|_| Error::LocalClosed)?;
 
         let result = if timeout_nano == 0 {
-            rx.recv()
-                .await
-                .ok_or_else(|| Error::RemoteClosed)?
+            rx.recv().await.ok_or_else(|| Error::RemoteClosed)?
         } else {
             tokio::time::timeout(
                 std::time::Duration::from_nanos(timeout_nano as u64),

src/asynchronous/connection.rs

@@ -59,6 +59,7 @@ where
 
         let (reader_delegate, mut writer_delegate) = builder.build();
 
+        // Long-running sender task
         let writer_task = tokio::spawn(async move {
             while let Some(mut sending_msg) = writer_delegate.recv().await {
                 trace!("write message: {:?}", sending_msg.msg);

src/asynchronous/stream.rs

@@ -11,6 +11,7 @@ use std::sync::{Arc, Mutex};
 
 use tokio::sync::mpsc;
 
+use super::Client;
 use crate::error::{Error, Result};
 use crate::proto::{
     Code, Codec, GenMessage, MessageHeader, Response, FLAG_NO_DATA, FLAG_REMOTE_CLOSED,
@@ -283,17 +284,20 @@ where
 pub struct ClientStreamReceiver<P> {
     inner: StreamReceiver,
     _recv: PhantomData<P>,
+    // Hold the req_tx in Client to keep receiver task running
+    _client_guard: Client,
 }
 
 impl<P> ClientStreamReceiver<P>
 where
     P: Codec,
     <P as Codec>::E: std::fmt::Display,
 {
-    pub fn new(inner: StreamInner) -> Self {
+    pub fn new(inner: StreamInner, _client_guard: Client) -> Self {
         Self {
             inner: inner.split().1,
             _recv: PhantomData,
+            _client_guard,
         }
     }
 

src/asynchronous/utils.rs

@@ -226,7 +226,7 @@ macro_rules! async_client_stream_receive {
         }
 
         let inner = $self.client.new_stream(creq, false, true).await?;
-        let stream = ::ttrpc::r#async::ClientStreamReceiver::new(inner);
+        let stream = ::ttrpc::r#async::ClientStreamReceiver::new(inner, $self.client.clone());
 
         return Ok(stream);
     };