Skip to content

Commit

Permalink
Avoid large reallocations when freezing BytesMut (#592)
Browse files Browse the repository at this point in the history
When we freeze a BytesMut, we turn it into a Vec, and then convert that
to a Bytes.  Currently, this happen using Vec::into_boxed_slice, which
reallocates to a slice of the same length as the Vev if the length and
the capacity are not equal.  This can pose a performance problem if the
Vec is large or if this happens many times in a loop.

Instead, let's compare the length and capacity, and if they're the same,
continue to handle this using into_boxed_slice.  Otherwise, since we
have a type of vtable which can handle a separate capacity, the shared
vtable, let's turn our Vec into that kind of Bytes.  While this does not
avoid allocation altogether, it performs a fixed size allocation and
avoids any need to memcpy.
  • Loading branch information
bk2204 authored Jan 31, 2023
1 parent f15bba3 commit 05e9d5c
Show file tree
Hide file tree
Showing 2 changed files with 75 additions and 2 deletions.
32 changes: 30 additions & 2 deletions src/bytes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -807,8 +807,36 @@ impl From<&'static str> for Bytes {

impl From<Vec<u8>> for Bytes {
fn from(vec: Vec<u8>) -> Bytes {
let slice = vec.into_boxed_slice();
slice.into()
let mut vec = vec;
let ptr = vec.as_mut_ptr();
let len = vec.len();
let cap = vec.capacity();

// Avoid an extra allocation if possible.
if len == cap {
return Bytes::from(vec.into_boxed_slice());
}

let shared = Box::new(Shared {
buf: ptr,
cap,
ref_cnt: AtomicUsize::new(1),
});
mem::forget(vec);

let shared = Box::into_raw(shared);
// The pointer should be aligned, so this assert should
// always succeed.
debug_assert!(
0 == (shared as usize & KIND_MASK),
"internal: Box<Shared> should have an aligned pointer",
);
Bytes {
ptr,
len,
data: AtomicPtr::new(shared as _),
vtable: &SHARED_VTABLE,
}
}
}

Expand Down
45 changes: 45 additions & 0 deletions tests/test_bytes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1163,3 +1163,48 @@ fn test_bytes_into_vec_promotable_even() {
assert_eq!(Vec::from(b2), vec[20..]);
assert_eq!(Vec::from(b1), vec[..20]);
}

#[test]
fn test_bytes_vec_conversion() {
let mut vec = Vec::with_capacity(10);
vec.extend(b"abcdefg");
let b = Bytes::from(vec);
let v = Vec::from(b);
assert_eq!(v.len(), 7);
assert_eq!(v.capacity(), 10);

let mut b = Bytes::from(v);
b.advance(1);
let v = Vec::from(b);
assert_eq!(v.len(), 6);
assert_eq!(v.capacity(), 10);
assert_eq!(v.as_slice(), b"bcdefg");
}

#[test]
fn test_bytes_mut_conversion() {
let mut b1 = BytesMut::with_capacity(10);
b1.extend(b"abcdefg");
let b2 = Bytes::from(b1);
let v = Vec::from(b2);
assert_eq!(v.len(), 7);
assert_eq!(v.capacity(), 10);

let mut b = Bytes::from(v);
b.advance(1);
let v = Vec::from(b);
assert_eq!(v.len(), 6);
assert_eq!(v.capacity(), 10);
assert_eq!(v.as_slice(), b"bcdefg");
}

#[test]
fn test_bytes_capacity_len() {
for cap in 0..100 {
for len in 0..=cap {
let mut v = Vec::with_capacity(cap);
v.resize(len, 0);
let _ = Bytes::from(v);
}
}
}

0 comments on commit 05e9d5c

Please sign in to comment.