Skip to content

Commit

Permalink
Address review: Add safety comments, update docs
Browse files Browse the repository at this point in the history
  • Loading branch information
shahn committed May 13, 2024
1 parent 3fa520b commit b9c8170
Showing 1 changed file with 58 additions and 35 deletions.
93 changes: 58 additions & 35 deletions src/bytes_mut.rs
Original file line number Diff line number Diff line change
Expand Up @@ -827,15 +827,18 @@ impl BytesMut {
}
}

/// Attempts to reclaim the whole allocation of the `BytesMut`.
/// Attempts to cheaply 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 capacity as long as there are no
/// other `Bytes` or `BytesMut` handles alive.
/// If no other `Bytes` or `BytesMut` handles to the underlying storage of this
/// `BytesMut` are alive, it might be possible to cheaply reclaim the full capacity
/// of the allocation.
///
/// Returns `true` if the allocation was reclaimed (or reclaiming was not
/// necessary, because the underlying storage was empty); returns `false` if the
/// `BytesMut` is not empty or there are other outstanding references to the
/// `try_reclaim` never copies bytes internally, therefore it always returns `false`
/// if the `BytesMut` is not empty.
///
/// Returns `true` if the allocation was reclaimed (or reclaiming was not necessary,
/// because the underlying storage was already fully available); returns `false` if the
/// `BytesMut` is not empty or there are other outstanding handles to the
/// underlying storage.
///
/// # Examples
Expand All @@ -845,15 +848,22 @@ impl BytesMut {
///
/// let mut buf = BytesMut::with_capacity(64);
/// assert_eq!(true, buf.try_reclaim());
/// 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());
/// assert_eq!(false, buf.try_reclaim());
///
/// drop(buf);
/// assert_eq!(false, split.try_reclaim());
///
/// split.clear();
/// assert_eq!(4, split.capacity());
/// assert_eq!(true, split.try_reclaim());
/// assert_eq!(64, split.capacity());
/// ```
#[must_use = "consider BytesMut::reserve if you need a specific capacity"]
pub fn try_reclaim(&mut self) -> bool {
Expand All @@ -863,41 +873,54 @@ impl BytesMut {

let kind = self.kind();
if kind == KIND_VEC {
unsafe {
// Short-circuit if we are already at the start of the vector.
let off = self.get_vec_pos();
if off == 0 {
return true;
}
// Otherwise, update info to point at the front of the vector
// again and reclaim capacity
let base_ptr = self.ptr.as_ptr().sub(off);
self.ptr = vptr(base_ptr);
self.set_vec_pos(0);
self.cap += off;
// Safety: self is of KIND_VEC, so calling get_vec_pos() is safe.
let off = unsafe { self.get_vec_pos() };

// Short-circuit if we are already at the start of the vector.
if off == 0 {
return true;
}
// 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;

unsafe {
// If there are other references to the underlying storage, we cannot reclaim
if !(*shared).is_unique() {
return false;
}
let v = &mut (*shared).vec;
let cap = v.capacity();
if cap == 0 {
// Short-circuit as an empty vector is trivially reclaimed already
return true;
}
// If there are other references to the underlying storage, we cannot reclaim
//
// Safety: self is of type KIND_ARC, so the Shared ptr is valid
if unsafe { !(*shared).is_unique() } {
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
// Safety: self is of type KIND_ARC and there are no other handles alive, so we
// can get a mut reference to the underlying storage
let v = unsafe { &mut (*shared).vec };
let cap = v.capacity();
if cap == 0 {
// Short-circuit as an empty vector is trivially reclaimed already
return true;
}

// 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
}

// private
Expand Down

0 comments on commit b9c8170

Please sign in to comment.