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

src/async_impl/decoder.rs

@@ -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 {
-    use std::fmt;
-    use std::future::Future;
-    use std::pin::Pin;
-    use std::task::{Context, Poll};
+#[cfg(feature = "gzip")]
+use async_compression::stream::GzipDecoder;
 
-    #[cfg(feature = "gzip")]
-    use async_compression::stream::GzipDecoder;
+#[cfg(feature = "brotli")]
+use async_compression::stream::BrotliDecoder;
 
-    #[cfg(feature = "brotli")]
-    use async_compression::stream::BrotliDecoder;
+use bytes::Bytes;
+use futures_core::Stream;
+use futures_util::stream::Peekable;
+use http::HeaderMap;
+use hyper::body::HttpBody;
 
-    use bytes::Bytes;
-    use futures_core::Stream;
-    use futures_util::stream::Peekable;
-    use http::HeaderMap;
+use super::super::Body;
+use crate::error;
 
-    use super::super::Body;
-    use crate::error;
-
-    /// A response decompressor over a non-blocking stream of chunks.
-    ///
-    /// The inner decoder may be constructed asynchronously.
-    pub(crate) struct Decoder {
+/// A response decompressor over a non-blocking stream of chunks.
+///
+/// The inner decoder may be constructed asynchronously.
+pub(crate) struct Decoder {
     inner: Inner,
-    }
+}
 
-    enum DecoderType {
+enum DecoderType {
     #[cfg(feature = "gzip")]
     Gzip,
     #[cfg(feature = "brotli")]
     Brotli,
-    }
+}
 
-    enum Inner {
+enum Inner {
     /// 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.
     #[cfg(feature = "gzip")]
@@ -49,20 +47,20 @@ mod imp {
     /// A decoder that doesn't have a value yet.
     #[cfg(any(feature = "brotli", feature = "gzip"))]
     Pending(Pending),
-    }
+}
 
-    /// A future attempt to poll the response body for EOF so we know whether to use gzip or not.
-    struct Pending(Peekable<IoStream>, DecoderType);
+/// 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 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 {
         f.debug_struct("Decoder").finish()
     }
-    }
+}
 
-    impl Decoder {
+impl Decoder {
     #[cfg(feature = "blocking")]
     pub(crate) fn empty() -> Decoder {
         Decoder {
@@ -205,9 +203,9 @@ mod imp {
 
         Decoder::plain_text(body)
     }
-    }
+}
 
-    impl Stream for Decoder {
+impl Stream for Decoder {
     type Item = Result<Bytes, error::Error>;
 
     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>;
 
     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)))),
         }
     }
-    }
+}
 
-    impl Stream for IoStream {
+impl Stream for IoStream {
     type Item = Result<Bytes, std::io::Error>;
 
     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),
         }
     }
-    }
 }

src/async_impl/response.rs

@@ -2,12 +2,11 @@ use std::borrow::Cow;
 use std::fmt;
 use std::net::SocketAddr;
 
-use bytes::{Bytes, BytesMut};
+use bytes::Bytes;
 use encoding_rs::{Encoding, UTF_8};
 use futures_util::stream::StreamExt;
 use http;
 use hyper::client::connect::HttpInfo;
-use hyper::header::CONTENT_LENGTH;
 use hyper::{HeaderMap, StatusCode, Version};
 use mime::Mime;
 #[cfg(feature = "json")]
@@ -89,13 +88,12 @@ impl Response {
     /// Reasons it may not be known:
     ///
     /// - 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).
     pub fn content_length(&self) -> Option<u64> {
-        self.headers()
-            .get(CONTENT_LENGTH)
-            .and_then(|ct_len| ct_len.to_str().ok())
-            .and_then(|ct_len| ct_len.parse().ok())
+        use hyper::body::HttpBody;
+
+        HttpBody::size_hint(&self.body).exact()
     }
 
     /// Retrieve the cookies contained in the response.
@@ -259,12 +257,8 @@ impl Response {
     /// # Ok(())
     /// # }
     /// ```
-    pub async fn bytes(mut self) -> crate::Result<Bytes> {
-        let mut buf = BytesMut::new();
-        while let Some(chunk) = self.body.next().await {
-            buf.extend(chunk?);
-        }
-        Ok(buf.freeze())
+    pub async fn bytes(self) -> crate::Result<Bytes> {
+        hyper::body::to_bytes(self.body).await
     }
 
     /// Stream a chunk of the response body.

tests/client.rs

@@ -68,10 +68,29 @@ async fn response_text() {
         .send()
         .await
         .expect("Failed to get");
+    assert_eq!(res.content_length(), Some(5));
     let text = res.text().await.expect("Failed to get 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]
 #[cfg(feature = "json")]
 async fn response_json() {