From d745a4f5a737a0c60bbaef6f84d14c9bd3e056a3 Mon Sep 17 00:00:00 2001 From: Sebastian Hahn Date: Sun, 31 Mar 2024 10:23:08 +0200 Subject: [PATCH 1/3] 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)); +} From 26c14f000d42f86b8642c75604a4a93dbe603e1b Mon Sep 17 00:00:00 2001 From: Sebastian Hahn Date: Mon, 10 Jun 2024 00:15:45 +0200 Subject: [PATCH 2/3] Allow copying bytes in `BytesMut::try_reclaim` This has the nice side effect of making the implementation simpler with regards to unsafe code, as it mostly follows what `reserve` does. --- src/bytes_mut.rs | 94 +++++++++++++-------------------------------- tests/test_bytes.rs | 12 +++++- 2 files changed, 36 insertions(+), 70 deletions(-) diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs index b00ea03b6..cb7d16bc4 100644 --- a/src/bytes_mut.rs +++ b/src/bytes_mut.rs @@ -589,12 +589,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. Significantly 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(); @@ -639,6 +640,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 = @@ -651,7 +655,7 @@ impl BytesMut { debug_assert_eq!(self.len, v.len() - off); } - return; + return true; } } @@ -693,6 +697,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); @@ -731,9 +738,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); @@ -756,21 +766,23 @@ 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 capacity for at least `additional` more bytes to be inserted - /// into the given `BytesMut` and returns `true` if it succeeded. + /// 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` 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. + /// 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 /// @@ -814,61 +826,7 @@ impl BytesMut { 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 + self.reserve_inner(additional, false) } /// Appends given bytes to this `BytesMut`. diff --git a/tests/test_bytes.rs b/tests/test_bytes.rs index 2fb547ec5..a20156345 100644 --- a/tests/test_bytes.rs +++ b/tests/test_bytes.rs @@ -1310,7 +1310,13 @@ 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); + 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()); } @@ -1336,6 +1342,8 @@ fn try_reserve_arc() { assert_eq!(6, buf.capacity()); assert_eq!(6, buf.len()); assert_eq!(false, buf.try_reclaim(6)); - buf.advance(6); + buf.advance(4); + assert_eq!(true, buf.try_reclaim(4)); + buf.advance(2); assert_eq!(true, buf.try_reclaim(6)); } From c5db37397ee3abdca5d7332737718b3f9ddbe786 Mon Sep 17 00:00:00 2001 From: Sebastian Hahn Date: Thu, 27 Jun 2024 22:44:43 +0200 Subject: [PATCH 3/3] Don't panic when trying to reclaim a huge number Implementing the suggestion by @conradludgate --- src/bytes_mut.rs | 6 +++++- tests/test_bytes.rs | 12 +++++++++--- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs index cb7d16bc4..10f51b1db 100644 --- a/src/bytes_mut.rs +++ b/src/bytes_mut.rs @@ -666,7 +666,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 diff --git a/tests/test_bytes.rs b/tests/test_bytes.rs index a20156345..aad162e0f 100644 --- a/tests/test_bytes.rs +++ b/tests/test_bytes.rs @@ -1285,7 +1285,7 @@ fn test_bytes_make_mut_promotable_even_arc_offset() { } #[test] -fn try_reserve_empty() { +fn try_reclaim_empty() { let mut buf = BytesMut::new(); assert_eq!(false, buf.try_reclaim(6)); buf.reserve(6); @@ -1306,9 +1306,12 @@ fn try_reserve_empty() { } #[test] -fn try_reserve_vec() { +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()); @@ -1322,11 +1325,14 @@ fn try_reserve_vec() { } #[test] -fn try_reserve_arc() { +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));