From f4a47ef71e5c0289e5ae25613912e08689c17d84 Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Sun, 18 Mar 2018 19:27:51 -0700
Subject: [PATCH] allow set client request/ws timeout

---
 CHANGES.md            |  4 ++++
 src/client/request.rs | 22 +++++++++++++++++++++-
 src/ws/client.rs      |  9 +++++++++
 3 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/CHANGES.md b/CHANGES.md
index 0c19ba0f..7ff63f66 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -4,6 +4,10 @@
 
 * Use `Error` instead of `InternalError` for `error::ErrorXXXX` methods
 
+* Allow to set client request timeout
+
+* Allow to set client websocket handshake timeout
+
 
 ## 0.4.9 (2018-03-16)
 
diff --git a/src/client/request.rs b/src/client/request.rs
index c5c95eb8..09f6d0c7 100644
--- a/src/client/request.rs
+++ b/src/client/request.rs
@@ -1,6 +1,7 @@
 use std::{fmt, mem};
 use std::fmt::Write as FmtWrite;
 use std::io::Write;
+use std::time::Duration;
 
 use actix::{Addr, Unsync};
 use cookie::{Cookie, CookieJar};
@@ -26,6 +27,7 @@ pub struct ClientRequest {
     body: Body,
     chunked: bool,
     upgrade: bool,
+    timeout: Option<Duration>,
     encoding: ContentEncoding,
     response_decompress: bool,
     buffer_capacity: usize,
@@ -49,6 +51,7 @@ impl Default for ClientRequest {
             body: Body::Empty,
             chunked: false,
             upgrade: false,
+            timeout: None,
             encoding: ContentEncoding::Auto,
             response_decompress: true,
             buffer_capacity: 32_768,
@@ -204,10 +207,16 @@ impl ClientRequest {
     ///
     /// This method returns future that resolves to a ClientResponse
     pub fn send(mut self) -> SendRequest {
-        match mem::replace(&mut self.conn, ConnectionType::Default) {
+        let timeout = self.timeout.take();
+        let send = match mem::replace(&mut self.conn, ConnectionType::Default) {
             ConnectionType::Default => SendRequest::new(self),
             ConnectionType::Connector(conn) => SendRequest::with_connector(self, conn),
             ConnectionType::Connection(conn) => SendRequest::with_connection(self, conn),
+        };
+        if let Some(timeout) = timeout {
+            send.timeout(timeout)
+        } else {
+            send
         }
     }
 }
@@ -476,6 +485,17 @@ impl ClientRequestBuilder {
         self
     }
 
+    /// Set request timeout
+    ///
+    /// Request timeout is a total time before response should be received.
+    /// Default value is 5 seconds.
+    pub fn timeout(&mut self, timeout: Duration) -> &mut Self {
+        if let Some(parts) = parts(&mut self.request, &self.err) {
+            parts.timeout = Some(timeout);
+        }
+        self
+    }
+
     /// Send request using custom connector
     pub fn with_connector(&mut self, conn: Addr<Unsync, ClientConnector>) -> &mut Self {
         if let Some(parts) = parts(&mut self.request, &self.err) {
diff --git a/src/ws/client.rs b/src/ws/client.rs
index cb406dbe..049930ea 100644
--- a/src/ws/client.rs
+++ b/src/ws/client.rs
@@ -209,6 +209,15 @@ impl Client {
         self
     }
 
+    /// Set websocket handshake timeout
+    ///
+    /// Handshake timeout is a total time for successful handshake.
+    /// Default value is 5 seconds.
+    pub fn timeout(mut self, timeout: Duration) -> Self {
+        self.request.timeout(timeout);
+        self
+    }
+
     /// Connect to websocket server and do ws handshake
     pub fn connect(&mut self) -> ClientHandshake {
         if let Some(e) = self.err.take() {