BytesMut::reserve should avoid small allocations

This change tracks the original capacity requested when `BytesMut` is
first created. This capacity is used when a `reserve` needs to allocate
due to the current view being too small. The newly allocated buffer will
be sized the same as the original allocation.

src/bytes.rs

@@ -244,11 +244,11 @@ pub struct BytesMut {
 // allocated yet and `self` is the only outstanding handle for the underlying
 // buffer.
 //
-// The lower two bits of `arc` are used as flags to track the storage state of
+// The lower two bits of `arc` are used to track the storage mode of `Inner`.
-// `Inner`. `0b01` indicates inline storage and `0b10` indicates static storage.
+// `0b01` indicates inline storage, `0b10` indicates static storage, and `0b11`
-// Since pointers to allocated structures are aligned, the lower two bits of a
+// indicates vector storage, not yet promoted to Arc.  Since pointers to
-// pointer will always be 0. This allows disambiguating between a pointer and
+// allocated structures are aligned, the lower two bits of a pointer will always
-// the two flags.
+// be 0. This allows disambiguating between a pointer and the two flags.
 //
 // When in "inlined" mode, the least significant byte of `arc` is also used to
 // store the length of the buffer view (vs. the capacity, which is a constant).
@@ -315,16 +315,22 @@ struct Inner2 {
 // other shenanigans to make it work.
 struct Shared {
     vec: Vec<u8>,
+    original_capacity: usize,
     ref_count: AtomicUsize,
 }
 
 // Buffer storage strategy flags.
+const KIND_ARC: usize = 0b00;
 const KIND_INLINE: usize = 0b01;
 const KIND_STATIC: usize = 0b10;
+const KIND_VEC: usize = 0b11;
+const KIND_MASK: usize = 0b11;
+
+const MAX_ORIGINAL_CAPACITY: usize = 1 << 16;
 
 // Bit op constants for extracting the inline length value from the `arc` field.
-const INLINE_LEN_MASK: usize = 0b11111110;
+const INLINE_LEN_MASK: usize = 0b11111100;
-const INLINE_LEN_OFFSET: usize = 1;
+const INLINE_LEN_OFFSET: usize = 2;
 
 // Byte offset from the start of `Inner` to where the inline buffer data
 // starts. On little endian platforms, the first byte of the struct is the
@@ -1302,7 +1308,7 @@ impl Inner {
     #[inline]
     fn empty() -> Inner {
         Inner {
-            arc: AtomicPtr::new(ptr::null_mut()),
+            arc: AtomicPtr::new(KIND_VEC as *mut Shared),
             ptr: ptr::null_mut(),
             len: 0,
             cap: 0,
@@ -1332,8 +1338,11 @@ impl Inner {
 
         mem::forget(src);
 
+        let original_capacity = cmp::min(cap, MAX_ORIGINAL_CAPACITY);
+        let arc = (original_capacity & !KIND_MASK) | KIND_VEC;
+
         Inner {
-            arc: AtomicPtr::new(ptr::null_mut()),
+            arc: AtomicPtr::new(arc as *mut Shared),
             ptr: ptr,
             len: len,
             cap: cap,
@@ -1544,37 +1553,28 @@ impl Inner {
 
     /// Checks if it is safe to mutate the memory
     fn is_mut_safe(&mut self) -> bool {
+        let kind = self.kind();
+
         // Always check `inline` first, because if the handle is using inline
         // data storage, all of the `Inner` struct fields will be gibberish.
-        if self.is_inline() {
+        if kind == KIND_INLINE {
             // Inlined buffers can always be mutated as the data is never shared
             // across handles.
             true
+        } else if kind == KIND_VEC {
+            true
+        } else if kind == KIND_STATIC {
+            false
         } else {
             // The function requires `&mut self`, which guarantees a unique
             // reference to the current handle. This means that the `arc` field
             // *cannot* be concurrently mutated. As such, `Relaxed` ordering is
             // fine (since we aren't synchronizing with anything).
-            //
-            // TODO: No ordering?
             let arc = self.arc.load(Relaxed);
 
-            // If the pointer is null, this is a non-shared handle and is mut
-            // safe.
-            if arc.is_null() {
-                return true;
-            }
-
-            // Check if this is a static buffer
-            if KIND_STATIC == arc as usize {
-                return false;
-            }
-
             // Otherwise, the underlying buffer is potentially shared with other
             // handles, so the ref_count needs to be checked.
-            unsafe {
+            unsafe { (*arc).is_unique() }
-                return (*arc).is_unique();
-            }
         }
     }
 
@@ -1607,10 +1607,10 @@ impl Inner {
             // that the current thread acquires the associated memory.
             let mut arc = self.arc.load(Acquire);
 
-            // If `arc` is null, then the buffer is still tracked in a
+            // If  the buffer is still tracked in a `Vec<u8>`. It is time to
-            // `Vec<u8>`. It is time to promote the vec to an `Arc`. This could
+            // promote the vec to an `Arc`. This could potentially be called
-            // potentially be called concurrently, so some care must be taken.
+            // concurrently, so some care must be taken.
-            if arc.is_null() {
+            if arc as usize & KIND_MASK == KIND_VEC {
                 unsafe {
                     // First, allocate a new `Shared` instance containing the
                     // `Vec` fields. It's important to note that `ptr`, `len`,
@@ -1621,6 +1621,7 @@ impl Inner {
                     // vector.
                     let shared = Box::new(Shared {
                         vec: Vec::from_raw_parts(self.ptr, self.len, self.cap),
+                        original_capacity: arc as usize & !KIND_MASK,
                         // Initialize refcount to 2. One for this reference, and one
                         // for the new clone that will be returned from
                         // `shallow_clone`.
@@ -1644,7 +1645,7 @@ impl Inner {
                     // pointed to by `actual` will be visible.
                     let actual = self.arc.compare_and_swap(arc, shared, AcqRel);
 
-                    if actual.is_null() {
+                    if actual == arc {
                         // The upgrade was successful, the new handle can be
                         // returned.
                         return Inner {
@@ -1663,7 +1664,7 @@ impl Inner {
                     // count update
                     arc = actual;
                 }
-            } else if KIND_STATIC == arc as usize {
+            } else if arc as usize & KIND_MASK == KIND_STATIC {
                 // Static buffer
                 return Inner {
                     arc: AtomicPtr::new(arc),
@@ -1701,9 +1702,11 @@ impl Inner {
             return;
         }
 
+        let kind = self.kind();
+
         // Always check `inline` first, because if the handle is using inline
         // data storage, all of the `Inner` struct fields will be gibberish.
-        if self.is_inline() {
+        if kind == KIND_INLINE {
             let new_cap = len + additional;
 
             // Promote to a vector
@@ -1713,18 +1716,16 @@ impl Inner {
             self.ptr = v.as_mut_ptr();
             self.len = v.len();
             self.cap = v.capacity();
-            self.arc = AtomicPtr::new(ptr::null_mut());
+
+            // Since the minimum capacity is `INLINE_CAP`, don't bother encoding
+            // the original capacity as INLINE_CAP
+            self.arc = AtomicPtr::new(KIND_VEC as *mut Shared);
 
             mem::forget(v);
             return;
         }
 
-        // `Relaxed` is Ok here (and really, no synchronization is necessary)
+        if kind == KIND_VEC {
-        // due to having a `&mut self` pointer. The `&mut self` pointer ensures
-        // that there is no concurrent access on `self`.
-        let arc = self.arc.load(Relaxed);
-
-        if arc.is_null() {
             // Currently backed by a vector, so just use `Vector::reserve`.
             unsafe {
                 let mut v = Vec::from_raw_parts(self.ptr, self.len, self.cap);
@@ -1742,15 +1743,23 @@ impl Inner {
             }
         }
 
-        debug_assert!(!self.is_static());
+        // `Relaxed` is Ok here (and really, no synchronization is necessary)
+        // due to having a `&mut self` pointer. The `&mut self` pointer ensures
+        // that there is no concurrent access on `self`.
+        let arc = self.arc.load(Relaxed);
+
+        debug_assert!(kind == KIND_ARC);
 
         // Reserving involves abandoning the currently shared buffer and
         // allocating a new vector with the requested capacity.
         //
         // Compute the new capacity
         let mut new_cap = len + additional;
+        let original_capacity;
 
         unsafe {
+            original_capacity = (*arc).original_capacity;
+
             // First, try to reclaim the buffer. This is possible if the current
             // handle is the only outstanding handle pointing to the buffer.
             if (*arc).is_unique() {
@@ -1775,7 +1784,15 @@ impl Inner {
                 // asking for more than the initial buffer capacity. Allocate more
                 // than requested if `new_cap` is not much bigger than the current
                 // capacity.
-                new_cap = cmp::max(v.capacity() << 1, new_cap);
+                //
+                // There are some situations, using `reserve_exact` that the
+                // buffer capacity could be below `original_capacity`, so do a
+                // check.
+                new_cap = cmp::max(
+                    cmp::max(v.capacity() << 1, new_cap),
+                    original_capacity);
+            } else {
+                new_cap = cmp::max(new_cap, original_capacity);
             }
         }
 
@@ -1793,7 +1810,10 @@ impl Inner {
         self.ptr = v.as_mut_ptr();
         self.len = v.len();
         self.cap = v.capacity();
-        self.arc = AtomicPtr::new(ptr::null_mut());
+
+        let arc = (original_capacity & !KIND_MASK) | KIND_VEC;
+
+        self.arc = AtomicPtr::new(arc as *mut Shared);
 
         // Forget the vector handle
         mem::forget(v);
@@ -1802,6 +1822,30 @@ impl Inner {
     /// Returns true if the buffer is stored inline
     #[inline]
     fn is_inline(&self) -> bool {
+        self.kind() == KIND_INLINE
+    }
+
+    /// Used for `debug_assert` statements. &mut is used to guarantee that it is
+    /// safe to check VEC_KIND
+    #[inline]
+    fn is_shared(&mut self) -> bool {
+        match self.kind() {
+            KIND_INLINE | KIND_ARC => true,
+            _ => false,
+        }
+    }
+
+    /// Used for `debug_assert` statements
+    #[inline]
+    fn is_static(&mut self) -> bool {
+        match self.kind() {
+            KIND_STATIC => true,
+            _ => false,
+        }
+    }
+
+    #[inline]
+    fn kind(&self) -> usize {
         // This function is going to probably raise some eyebrows. The function
         // returns true if the buffer is stored inline. This is done by checking
         // the least significant bit in the `arc` field.
@@ -1822,67 +1866,40 @@ impl Inner {
 
         #[cfg(target_endian = "little")]
         #[inline]
-        fn imp(arc: &AtomicPtr<Shared>) -> bool {
+        fn imp(arc: &AtomicPtr<Shared>) -> usize {
             unsafe {
                 let p: &u8 = mem::transmute(arc);
-                *p & (KIND_INLINE as u8) == (KIND_INLINE as u8)
+                (*p as usize) & KIND_MASK
             }
         }
 
         #[cfg(target_endian = "big")]
         #[inline]
-        fn imp(arc: &AtomicPtr<Shared>) -> bool {
+        fn imp(arc: &AtomicPtr<Shared>) -> usize {
             unsafe {
                 let p: &usize = mem::transmute(arc);
-                *p & KIND_INLINE == KIND_INLINE
+                *p & KIND_MASK
             }
         }
 
         imp(&self.arc)
     }
-
-    /// Used for `debug_assert` statements
-    #[inline]
-    fn is_shared(&self) -> bool {
-        self.is_inline() ||
-            !self.arc.load(Relaxed).is_null()
-    }
-
-    /// Used for `debug_assert` statements
-    #[inline]
-    fn is_static(&self) -> bool {
-        !self.is_inline() &&
-            self.arc.load(Relaxed) as usize == KIND_STATIC
-    }
 }
 
 impl Drop for Inner2 {
     fn drop(&mut self) {
-        // Always check `inline` first, because if the handle is using inline
+        let kind = self.kind();
-        // data storage, all of the `Inner` struct fields will be gibberish.
-        if self.is_inline() {
-            return;
-        }
 
-        // Acquire is needed here to ensure that the `Shared` memory is
+        if kind == KIND_VEC {
-        // visible.
-        let arc = self.arc.load(Acquire);
-
-        if arc as usize == KIND_STATIC {
-            // Static buffer, no work to do
-            return;
-        }
-
-        if arc.is_null() {
             // Vector storage, free the vector
             unsafe {
                 let _ = Vec::from_raw_parts(self.ptr, self.len, self.cap);
             }
-
+        } else if kind == KIND_ARC {
-            return;
+            // &mut self guarantees correct ordering
+            let arc = self.arc.load(Relaxed);
+            release_shared(arc);
         }
-
-        release_shared(arc);
     }
 }
 

tests/test_bytes.rs

@@ -251,6 +251,36 @@ fn reserve_growth() {
     assert_eq!(bytes.capacity(), 128);
 }
 
+#[test]
+fn reserve_allocates_at_least_original_capacity() {
+    let mut bytes = BytesMut::with_capacity(128);
+
+    for i in 0..120 {
+        bytes.put(i as u8);
+    }
+
+    let _other = bytes.take();
+
+    bytes.reserve(16);
+    assert_eq!(bytes.capacity(), 128);
+}
+
+#[test]
+fn reserve_max_original_capacity_value() {
+    const SIZE: usize = 128 * 1024;
+
+    let mut bytes = BytesMut::with_capacity(SIZE);
+
+    for _ in 0..SIZE {
+        bytes.put(0u8);
+    }
+
+    let _other = bytes.take();
+
+    bytes.reserve(16);
+    assert_eq!(bytes.capacity(), 64 * 1024);
+}
+
 #[test]
 fn inline_storage() {
     let mut bytes = BytesMut::with_capacity(inline_cap());