Improve handling of socket errors in the async client

Previously, socket errors in `ttrpc::r#async::Client` were not properly
handled. If such errors were occurred during an RPC, the resources
associated with the RPC will become leaked, causing the RPC to hang forever.

This commit tries to properly deliver socket errors to all pending RPC
callers, as well as stop any further RPC requests from being processed.

Signed-off-by: Yu Chen <[email protected]>

Author: Kazurin-775

src/asynchronous/client.rs

@@ -54,7 +54,7 @@ impl Client {
         let notify2 = notify.clone();
 
         // Request sender
-        tokio::spawn(async move {
+        let request_sender = tokio::spawn(async move {
             let mut stream_id: u32 = 1;
 
             while let Some((body, resp_tx)) = rx.recv().await {
@@ -93,14 +93,14 @@ impl Client {
         // Response receiver
         tokio::spawn(async move {
             loop {
-                let req_map = req_map.clone();
                 tokio::select! {
                     _ = notify2.notified() => {
                         break;
                     }
                     res = receive(&mut reader) => {
                         match res {
                             Ok((header, body)) => {
+                                let req_map = req_map.clone();
                                 tokio::spawn(async move {
                                     let resp_tx2;
                                     {
@@ -135,7 +135,23 @@ impl Client {
                                 });
                             }
                             Err(e) => {
-                                trace!("error {:?}", e);
+                                debug!("Connection closed by the ttRPC server: {}", e);
+
+                                // Abort the request sender task to prevent incoming RPC requests
+                                // from being processed.
+                                request_sender.abort();
+                                let _ = request_sender.await;
+
+                                // Take all items out of `req_map`.
+                                let mut map = std::mem::take(&mut *req_map.lock().unwrap());
+                                // Terminate outstanding RPC requests with the error.
+                                for (_stream_id, resp_tx) in map.drain() {
+                                    if let Err(_e) = resp_tx.send(Err(e.clone())).await {
+                                        warn!("Failed to terminate pending RPC: \
+                                               the request has returned");
+                                    }
+                                }
+
                                 break;
                             }
                         }

src/error.rs

@@ -19,7 +19,7 @@ use std::result;
 use thiserror::Error;
 
 /// The error type for ttrpc.
-#[derive(Error, Debug)]
+#[derive(Error, Debug, Clone)]
 pub enum Error {
     #[error("socket err: {0}")]
     Socket(String),