From fac9987f529f97b92a75be4ed181907d1872e4f8 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sun, 19 Nov 2023 13:49:36 +0800 Subject: [PATCH] refactor(oio): Polish IncomingBody::bytes Signed-off-by: Xuanwo --- core/src/raw/http_util/body.rs | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/core/src/raw/http_util/body.rs b/core/src/raw/http_util/body.rs index fec20f9eec73..faaa3b2c4d53 100644 --- a/core/src/raw/http_util/body.rs +++ b/core/src/raw/http_util/body.rs @@ -105,12 +105,12 @@ impl IncomingAsyncBody { /// Consume the response to bytes. /// - /// This code is Inspired from hyper's [`to_bytes`](https://docs.rs/hyper/latest/hyper/body/fn.to_bytes.html). + /// This code is inspired from hyper's [`to_bytes`](https://docs.rs/hyper/0.14.23/hyper/body/fn.to_bytes.html). pub async fn bytes(mut self) -> Result { use oio::ReadExt; // If there's only 1 chunk, we can just return Buf::to_bytes() - let mut first = if let Some(buf) = self.next().await { + let first = if let Some(buf) = self.next().await { buf? } else { return Ok(Bytes::new()); @@ -119,11 +119,19 @@ impl IncomingAsyncBody { let second = if let Some(buf) = self.next().await { buf? } else { - return Ok(first.copy_to_bytes(first.remaining())); + return Ok(first); }; // With more than 1 buf, we gotta flatten into a Vec first. - let cap = first.remaining() + second.remaining() + self.size.unwrap_or_default() as usize; + let cap = if let Some(size) = self.size { + // The convert from u64 to usize could fail, but it's unlikely. + // Let's just make it overflow. + size as usize + } else { + // It's highly possible that we have more data to read. + // Add extra 16K buffer to avoid another allocation. + first.remaining() + second.remaining() + 16 * 1024 + }; let mut vec = Vec::with_capacity(cap); vec.put(first); vec.put(second);