Skip to content

Commit

Permalink
Allow reclaiming the current allocation
Browse files Browse the repository at this point in the history
This is based on #680, where it was noted that it is hard to use
BytesMut without additional allocations in some circumstances.
  • Loading branch information
shahn committed Mar 31, 2024
1 parent ce8d8a0 commit 3d42904
Show file tree
Hide file tree
Showing 2 changed files with 115 additions and 0 deletions.
68 changes: 68 additions & 0 deletions src/bytes_mut.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ use alloc::{
vec,
vec::Vec,
};
use std::num::NonZeroUsize;

use crate::buf::{IntoIter, UninitSlice};
use crate::bytes::Vtable;
Expand Down Expand Up @@ -819,6 +820,73 @@ impl BytesMut {
}
}

/// Attempts to reclaim the whole allocation of the `BytesMut`.
///
/// If the `BytesMut` is empty but the underlying storage has been used before,
/// it might be possible to cheaply reclaim space by just updating a few indices.
/// Returns `None` if the `BytesMut` is not empty, there is nothing to reclaim
/// (no underlying storage has been allocated), or there are any other live
/// references to the underlying storage. Otherwise, returns the available
/// capacity after reclaiming.
///
/// # Examples
///
/// ```
/// use bytes::BytesMut;
/// use core::num::NonZeroUsize;
///
/// let mut buf = BytesMut::with_capacity(64);
/// assert_eq!(None, buf.try_reclaim());
///
/// buf.extend_from_slice(b"abcd");
/// let mut split = buf.split();
/// assert_eq!(None, split.try_reclaim());
/// assert_eq!(None, buf.try_reclaim());
/// drop(buf);
/// assert_eq!(None, split.try_reclaim());
/// split.clear();
/// assert_eq!(Some(64), split.try_reclaim().map(NonZeroUsize::into));
/// ```
pub fn try_reclaim(&mut self) -> Option<NonZeroUsize> {
if !self.is_empty() {
return None;
}

let kind = self.kind();
if kind == KIND_VEC {
unsafe {
let off = self.get_vec_pos();
if off == 0 {
return None;
}

let base_ptr = self.ptr.as_ptr().sub(off);
self.ptr = vptr(base_ptr);
self.set_vec_pos(0);
self.cap += off;
debug_assert!(self.capacity() > 0);
return Some(NonZeroUsize::new_unchecked(self.capacity()));
}
}
let shared: *mut Shared = self.data;

unsafe {
if !(*shared).is_unique() {
return None;
}
let v = &mut (*shared).vec;
let cap = v.capacity();
if cap == 0 {
return None;
}

let ptr = v.as_mut_ptr();
self.ptr = vptr(ptr);
self.cap = cap;
Some(NonZeroUsize::new_unchecked(cap))
}
}

// private

// For now, use a `Vec` to manage the memory for us, but we may want to
Expand Down
47 changes: 47 additions & 0 deletions tests/test_bytes.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
#![warn(rust_2018_idioms)]

use bytes::{Buf, BufMut, Bytes, BytesMut};
use std::num::NonZeroUsize;

use std::usize;

Expand Down Expand Up @@ -1172,3 +1173,49 @@ fn shared_is_unique() {
drop(b);
assert!(c.is_unique());
}

#[test]
fn try_reclaim_empty() {
let mut buf = BytesMut::new();
assert_eq!(None, buf.try_reclaim());
buf.reserve(6);
assert_eq!(None, buf.try_reclaim());

let mut buf = BytesMut::new();
let mut split = buf.split();
drop(buf);
assert_eq!(0, split.capacity());
assert_eq!(None, split.try_reclaim());
}

#[test]
fn try_reclaim_vec() {
let mut buf = BytesMut::with_capacity(6);
buf.put_slice(b"abc");
buf.advance(3);
assert_eq!(3, buf.capacity());
assert_eq!(Some(NonZeroUsize::new(6).unwrap()), buf.try_reclaim());
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");
let y = buf.split().freeze();
let z = y.clone();
assert_eq!(None, buf.try_reclaim());
drop(x);
drop(z);
assert_eq!(None, buf.try_reclaim());
drop(y);
assert_eq!(Some(NonZeroUsize::new(6).unwrap()), buf.try_reclaim());
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());
}

0 comments on commit 3d42904

Please sign in to comment.