From 5adde38a65fb13871e2c397de169f58e2918fe51 Mon Sep 17 00:00:00 2001
From: Sebastian Thiel <byronimo@gmail.com>
Date: Sat, 11 Feb 2017 09:01:07 +0100
Subject: [PATCH] Allocation in get_mut() always uses 8 * 1024 bytes

The previous implementation would always use the capacity
of the previous buffer, which would effectively prevent it
from ever shrinking.

This also means that protocol with greater variance in
possible frame sizes would likely be heavily over-allocating.
If these implementations use zero-copy, this would imply
that even small frames kept alive by the client would use
large amounts of memory.

The change is motivated by the implementation of the
cassandra-protocol, which allows frames of up to 256MB
in size, which solely depend on the kind of query.
---
 src/io/frame.rs | 20 ++++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/src/io/frame.rs b/src/io/frame.rs
index 479433723..44fa5a1b9 100644
--- a/src/io/frame.rs
+++ b/src/io/frame.rs
@@ -8,6 +8,8 @@ use futures::{Async, Poll, Stream, Sink, StartSend, AsyncSink};
 
 use io::Io;
 
+const INITIAL_CAPACITY: usize = 8 * 1024;
+
 /// A reference counted buffer of bytes.
 ///
 /// An `EasyBuf` is a representation of a byte buffer where sub-slices of it can
@@ -31,7 +33,7 @@ pub struct EasyBufMut<'a> {
 impl EasyBuf {
     /// Creates a new EasyBuf with no data and the default capacity.
     pub fn new() -> EasyBuf {
-        EasyBuf::with_capacity(8 * 1024)
+        EasyBuf::with_capacity(INITIAL_CAPACITY)
     }
 
     /// Creates a new EasyBuf with `cap` capacity.
@@ -126,8 +128,8 @@ impl EasyBuf {
     /// If this `EasyBuf` is the only instance pointing at the underlying buffer
     /// of bytes, a direct mutable reference will be returned. Otherwise the
     /// contents of this `EasyBuf` will be reallocated in a fresh `Vec<u8>`
-    /// allocation with the same capacity as this allocation, and that
-    /// allocation will be returned.
+    /// allocation with the same capacity as an `EasyBuf` created with `EasyBuf::new()`,
+    /// and that allocation will be returned.
     ///
     /// This operation **is not O(1)** as it may clone the entire contents of
     /// this buffer.
@@ -150,7 +152,7 @@ impl EasyBuf {
 
         // If we couldn't get access above then we give ourself a new buffer
         // here.
-        let mut v = Vec::with_capacity(self.buf.capacity());
+        let mut v = Vec::with_capacity(INITIAL_CAPACITY);
         v.extend_from_slice(self.as_ref());
         self.start = 0;
         self.buf = Arc::new(v);
@@ -345,9 +347,10 @@ impl<T: Io, C: Codec> Sink for Framed<T, C> {
     fn start_send(&mut self, item: C::Out) -> StartSend<C::Out, io::Error> {
         // If the buffer is already over 8KiB, then attempt to flush it. If after flushing it's
         // *still* over 8KiB, then apply backpressure (reject the send).
-        if self.wr.len() > 8 * 1024 {
+        const BACKPRESSURE_BOUNDARY: usize = INITIAL_CAPACITY;
+        if self.wr.len() > BACKPRESSURE_BOUNDARY {
             try!(self.poll_complete());
-            if self.wr.len() > 8 * 1024 {
+            if self.wr.len() > BACKPRESSURE_BOUNDARY {
                 return Ok(AsyncSink::NotReady(item));
             }
         }
@@ -384,7 +387,7 @@ pub fn framed<T, C>(io: T, codec: C) -> Framed<T, C> {
         eof: false,
         is_readable: false,
         rd: EasyBuf::new(),
-        wr: Vec::with_capacity(8 * 1024),
+        wr: Vec::with_capacity(INITIAL_CAPACITY),
     }
 }
 
@@ -421,7 +424,7 @@ impl<T, C> Framed<T, C> {
 
 #[cfg(test)]
 mod tests {
-    use super::EasyBuf;
+    use super::{INITIAL_CAPACITY, EasyBuf};
     use std::mem;
 
     #[test]
@@ -469,6 +472,7 @@ mod tests {
         // Clone to make shared
         let clone = buf.clone();
         assert_eq!(*buf.get_mut(), [3, 4, 5, 6, 7, 8]);
+        assert_eq!(buf.get_mut().buf.capacity(), INITIAL_CAPACITY);
         mem::drop(clone); // prevent unused warning
     }