Skip to content

Commit 9fdf0eb

Browse files
authored
(de)compression: reduce memory allocation to improve performance (#521)
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
1 parent aeca262 commit 9fdf0eb

File tree

1 file changed

+16
-3
lines changed

1 file changed

+16
-3
lines changed

tower-http/src/compression_utils.rs

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -141,10 +141,17 @@ pin_project! {
141141
// rust-analyer thinks this field is private if its `pub(crate)` but works fine when its
142142
// `pub`
143143
pub read: M::Output,
144+
// A buffer to temporarily store the data read from the underlying body.
145+
// Reused as much as possible to optimize allocations.
146+
buf: BytesMut,
144147
read_all_data: bool,
145148
}
146149
}
147150

151+
impl<M: DecorateAsyncRead> WrapBody<M> {
152+
const INTERNAL_BUF_CAPACITY: usize = 4096;
153+
}
154+
148155
impl<M: DecorateAsyncRead> WrapBody<M> {
149156
#[allow(dead_code)]
150157
pub(crate) fn new<B>(body: B, quality: CompressionLevel) -> Self
@@ -167,6 +174,7 @@ impl<M: DecorateAsyncRead> WrapBody<M> {
167174

168175
Self {
169176
read,
177+
buf: BytesMut::with_capacity(Self::INTERNAL_BUF_CAPACITY),
170178
read_all_data: false,
171179
}
172180
}
@@ -186,16 +194,21 @@ where
186194
cx: &mut Context<'_>,
187195
) -> Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
188196
let mut this = self.project();
189-
let mut buf = BytesMut::new();
197+
190198
if !*this.read_all_data {
191-
let result = tokio_util::io::poll_read_buf(this.read.as_mut(), cx, &mut buf);
199+
if this.buf.capacity() == 0 {
200+
this.buf.reserve(Self::INTERNAL_BUF_CAPACITY);
201+
}
202+
203+
let result = tokio_util::io::poll_read_buf(this.read.as_mut(), cx, &mut this.buf);
192204

193205
match ready!(result) {
194206
Ok(0) => {
195207
*this.read_all_data = true;
196208
}
197209
Ok(_) => {
198-
return Poll::Ready(Some(Ok(Frame::data(buf.freeze()))));
210+
let chunk = this.buf.split().freeze();
211+
return Poll::Ready(Some(Ok(Frame::data(chunk))));
199212
}
200213
Err(err) => {
201214
let body_error: Option<B::Error> = M::get_pin_mut(this.read)

0 commit comments

Comments
 (0)