diff --git a/tower-http/src/compression_utils.rs b/tower-http/src/compression_utils.rs index 153ae324..58aa6a08 100644 --- a/tower-http/src/compression_utils.rs +++ b/tower-http/src/compression_utils.rs @@ -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 WrapBody { + const INTERNAL_BUF_CAPACITY: usize = 4096; +} + impl WrapBody { #[allow(dead_code)] pub(crate) fn new(body: B, quality: CompressionLevel) -> Self @@ -167,6 +174,7 @@ impl WrapBody { Self { read, + buf: BytesMut::with_capacity(Self::INTERNAL_BUF_CAPACITY), read_all_data: false, } } @@ -186,16 +194,21 @@ where cx: &mut Context<'_>, ) -> Poll, 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 = M::get_pin_mut(this.read)