Improve performance of Response::bytes() (#827)
This commit is contained in:
@@ -1,42 +1,40 @@
|
|||||||
pub(crate) use self::imp::Decoder;
|
use std::fmt;
|
||||||
|
use std::future::Future;
|
||||||
|
use std::pin::Pin;
|
||||||
|
use std::task::{Context, Poll};
|
||||||
|
|
||||||
mod imp {
|
#[cfg(feature = "gzip")]
|
||||||
use std::fmt;
|
use async_compression::stream::GzipDecoder;
|
||||||
use std::future::Future;
|
|
||||||
use std::pin::Pin;
|
|
||||||
use std::task::{Context, Poll};
|
|
||||||
|
|
||||||
#[cfg(feature = "gzip")]
|
#[cfg(feature = "brotli")]
|
||||||
use async_compression::stream::GzipDecoder;
|
use async_compression::stream::BrotliDecoder;
|
||||||
|
|
||||||
#[cfg(feature = "brotli")]
|
use bytes::Bytes;
|
||||||
use async_compression::stream::BrotliDecoder;
|
use futures_core::Stream;
|
||||||
|
use futures_util::stream::Peekable;
|
||||||
|
use http::HeaderMap;
|
||||||
|
use hyper::body::HttpBody;
|
||||||
|
|
||||||
use bytes::Bytes;
|
use super::super::Body;
|
||||||
use futures_core::Stream;
|
use crate::error;
|
||||||
use futures_util::stream::Peekable;
|
|
||||||
use http::HeaderMap;
|
|
||||||
|
|
||||||
use super::super::Body;
|
/// A response decompressor over a non-blocking stream of chunks.
|
||||||
use crate::error;
|
///
|
||||||
|
/// The inner decoder may be constructed asynchronously.
|
||||||
/// A response decompressor over a non-blocking stream of chunks.
|
pub(crate) struct Decoder {
|
||||||
///
|
|
||||||
/// The inner decoder may be constructed asynchronously.
|
|
||||||
pub(crate) struct Decoder {
|
|
||||||
inner: Inner,
|
inner: Inner,
|
||||||
}
|
}
|
||||||
|
|
||||||
enum DecoderType {
|
enum DecoderType {
|
||||||
#[cfg(feature = "gzip")]
|
#[cfg(feature = "gzip")]
|
||||||
Gzip,
|
Gzip,
|
||||||
#[cfg(feature = "brotli")]
|
#[cfg(feature = "brotli")]
|
||||||
Brotli,
|
Brotli,
|
||||||
}
|
}
|
||||||
|
|
||||||
enum Inner {
|
enum Inner {
|
||||||
/// A `PlainText` decoder just returns the response content as is.
|
/// A `PlainText` decoder just returns the response content as is.
|
||||||
PlainText(super::super::body::ImplStream),
|
PlainText(super::body::ImplStream),
|
||||||
|
|
||||||
/// A `Gzip` decoder will uncompress the gzipped response content before returning it.
|
/// A `Gzip` decoder will uncompress the gzipped response content before returning it.
|
||||||
#[cfg(feature = "gzip")]
|
#[cfg(feature = "gzip")]
|
||||||
@@ -49,20 +47,20 @@ mod imp {
|
|||||||
/// A decoder that doesn't have a value yet.
|
/// A decoder that doesn't have a value yet.
|
||||||
#[cfg(any(feature = "brotli", feature = "gzip"))]
|
#[cfg(any(feature = "brotli", feature = "gzip"))]
|
||||||
Pending(Pending),
|
Pending(Pending),
|
||||||
}
|
}
|
||||||
|
|
||||||
/// A future attempt to poll the response body for EOF so we know whether to use gzip or not.
|
/// A future attempt to poll the response body for EOF so we know whether to use gzip or not.
|
||||||
struct Pending(Peekable<IoStream>, DecoderType);
|
struct Pending(Peekable<IoStream>, DecoderType);
|
||||||
|
|
||||||
struct IoStream(super::super::body::ImplStream);
|
struct IoStream(super::body::ImplStream);
|
||||||
|
|
||||||
impl fmt::Debug for Decoder {
|
impl fmt::Debug for Decoder {
|
||||||
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
||||||
f.debug_struct("Decoder").finish()
|
f.debug_struct("Decoder").finish()
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Decoder {
|
impl Decoder {
|
||||||
#[cfg(feature = "blocking")]
|
#[cfg(feature = "blocking")]
|
||||||
pub(crate) fn empty() -> Decoder {
|
pub(crate) fn empty() -> Decoder {
|
||||||
Decoder {
|
Decoder {
|
||||||
@@ -205,9 +203,9 @@ mod imp {
|
|||||||
|
|
||||||
Decoder::plain_text(body)
|
Decoder::plain_text(body)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Stream for Decoder {
|
impl Stream for Decoder {
|
||||||
type Item = Result<Bytes, error::Error>;
|
type Item = Result<Bytes, error::Error>;
|
||||||
|
|
||||||
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
|
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
|
||||||
@@ -243,9 +241,37 @@ mod imp {
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl HttpBody for Decoder {
|
||||||
|
type Data = Bytes;
|
||||||
|
type Error = crate::Error;
|
||||||
|
|
||||||
|
fn poll_data(
|
||||||
|
self: Pin<&mut Self>,
|
||||||
|
cx: &mut Context,
|
||||||
|
) -> Poll<Option<Result<Self::Data, Self::Error>>> {
|
||||||
|
self.poll_next(cx)
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Future for Pending {
|
fn poll_trailers(
|
||||||
|
self: Pin<&mut Self>,
|
||||||
|
_cx: &mut Context,
|
||||||
|
) -> Poll<Result<Option<http::HeaderMap>, Self::Error>> {
|
||||||
|
Poll::Ready(Ok(None))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn size_hint(&self) -> http_body::SizeHint {
|
||||||
|
match self.inner {
|
||||||
|
Inner::PlainText(ref body) => HttpBody::size_hint(body),
|
||||||
|
// the rest are "unknown", so default
|
||||||
|
#[cfg(any(feature = "brotli", feature = "gzip"))]
|
||||||
|
_ => http_body::SizeHint::default(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Future for Pending {
|
||||||
type Output = Result<Inner, std::io::Error>;
|
type Output = Result<Inner, std::io::Error>;
|
||||||
|
|
||||||
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
|
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
|
||||||
@@ -278,9 +304,9 @@ mod imp {
|
|||||||
DecoderType::Gzip => Poll::Ready(Ok(Inner::Gzip(GzipDecoder::new(_body)))),
|
DecoderType::Gzip => Poll::Ready(Ok(Inner::Gzip(GzipDecoder::new(_body)))),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Stream for IoStream {
|
impl Stream for IoStream {
|
||||||
type Item = Result<Bytes, std::io::Error>;
|
type Item = Result<Bytes, std::io::Error>;
|
||||||
|
|
||||||
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
|
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
|
||||||
@@ -290,5 +316,4 @@ mod imp {
|
|||||||
None => Poll::Ready(None),
|
None => Poll::Ready(None),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -2,12 +2,11 @@ use std::borrow::Cow;
|
|||||||
use std::fmt;
|
use std::fmt;
|
||||||
use std::net::SocketAddr;
|
use std::net::SocketAddr;
|
||||||
|
|
||||||
use bytes::{Bytes, BytesMut};
|
use bytes::Bytes;
|
||||||
use encoding_rs::{Encoding, UTF_8};
|
use encoding_rs::{Encoding, UTF_8};
|
||||||
use futures_util::stream::StreamExt;
|
use futures_util::stream::StreamExt;
|
||||||
use http;
|
use http;
|
||||||
use hyper::client::connect::HttpInfo;
|
use hyper::client::connect::HttpInfo;
|
||||||
use hyper::header::CONTENT_LENGTH;
|
|
||||||
use hyper::{HeaderMap, StatusCode, Version};
|
use hyper::{HeaderMap, StatusCode, Version};
|
||||||
use mime::Mime;
|
use mime::Mime;
|
||||||
#[cfg(feature = "json")]
|
#[cfg(feature = "json")]
|
||||||
@@ -89,13 +88,12 @@ impl Response {
|
|||||||
/// Reasons it may not be known:
|
/// Reasons it may not be known:
|
||||||
///
|
///
|
||||||
/// - The server didn't send a `content-length` header.
|
/// - The server didn't send a `content-length` header.
|
||||||
/// - The response is gzipped and automatically decoded (thus changing
|
/// - The response is compressed and automatically decoded (thus changing
|
||||||
/// the actual decoded length).
|
/// the actual decoded length).
|
||||||
pub fn content_length(&self) -> Option<u64> {
|
pub fn content_length(&self) -> Option<u64> {
|
||||||
self.headers()
|
use hyper::body::HttpBody;
|
||||||
.get(CONTENT_LENGTH)
|
|
||||||
.and_then(|ct_len| ct_len.to_str().ok())
|
HttpBody::size_hint(&self.body).exact()
|
||||||
.and_then(|ct_len| ct_len.parse().ok())
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Retrieve the cookies contained in the response.
|
/// Retrieve the cookies contained in the response.
|
||||||
@@ -259,12 +257,8 @@ impl Response {
|
|||||||
/// # Ok(())
|
/// # Ok(())
|
||||||
/// # }
|
/// # }
|
||||||
/// ```
|
/// ```
|
||||||
pub async fn bytes(mut self) -> crate::Result<Bytes> {
|
pub async fn bytes(self) -> crate::Result<Bytes> {
|
||||||
let mut buf = BytesMut::new();
|
hyper::body::to_bytes(self.body).await
|
||||||
while let Some(chunk) = self.body.next().await {
|
|
||||||
buf.extend(chunk?);
|
|
||||||
}
|
|
||||||
Ok(buf.freeze())
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Stream a chunk of the response body.
|
/// Stream a chunk of the response body.
|
||||||
|
|||||||
@@ -68,10 +68,29 @@ async fn response_text() {
|
|||||||
.send()
|
.send()
|
||||||
.await
|
.await
|
||||||
.expect("Failed to get");
|
.expect("Failed to get");
|
||||||
|
assert_eq!(res.content_length(), Some(5));
|
||||||
let text = res.text().await.expect("Failed to get text");
|
let text = res.text().await.expect("Failed to get text");
|
||||||
assert_eq!("Hello", text);
|
assert_eq!("Hello", text);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn response_bytes() {
|
||||||
|
let _ = env_logger::try_init();
|
||||||
|
|
||||||
|
let server = server::http(move |_req| async { http::Response::new("Hello".into()) });
|
||||||
|
|
||||||
|
let client = Client::new();
|
||||||
|
|
||||||
|
let res = client
|
||||||
|
.get(&format!("http://{}/bytes", server.addr()))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.expect("Failed to get");
|
||||||
|
assert_eq!(res.content_length(), Some(5));
|
||||||
|
let bytes = res.bytes().await.expect("res.bytes()");
|
||||||
|
assert_eq!("Hello", bytes);
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
#[cfg(feature = "json")]
|
#[cfg(feature = "json")]
|
||||||
async fn response_json() {
|
async fn response_json() {
|
||||||
|
|||||||
Reference in New Issue
Block a user