From b56be8e5718367a9242035ba71edf5245e7c185e Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Fri, 9 Mar 2018 10:09:13 -0800
Subject: [PATCH] write buffer capacity for client

---
 src/client/request.rs | 16 ++++++++--------
 src/client/writer.rs  | 23 ++++++-----------------
 src/ws/client.rs      |  8 ++++++++
 3 files changed, 22 insertions(+), 25 deletions(-)

diff --git a/src/client/request.rs b/src/client/request.rs
index 095f8d7f6..1f564753a 100644
--- a/src/client/request.rs
+++ b/src/client/request.rs
@@ -28,7 +28,7 @@ pub struct ClientRequest {
     upgrade: bool,
     encoding: ContentEncoding,
     response_decompress: bool,
-    buffer_capacity: Option<(usize, usize)>,
+    buffer_capacity: usize,
     conn: ConnectionType,
 }
 
@@ -51,7 +51,7 @@ impl Default for ClientRequest {
             upgrade: false,
             encoding: ContentEncoding::Auto,
             response_decompress: true,
-            buffer_capacity: None,
+            buffer_capacity: 32_768,
             conn: ConnectionType::Default,
         }
     }
@@ -179,7 +179,8 @@ impl ClientRequest {
         self.response_decompress
     }
 
-    pub fn buffer_capacity(&self) -> Option<(usize, usize)> {
+    /// Requested write buffer capacity
+    pub fn write_buffer_capacity(&self) -> usize {
         self.buffer_capacity
     }
     
@@ -466,12 +467,11 @@ impl ClientRequestBuilder {
     }
 
     /// Set write buffer capacity
-    pub fn buffer_capacity(&mut self,
-                           low_watermark: usize,
-                           high_watermark: usize) -> &mut Self
-    {
+    ///
+    /// Default buffer capacity is 32kb
+    pub fn write_buffer_capacity(&mut self, cap: usize) -> &mut Self {
         if let Some(parts) = parts(&mut self.request, &self.err) {
-            parts.buffer_capacity = Some((low_watermark, high_watermark));
+            parts.buffer_capacity = cap;
         }
         self
     }
diff --git a/src/client/writer.rs b/src/client/writer.rs
index f67bd7261..7cd522113 100644
--- a/src/client/writer.rs
+++ b/src/client/writer.rs
@@ -24,8 +24,6 @@ use server::encoding::{ContentEncoder, TransferEncoding};
 use client::ClientRequest;
 
 
-const LOW_WATERMARK: usize = 1024;
-const HIGH_WATERMARK: usize = 8 * LOW_WATERMARK;
 const AVERAGE_HEADER_SIZE: usize = 30;
 
 bitflags! {
@@ -42,9 +40,8 @@ pub(crate) struct HttpClientWriter {
     written: u64,
     headers_size: u32,
     buffer: SharedBytes,
+    buffer_capacity: usize,
     encoder: ContentEncoder,
-    low: usize,
-    high: usize,
 }
 
 impl HttpClientWriter {
@@ -55,10 +52,9 @@ impl HttpClientWriter {
             flags: Flags::empty(),
             written: 0,
             headers_size: 0,
+            buffer_capacity: 0,
             buffer,
             encoder,
-            low: LOW_WATERMARK,
-            high: HIGH_WATERMARK,
         }
     }
 
@@ -70,12 +66,6 @@ impl HttpClientWriter {
     //    self.flags.contains(Flags::KEEPALIVE) && !self.flags.contains(Flags::UPGRADE)
     // }
 
-    /// Set write buffer capacity
-    pub fn set_buffer_capacity(&mut self, low_watermark: usize, high_watermark: usize) {
-        self.low = low_watermark;
-        self.high = high_watermark;
-    }
-
     fn write_to_stream<T: AsyncWrite>(&mut self, stream: &mut T) -> io::Result<WriterState> {
         while !self.buffer.is_empty() {
             match stream.write(self.buffer.as_ref()) {
@@ -87,7 +77,7 @@ impl HttpClientWriter {
                     let _ = self.buffer.split_to(n);
                 },
                 Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
-                    if self.buffer.len() > self.high {
+                    if self.buffer.len() > self.buffer_capacity {
                         return Ok(WriterState::Pause)
                     } else {
                         return Ok(WriterState::Done)
@@ -106,9 +96,6 @@ impl HttpClientWriter {
         // prepare task
         self.flags.insert(Flags::STARTED);
         self.encoder = content_encoder(self.buffer.clone(), msg);
-        if let Some(capacity) = msg.buffer_capacity() {
-            self.set_buffer_capacity(capacity.0, capacity.1);
-        }
 
         // render message
         {
@@ -153,6 +140,8 @@ impl HttpClientWriter {
                     self.written += bytes.len() as u64;
                     self.encoder.write(bytes)?;
                 }
+            } else {
+                self.buffer_capacity = msg.write_buffer_capacity();
             }
         }
         Ok(())
@@ -168,7 +157,7 @@ impl HttpClientWriter {
             }
         }
 
-        if self.buffer.len() > self.high {
+        if self.buffer.len() > self.buffer_capacity {
             Ok(WriterState::Pause)
         } else {
             Ok(WriterState::Done)
diff --git a/src/ws/client.rs b/src/ws/client.rs
index 42dacc847..80169b324 100644
--- a/src/ws/client.rs
+++ b/src/ws/client.rs
@@ -192,6 +192,14 @@ impl Client {
         self
     }
 
+    /// Set write buffer capacity
+    ///
+    /// Default buffer capacity is 32kb
+    pub fn write_buffer_capacity(mut self, cap: usize) -> Self {
+        self.request.write_buffer_capacity(cap);
+        self
+    }
+
     /// Set request header
     pub fn header<K, V>(mut self, key: K, value: V) -> Self
         where HeaderName: HttpTryFrom<K>, V: IntoHeaderValue