diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs index ec74c4e97..e60709fa9 100644 --- a/src/bytes_mut.rs +++ b/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 { let len = self.len(); let kind = self.kind(); @@ -649,13 +715,13 @@ 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)); @@ -663,7 +729,7 @@ impl BytesMut { 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`.