From d745a4f5a737a0c60bbaef6f84d14c9bd3e056a3 Mon Sep 17 00:00:00 2001 From: Sebastian Hahn Date: Sun, 31 Mar 2024 10:23:08 +0200 Subject: [PATCH] Allow reclaiming the current allocation This fixes #680, where it was noted that it is hard to use BytesMut without additional allocations in some circumstances. --- src/bytes_mut.rs | 115 +++++++++++++++++++++++++++++++++++++++++++- tests/test_bytes.rs | 56 +++++++++++++++++++++ 2 files changed, 170 insertions(+), 1 deletion(-) diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs index 537f01ad3..b00ea03b6 100644 --- a/src/bytes_mut.rs +++ b/src/bytes_mut.rs @@ -593,7 +593,7 @@ impl BytesMut { } // In separate function to allow the short-circuits in `reserve` to - // be inline-able. Significant helps performance. + // be inline-able. Significantly helps performance. fn reserve_inner(&mut self, additional: usize) { let len = self.len(); let kind = self.kind(); @@ -758,6 +758,119 @@ impl BytesMut { debug_assert_eq!(self.len, v.len()); } + /// Attempts to cheaply reclaim capacity for at least `additional` more bytes to be inserted + /// into the given `BytesMut` and returns `true` if it succeeded. + /// + /// `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` is currently empty and + /// there are no outstanding references through other `BytesMut`s or `Bytes` which point to the + /// same underlying storage. + /// + /// A call to `try_reclaim` never copies inside the allocation nor allocates new 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; + } + + if !self.is_empty() { + // try_reclaim never copies any bytes but there are some bytes stored already + return false; + } + + let kind = self.kind(); + if kind == KIND_VEC { + // Safety: self is of KIND_VEC, so calling get_vec_pos() is safe. + let off = unsafe { self.get_vec_pos() }; + // The whole allocation is too small to fit the request + if additional > rem + off { + return false; + } + + // Otherwise, update info to point at the front of the vector + // again and reclaim capacity + // + // Safety: Offset `off` means self.ptr was moved `off` bytes from the base + // pointer of the allocation. Going back to the base pointer stays within + // that same allocation. + let base_ptr = unsafe { self.ptr.as_ptr().sub(off) }; + self.ptr = vptr(base_ptr); + + // Safety: Resetting the offset to 0 is safe as we reset the storage + // pointer to the base pointer of the allocation above + unsafe { self.set_vec_pos(0) } + self.cap += off; + + return true; + } + + debug_assert_eq!(kind, KIND_ARC); + let shared: *mut Shared = self.data; + + // If there are other references to the underlying storage, we cannot reclaim it + // + // Safety: self is of type KIND_ARC, so the Shared ptr is valid + if unsafe { !(*shared).is_unique() } { + return false; + } + + // Safety: self is of type KIND_ARC and there are no other handles alive, so we + // can get a mut reference to the underlying storageq + let v = unsafe { &mut (*shared).vec }; + let cap = v.capacity(); + if additional > cap { + // The underlying storage does not have enough capacity + return false; + } + + // Update info to point at start of allocation and reclaim capacity + let ptr = v.as_mut_ptr(); + self.ptr = vptr(ptr); + self.cap = cap; + true + } + /// Appends given bytes to this `BytesMut`. /// /// If this `BytesMut` object does not have enough capacity, it is resized diff --git a/tests/test_bytes.rs b/tests/test_bytes.rs index 2f283af2f..2fb547ec5 100644 --- a/tests/test_bytes.rs +++ b/tests/test_bytes.rs @@ -1283,3 +1283,59 @@ fn test_bytes_make_mut_promotable_even_arc_offset() { assert_eq!(b2m, vec[20..]); assert_eq!(b1m, vec[..20]); } + +#[test] +fn try_reserve_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_reserve_vec() { + let mut buf = BytesMut::with_capacity(6); + buf.put_slice(b"abc"); + assert_eq!(false, buf.try_reclaim(6)); + buf.advance(3); + assert_eq!(true, buf.try_reclaim(6)); + assert_eq!(6, buf.capacity()); +} + +#[test] +fn try_reserve_arc() { + let mut buf = BytesMut::with_capacity(6); + buf.put_slice(b"abc"); + let x = buf.split().freeze(); + buf.put_slice(b"def"); + 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(6); + assert_eq!(true, buf.try_reclaim(6)); +}