Improve performance of Response::bytes() (#827)

This commit is contained in:
Sean McArthur
2020-02-27 12:44:04 -08:00
committed by GitHub
parent 41722a14fd
commit c916dc03cc
3 changed files with 317 additions and 279 deletions

View File

@@ -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),
} }
} }
}
} }

View File

@@ -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.

View File

@@ -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() {