feat(body): add body::aggregate and body::to_bytes functions

Adds utility functions to `hyper::body` to help asynchronously
collecting all the buffers of some `HttpBody` into one.

- `aggregate` will collect all into an `impl Buf` without copying the
  contents. This is ideal if you don't need a contiguous buffer.
- `to_bytes` will copy all the data into a single contiguous `Bytes`
  buffer.
This commit is contained in:
Sean McArthur
2019-12-05 17:51:37 -08:00
parent 5a59875742
commit 8ba9a8d2c4
15 changed files with 282 additions and 128 deletions

25
src/body/aggregate.rs Normal file
View File

@@ -0,0 +1,25 @@
use bytes::Buf;
use super::HttpBody;
use crate::common::buf::BufList;
/// Aggregate the data buffers from a body asynchronously.
///
/// The returned `impl Buf` groups the `Buf`s from the `HttpBody` without
/// copying them. This is ideal if you don't require a contiguous buffer.
pub async fn aggregate<T>(body: T) -> Result<impl Buf, T::Error>
where
T: HttpBody,
{
let mut bufs = BufList::new();
futures_util::pin_mut!(body);
while let Some(buf) = body.data().await {
let buf = buf?;
if buf.has_remaining() {
bufs.push(buf);
}
}
Ok(bufs)
}

View File

@@ -18,11 +18,16 @@
pub use bytes::{Buf, Bytes};
pub use http_body::Body as HttpBody;
pub use self::aggregate::aggregate;
pub use self::body::{Body, Sender};
pub use self::to_bytes::to_bytes;
pub(crate) use self::payload::Payload;
mod aggregate;
mod body;
mod payload;
mod to_bytes;
/// An optimization to try to take a full body if immediately available.
///

36
src/body/to_bytes.rs Normal file
View File

@@ -0,0 +1,36 @@
use bytes::{Buf, BufMut, Bytes};
use super::HttpBody;
/// dox
pub async fn to_bytes<T>(body: T) -> Result<Bytes, T::Error>
where
T: HttpBody,
{
futures_util::pin_mut!(body);
// If there's only 1 chunk, we can just return Buf::to_bytes()
let mut first = if let Some(buf) = body.data().await {
buf?
} else {
return Ok(Bytes::new());
};
let second = if let Some(buf) = body.data().await {
buf?
} else {
return Ok(first.to_bytes());
};
// With more than 1 buf, we gotta flatten into a Vec first.
let cap = first.remaining() + second.remaining() + body.size_hint().lower() as usize;
let mut vec = Vec::with_capacity(cap);
vec.put(first);
vec.put(second);
while let Some(buf) = body.data().await {
vec.put(buf?);
}
Ok(vec.into())
}