Add BytesMut::try_reserve() #484

src/bytes_mut.rs

@@ -7,6 +7,7 @@ use core::{cmp, fmt, hash, isize, slice, usize};
 use alloc::{
     borrow::{Borrow, BorrowMut},
     boxed::Box,
+    collections::TryReserveError,
     string::String,
     vec,
     vec::Vec,
@@ -588,22 +589,87 @@ impl BytesMut {
     /// Panics if the new capacity overflows `usize`.
     #[inline]
     pub fn reserve(&mut self, additional: usize) {
+        match self.try_reserve(additional) {
+            Err(err) => panic!("fail to reserve: {}", err),
+            Ok(_) => {}
+        }
+    }
+
+    /// Tries to reserves capacity for at least `additional` more bytes to be inserted
+    /// into the given `BytesMut`.
+    ///
+    /// More than `additional` bytes may be reserved in order to avoid frequent
+    /// reallocations. A call to `try_reserve` may result in an allocation.
+    ///
+    /// Before allocating new buffer space, the function will attempt to reclaim
+    /// space in the existing buffer. If the current handle references a small
+    /// view in the original buffer and all other handles have been dropped,
+    /// and the requested capacity is less than or equal to the existing
+    /// buffer's capacity, then the current view will be copied to the front of
+    /// the buffer and the handle will take ownership of the full buffer.
+    ///
+    /// # Errors
+    ///
+    /// If the capacity overflows, or the allocator reports a failure, then an error is returned.
+    ///
+    /// # Examples
+    ///
+    /// In the following example, a new buffer is allocated.
+    ///
+    /// ```
+    /// use bytes::BytesMut;
+    ///
+    /// let mut buf = BytesMut::from(&b"hello"[..]);
+    /// let res = buf.try_reserve(64);
+    /// assert!(res.is_ok());
+    /// assert!(buf.capacity() >= 69);
+    /// ```
+    ///
+    /// In the following example, the existing buffer is reclaimed.
+    ///
+    /// ```
+    /// use bytes::{BytesMut, BufMut};
+    ///
+    /// let mut buf = BytesMut::with_capacity(128);
+    /// buf.put(&[0; 64][..]);
+    ///
+    /// let ptr = buf.as_ptr();
+    /// let other = buf.split();
+    ///
+    /// assert!(buf.is_empty());
+    /// assert_eq!(buf.capacity(), 64);
+    ///
+    /// drop(other);
+    /// let res = buf.try_reserve(128);
+    ///
+    /// assert!(res.is_ok());
+    /// assert_eq!(buf.capacity(), 128);
+    /// assert_eq!(buf.as_ptr(), ptr);
+    /// ```
+    #[inline]
+    pub fn try_reserve(&mut self, additional: usize) -> Result<(), TryReserveError> {
         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;
+            return Ok(());
         }
 
-        // will always succeed
-        let _ = self.reserve_inner(additional, true);
+        self.reserve_inner(additional, true).map(
+            // will always succeed
+            |_| (),
+        )
     }
 
     // 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 {
+    fn reserve_inner(
+        &mut self,
+        additional: usize,
+        allocate: bool,
+    ) -> Result<bool, TryReserveError> {
         let len = self.len();
         let kind = self.kind();
 
@@ -649,21 +715,21 @@ impl BytesMut {
                     self.cap += off;
                 } else {
                     if !allocate {
-                        return false;
+                        return Ok(false);
                     }
                     // Not enough space, or reusing might be too much overhead:
                     // allocate more space!
                     let mut v =
                         ManuallyDrop::new(rebuild_vec(self.ptr.as_ptr(), self.len, self.cap, off));
-                    v.reserve(additional);
+                    v.try_reserve(additional)?;
 
                     // Update the info
                     self.ptr = vptr(v.as_mut_ptr().add(off));
                     self.cap = v.capacity() - off;
                     debug_assert_eq!(self.len, v.len() - off);
                 }
 
-                return true;
+                return Ok(true);
             }
         }
 
@@ -676,7 +742,7 @@ impl BytesMut {
         // Compute the new capacity
         let mut new_cap = match len.checked_add(additional) {
             Some(new_cap) => new_cap,
-            None if !allocate => return false,
+            None if !allocate => return Ok(false),
             None => panic!("overflow"),
         };
 
@@ -710,7 +776,7 @@ impl BytesMut {
                     self.cap = v.capacity();
                 } else {
                     if !allocate {
-                        return false;
+                        return Ok(false);
                     }
                     // calculate offset
                     let off = (self.ptr.as_ptr() as usize) - (v.as_ptr() as usize);
@@ -743,18 +809,18 @@ impl BytesMut {
                     // care about in the unused capacity before calling `reserve`.
                     debug_assert!(off + len <= v.capacity());
                     v.set_len(off + len);
-                    v.reserve(new_cap - v.len());
+                    v.try_reserve(new_cap - v.len())?;
 
                     // Update the info
                     self.ptr = vptr(v.as_mut_ptr().add(off));
                     self.cap = v.capacity() - off;
                 }
 
-                return true;
+                return Ok(true);
             }
         }
         if !allocate {
-            return false;
+            return Ok(false);
         }
 
         let original_capacity_repr = unsafe { (*shared).original_capacity_repr };
@@ -763,7 +829,9 @@ impl BytesMut {
         new_cap = cmp::max(new_cap, original_capacity);
 
         // Create a new vector to store the data
-        let mut v = ManuallyDrop::new(Vec::with_capacity(new_cap));
+        let mut v = Vec::new();
+        v.try_reserve(new_cap)?;
+        let mut v = ManuallyDrop::new(v);
 
         // Copy the bytes
         v.extend_from_slice(self.as_ref());
@@ -778,7 +846,7 @@ impl BytesMut {
         self.ptr = vptr(v.as_mut_ptr());
         self.cap = v.capacity();
         debug_assert_eq!(self.len, v.len());
-        return true;
+        Ok(true)
     }
 
     /// Attempts to cheaply reclaim already allocated capacity for at least `additional` more
@@ -838,7 +906,10 @@ impl BytesMut {
             return true;
         }
 
-        self.reserve_inner(additional, false)
+        match self.reserve_inner(additional, false) {
+            Err(err) => panic!("fail to reserve: {}", err),
+            Ok(r) => r,
+        }
     }
 
     /// Appends given bytes to this `BytesMut`.