Skip to content

Commit

Permalink
Fix reserve over allocating underlying buffer (#560)
Browse files Browse the repository at this point in the history
Fixes calls to `reserve` when the underlying shared buffer was already
big enough to fit the requested capacity. Previously a new even larger
buffer was created anyways. This could eventually lead to an OOM
condition.
  • Loading branch information
schultetwin1 authored Jul 30, 2022
1 parent 38fd42a commit d6e1999
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 1 deletion.
5 changes: 4 additions & 1 deletion src/bytes_mut.rs
Original file line number Diff line number Diff line change
Expand Up @@ -670,7 +670,10 @@ impl BytesMut {

// Compare the condition in the `kind == KIND_VEC` case above
// for more details.
if v_capacity >= new_cap && offset >= len {
if v_capacity >= new_cap + offset {
self.cap = new_cap;
// no copy is necessary
} else if v_capacity >= new_cap && offset >= len {
// The capacity is sufficient, and copying is not too much
// overhead: reclaim the buffer!

Expand Down
28 changes: 28 additions & 0 deletions tests/test_bytes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -515,6 +515,34 @@ fn reserve_in_arc_unique_doubles() {
assert_eq!(2000, bytes.capacity());
}

#[test]
fn reserve_in_arc_unique_does_not_overallocate_after_split() {
let mut bytes = BytesMut::from(LONG);
let orig_capacity = bytes.capacity();
drop(bytes.split_off(LONG.len() / 2));

// now bytes is Arc and refcount == 1

let new_capacity = bytes.capacity();
bytes.reserve(orig_capacity - new_capacity);
assert_eq!(bytes.capacity(), orig_capacity);
}

#[test]
fn reserve_in_arc_unique_does_not_overallocate_after_multiple_splits() {
let mut bytes = BytesMut::from(LONG);
let orig_capacity = bytes.capacity();
for _ in 0..10 {
drop(bytes.split_off(LONG.len() / 2));

// now bytes is Arc and refcount == 1

let new_capacity = bytes.capacity();
bytes.reserve(orig_capacity - new_capacity);
}
assert_eq!(bytes.capacity(), orig_capacity);
}

#[test]
fn reserve_in_arc_nonunique_does_not_overallocate() {
let mut bytes = BytesMut::with_capacity(1000);
Expand Down

0 comments on commit d6e1999

Please sign in to comment.