Skip to content

Commit

Permalink
refactor(oio): Polish IncomingAsyncBody::bytes (#3621)
Browse files Browse the repository at this point in the history
  • Loading branch information
Xuanwo authored Nov 19, 2023
1 parent a53a7e5 commit 84dfbfb
Showing 1 changed file with 12 additions and 4 deletions.
16 changes: 12 additions & 4 deletions core/src/raw/http_util/body.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<Bytes> {
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());
Expand All @@ -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);
Expand Down

0 comments on commit 84dfbfb

Please sign in to comment.