From e7c3ef525086eec00e41ec84a7678e7d92f527b9 Mon Sep 17 00:00:00 2001
From: Alexey Galakhov <agalakhov@snapview.de>
Date: Tue, 23 May 2017 03:11:34 +0200
Subject: [PATCH] Set TCP_NODELAY on connect.

Signed-off-by: Alexey Galakhov <agalakhov@snapview.de>
---
 Cargo.toml     |  5 +++--
 src/connect.rs | 26 ++++++++++++++++++++++++++
 src/lib.rs     |  5 +++--
 src/stream.rs  | 17 +++++++++++++++++
 4 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 57dcaac..1c9c97e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -12,8 +12,9 @@ version = "0.2.1"
 
 [features]
 default = ["connect", "tls"]
-connect = ["tokio-dns-unofficial", "tokio-core"]
-tls = ["tokio-tls", "native-tls", "bytes"]
+connect = ["tokio-dns-unofficial", "tokio-core", "stream"]
+tls = ["tokio-tls", "native-tls", "stream"]
+stream = ["bytes"]
 
 [dependencies]
 futures = "*"
diff --git a/src/connect.rs b/src/connect.rs
index 6743019..4d43694 100644
--- a/src/connect.rs
+++ b/src/connect.rs
@@ -6,12 +6,23 @@ extern crate tokio_core;
 use self::tokio_dns::tcp_connect;
 use self::tokio_core::reactor::Remote;
 
+use std::io::Result as IoResult;
+
 use futures::{Future, BoxFuture};
 use futures::future;
 
 use super::{WebSocketStream, Request, client_async};
 use tungstenite::Error;
 use tungstenite::client::url_mode;
+use stream::NoDelay;
+
+use self::tokio_core::net::TcpStream;
+
+impl NoDelay for TcpStream {
+    fn set_nodelay(&mut self, nodelay: bool) -> IoResult<()> {
+        TcpStream::set_nodelay(self, nodelay)
+    }
+}
 
 #[cfg(feature="tls")]
 mod encryption {
@@ -23,15 +34,25 @@ mod encryption {
     use self::native_tls::TlsConnector;
     use self::tokio_tls::{TlsConnectorExt, TlsStream};
 
+    use std::io::{Read, Write, Result as IoResult};
+
     use futures::{Future, BoxFuture};
     use futures::future;
 
     use tungstenite::Error;
     use tungstenite::stream::Mode;
 
+    use stream::NoDelay;
+
     pub use stream::Stream as StreamSwitcher;
     pub type AutoStream = StreamSwitcher<TcpStream, TlsStream<TcpStream>>;
 
+    impl<T: Read + Write + NoDelay> NoDelay for TlsStream<T> {
+        fn set_nodelay(&mut self, nodelay: bool) -> IoResult<()> {
+            self.get_mut().get_mut().set_nodelay(nodelay)
+        }
+    }
+
     pub fn wrap_stream(socket: TcpStream, domain: String, mode: Mode) -> BoxFuture<AutoStream, Error> {
         match mode {
             Mode::Plain => future::ok(StreamSwitcher::Plain(socket)).boxed(),
@@ -88,6 +109,11 @@ where R: Into<Request<'static>>
 
     tcp_connect((domain.as_str(), port), handle).map_err(|e| e.into())
     .and_then(move |socket| wrap_stream(socket, domain, mode))
+    .and_then(|mut stream| {
+        NoDelay::set_nodelay(&mut stream, true)
+        .map(move |()| stream)
+        .map_err(|e| e.into())
+     })
     .and_then(move |stream| client_async(request, stream))
     .boxed()
 }
diff --git a/src/lib.rs b/src/lib.rs
index 6076272..203e1a7 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -27,8 +27,9 @@ extern crate url;
 
 #[cfg(feature="connect")]
 mod connect;
-#[cfg(all(feature="connect", feature="tls"))]
-mod stream;
+
+#[cfg(feature="stream")]
+pub mod stream;
 
 use std::io::ErrorKind;
 
diff --git a/src/stream.rs b/src/stream.rs
index 51992e3..d06cb1e 100644
--- a/src/stream.rs
+++ b/src/stream.rs
@@ -12,9 +12,17 @@ use self::bytes::{Buf, BufMut};
 use futures::Poll;
 use tokio_io::{AsyncRead, AsyncWrite};
 
+/// Trait to switch TCP_NODELAY.
+pub trait NoDelay {
+    /// Set the TCP_NODELAY option to the given value.
+    fn set_nodelay(&mut self, nodelay: bool) -> IoResult<()>;
+}
+
 /// Stream, either plain TCP or TLS.
 pub enum Stream<S, T> {
+    /// Unencrypted socket stream.
     Plain(S),
+    /// Encrypted socket stream.
     Tls(T),
 }
 
@@ -42,6 +50,15 @@ impl<S: Write, T: Write> Write for Stream<S, T> {
     }
 }
 
+impl<S: NoDelay, T: NoDelay> NoDelay for Stream<S, T> {
+    fn set_nodelay(&mut self, nodelay: bool) -> IoResult<()> {
+        match *self {
+            Stream::Plain(ref mut s) => s.set_nodelay(nodelay),
+            Stream::Tls(ref mut s) => s.set_nodelay(nodelay),
+        }
+    }
+}
+
 impl<S: AsyncRead, T: AsyncRead> AsyncRead for Stream<S, T> {
     unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool {
         match *self {