Merge pull request #51 from Tim-Zhang/support-timeout-for-async

Async: Add support for timeout in ttrpc request

Author: lifupan

Cargo.toml

@@ -20,7 +20,7 @@ byteorder = "1.3.2"
 thiserror = "1.0"
 
 async-trait = { version = "0.1.31", optional = true }
-tokio = { version = "0.2", features = ["rt-threaded", "sync", "uds", "stream", "macros", "io-util"], optional = true }
+tokio = { version = "0.2", features = ["rt-threaded", "sync", "uds", "stream", "macros", "io-util", "time"], optional = true }
 futures = { version = "0.3", optional = true }
 tokio-vsock = { version = "0.2.1", optional = true }
 

src/asynchronous/client.rs

@@ -8,6 +8,7 @@ use protobuf::{CodedInputStream, CodedOutputStream, Message};
 use std::collections::HashMap;
 use std::os::unix::io::RawFd;
 use std::sync::{Arc, Mutex};
+use std::time::Duration;
 
 use crate::common::MESSAGE_TYPE_RESPONSE;
 use crate::error::{Error, Result};
@@ -154,10 +155,22 @@ impl Client {
             .await
             .map_err(|e| Error::Others(format!("Send packet to sender error {:?}", e)))?;
 
-        let result = rx
-            .recv()
-            .await
-            .ok_or_else(|| Error::Others("Recive packet from recver error".to_string()))?;
+        let result: Result<Vec<u8>> = if req.timeout_nano == 0 {
+            rx.recv()
+                .await
+                .ok_or_else(|| Error::Others("Recive packet from recver error".to_string()))?
+        } else {
+            let timeout = tokio::time::delay_for(Duration::from_nanos(req.timeout_nano as u64));
+
+            tokio::select! {
+                result = rx.recv() => {
+                    result.ok_or_else(|| Error::Others("Recive packet from recver error".to_string()))?
+                }
+                _ = timeout => {
+                    return Err(Error::Others("Recive packet from recver error".to_string()));
+                }
+            }
+        };
 
         let buf = result?;
         let mut s = CodedInputStream::from_bytes(&buf);