Allow reclaiming the current allocation (#686)

Author: shahn

src/bytes_mut.rs

@@ -597,12 +597,13 @@ impl BytesMut {
             return;
         }
 
-        self.reserve_inner(additional);
+        // will always succeed
+        let _ = self.reserve_inner(additional, true);
     }
 
-    // In separate function to allow the short-circuits in `reserve` to
-    // be inline-able. Significant helps performance.
-    fn reserve_inner(&mut self, additional: usize) {
+    // In separate function to allow the short-circuits in `reserve` and `try_reclaim` to
+    // be inline-able. Significantly helps performance. Returns false if it did not succeed.
+    fn reserve_inner(&mut self, additional: usize, allocate: bool) -> bool {
         let len = self.len();
         let kind = self.kind();
 
@@ -647,6 +648,9 @@ impl BytesMut {
                     // can gain capacity back.
                     self.cap += off;
                 } else {
+                    if !allocate {
+                        return false;
+                    }
                     // Not enough space, or reusing might be too much overhead:
                     // allocate more space!
                     let mut v =
@@ -659,7 +663,7 @@ impl BytesMut {
                     debug_assert_eq!(self.len, v.len() - off);
                 }
 
-                return;
+                return true;
             }
         }
 
@@ -670,7 +674,11 @@ impl BytesMut {
         // allocating a new vector with the requested capacity.
         //
         // Compute the new capacity
-        let mut new_cap = len.checked_add(additional).expect("overflow");
+        let mut new_cap = match len.checked_add(additional) {
+            Some(new_cap) => new_cap,
+            None if !allocate => return false,
+            None => panic!("overflow"),
+        };
 
         unsafe {
             // First, try to reclaim the buffer. This is possible if the current
@@ -701,6 +709,9 @@ impl BytesMut {
                     self.ptr = vptr(ptr);
                     self.cap = v.capacity();
                 } else {
+                    if !allocate {
+                        return false;
+                    }
                     // calculate offset
                     let off = (self.ptr.as_ptr() as usize) - (v.as_ptr() as usize);
 
@@ -739,9 +750,12 @@ impl BytesMut {
                     self.cap = v.capacity() - off;
                 }
 
-                return;
+                return true;
             }
         }
+        if !allocate {
+            return false;
+        }
 
         let original_capacity_repr = unsafe { (*shared).original_capacity_repr };
         let original_capacity = original_capacity_from_repr(original_capacity_repr);
@@ -764,6 +778,67 @@ impl BytesMut {
         self.ptr = vptr(v.as_mut_ptr());
         self.cap = v.capacity();
         debug_assert_eq!(self.len, v.len());
+        return true;
+    }
+
+    /// Attempts to cheaply reclaim already allocated capacity for at least `additional` more
+    /// bytes to be inserted into the given `BytesMut` and returns `true` if it succeeded.
+    ///
+    /// `try_reclaim` behaves exactly like `reserve`, except that it never allocates new storage
+    /// and returns a `bool` indicating whether it was successful in doing so:
+    ///
+    /// `try_reclaim` returns false under these conditions:
+    ///  - The spare capacity left is less than `additional` bytes AND
+    ///  - The existing allocation cannot be reclaimed cheaply or it was less than
+    ///    `additional` bytes in size
+    ///
+    /// Reclaiming the allocation cheaply is possible if the `BytesMut` has no outstanding
+    /// references through other `BytesMut`s or `Bytes` which point to the same underlying
+    /// storage.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bytes::BytesMut;
+    ///
+    /// let mut buf = BytesMut::with_capacity(64);
+    /// assert_eq!(true, buf.try_reclaim(64));
+    /// assert_eq!(64, buf.capacity());
+    ///
+    /// buf.extend_from_slice(b"abcd");
+    /// let mut split = buf.split();
+    /// assert_eq!(60, buf.capacity());
+    /// assert_eq!(4, split.capacity());
+    /// assert_eq!(false, split.try_reclaim(64));
+    /// assert_eq!(false, buf.try_reclaim(64));
+    /// // The split buffer is filled with "abcd"
+    /// assert_eq!(false, split.try_reclaim(4));
+    /// // buf is empty and has capacity for 60 bytes
+    /// assert_eq!(true, buf.try_reclaim(60));
+    ///
+    /// drop(buf);
+    /// assert_eq!(false, split.try_reclaim(64));
+    ///
+    /// split.clear();
+    /// assert_eq!(4, split.capacity());
+    /// assert_eq!(true, split.try_reclaim(64));
+    /// assert_eq!(64, split.capacity());
+    /// ```
+    // I tried splitting out try_reclaim_inner after the short circuits, but it was inlined
+    // regardless with Rust 1.78.0 so probably not worth it
+    #[inline]
+    #[must_use = "consider BytesMut::reserve if you need an infallible reservation"]
+    pub fn try_reclaim(&mut self, additional: usize) -> bool {
+        let len = self.len();
+        let rem = self.capacity() - len;
+
+        if additional <= rem {
+            // The handle can already store at least `additional` more bytes, so
+            // there is no further work needed to be done.
+            return true;
+        }
+
+        self.reserve_inner(additional, false)
     }
 
     /// Appends given bytes to this `BytesMut`.

tests/test_bytes.rs

@@ -1283,3 +1283,73 @@ fn test_bytesmut_from_bytes_promotable_even_arc_offset() {
     assert_eq!(b2m, vec[20..]);
     assert_eq!(b1m, vec[..20]);
 }
+
+#[test]
+fn try_reclaim_empty() {
+    let mut buf = BytesMut::new();
+    assert_eq!(false, buf.try_reclaim(6));
+    buf.reserve(6);
+    assert_eq!(true, buf.try_reclaim(6));
+    let cap = buf.capacity();
+    assert!(cap >= 6);
+    assert_eq!(false, buf.try_reclaim(cap + 1));
+
+    let mut buf = BytesMut::new();
+    buf.reserve(6);
+    let cap = buf.capacity();
+    assert!(cap >= 6);
+    let mut split = buf.split();
+    drop(buf);
+    assert_eq!(0, split.capacity());
+    assert_eq!(true, split.try_reclaim(6));
+    assert_eq!(false, split.try_reclaim(cap + 1));
+}
+
+#[test]
+fn try_reclaim_vec() {
+    let mut buf = BytesMut::with_capacity(6);
+    buf.put_slice(b"abc");
+    // Reclaiming a ludicrous amount of space should calmly return false
+    assert_eq!(false, buf.try_reclaim(usize::MAX));
+
+    assert_eq!(false, buf.try_reclaim(6));
+    buf.advance(2);
+    assert_eq!(4, buf.capacity());
+    // We can reclaim 5 bytes, because the byte in the buffer can be moved to the front. 6 bytes
+    // cannot be reclaimed because there is already one byte stored
+    assert_eq!(false, buf.try_reclaim(6));
+    assert_eq!(true, buf.try_reclaim(5));
+    buf.advance(1);
+    assert_eq!(true, buf.try_reclaim(6));
+    assert_eq!(6, buf.capacity());
+}
+
+#[test]
+fn try_reclaim_arc() {
+    let mut buf = BytesMut::with_capacity(6);
+    buf.put_slice(b"abc");
+    let x = buf.split().freeze();
+    buf.put_slice(b"def");
+    // Reclaiming a ludicrous amount of space should calmly return false
+    assert_eq!(false, buf.try_reclaim(usize::MAX));
+
+    let y = buf.split().freeze();
+    let z = y.clone();
+    assert_eq!(false, buf.try_reclaim(6));
+    drop(x);
+    drop(z);
+    assert_eq!(false, buf.try_reclaim(6));
+    drop(y);
+    assert_eq!(true, buf.try_reclaim(6));
+    assert_eq!(6, buf.capacity());
+    assert_eq!(0, buf.len());
+    buf.put_slice(b"abc");
+    buf.put_slice(b"def");
+    assert_eq!(6, buf.capacity());
+    assert_eq!(6, buf.len());
+    assert_eq!(false, buf.try_reclaim(6));
+    buf.advance(4);
+    assert_eq!(true, buf.try_reclaim(4));
+    buf.advance(2);
+    assert_eq!(true, buf.try_reclaim(6));
+}