Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

(de)compression: reduce memory allocation to improve performance #521

Merged
merged 1 commit into from
Sep 23, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
(de)compression: reduce memory allocation to improve performance
Currently, every time `WrapBody::poll_frame` is called, new instance of
`BytesMut` is created with the default capacity, which is effectively
64 bytes. This ends up with a lot of memory allocation in certain
situations, making the throughput significantly worse.

To optimize memory allocation, `WrapBody` now gets `BytesMut` as its
field, with initial capacity of 4096 bytes. This buffer will be reused
as much as possible across multiple `poll_frame` calls, and only when
its capacity becomes 0, new allocation of another 4096 bytes is
performed.

Fixes: #520
magurotuna committed Sep 22, 2024

Verified

This commit was signed with the committer’s verified signature.
eljamm Fedi Jamoussi
commit d84e6c0b0d1427edb366704b9a3bfd7e48f9d3d4
19 changes: 16 additions & 3 deletions tower-http/src/compression_utils.rs
Original file line number Diff line number Diff line change
@@ -141,10 +141,17 @@ pin_project! {
// rust-analyer thinks this field is private if its `pub(crate)` but works fine when its
// `pub`
pub read: M::Output,
// A buffer to temporarily store the data read from the underlying body.
// Reused as much as possible to optimize allocations.
buf: BytesMut,
read_all_data: bool,
}
}

impl<M: DecorateAsyncRead> WrapBody<M> {
const INTERNAL_BUF_CAPACITY: usize = 4096;
}

impl<M: DecorateAsyncRead> WrapBody<M> {
#[allow(dead_code)]
pub(crate) fn new<B>(body: B, quality: CompressionLevel) -> Self
@@ -167,6 +174,7 @@ impl<M: DecorateAsyncRead> WrapBody<M> {

Self {
read,
buf: BytesMut::with_capacity(Self::INTERNAL_BUF_CAPACITY),
read_all_data: false,
}
}
@@ -186,16 +194,21 @@ where
cx: &mut Context<'_>,
) -> Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
let mut this = self.project();
let mut buf = BytesMut::new();

if !*this.read_all_data {
let result = tokio_util::io::poll_read_buf(this.read.as_mut(), cx, &mut buf);
if this.buf.capacity() == 0 {
this.buf.reserve(Self::INTERNAL_BUF_CAPACITY);
}

let result = tokio_util::io::poll_read_buf(this.read.as_mut(), cx, &mut this.buf);

match ready!(result) {
Ok(0) => {
*this.read_all_data = true;
}
Ok(_) => {
return Poll::Ready(Some(Ok(Frame::data(buf.freeze()))));
let chunk = this.buf.split().freeze();
return Poll::Ready(Some(Ok(Frame::data(chunk))));
}
Err(err) => {
let body_error: Option<B::Error> = M::get_pin_mut(this.read)