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.

src/body/aggregate.rs

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

src/body/mod.rs

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

src/body/to_bytes.rs

@@ -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())
+}

src/common/buf.rs

@@ -0,0 +1,75 @@
+use std::collections::VecDeque;
+use std::io::IoSlice;
+
+use bytes::Buf;
+
+pub(crate) struct BufList<T> {
+    bufs: VecDeque<T>,
+}
+
+impl<T: Buf> BufList<T> {
+    pub(crate) fn new() -> BufList<T> {
+        BufList {
+            bufs: VecDeque::new(),
+        }
+    }
+
+    #[inline]
+    pub(crate) fn push(&mut self, buf: T) {
+        debug_assert!(buf.has_remaining());
+        self.bufs.push_back(buf);
+    }
+
+    #[inline]
+    pub(crate) fn bufs_cnt(&self) -> usize {
+        self.bufs.len()
+    }
+}
+
+impl<T: Buf> Buf for BufList<T> {
+    #[inline]
+    fn remaining(&self) -> usize {
+        self.bufs.iter().map(|buf| buf.remaining()).sum()
+    }
+
+    #[inline]
+    fn bytes(&self) -> &[u8] {
+        for buf in &self.bufs {
+            return buf.bytes();
+        }
+        &[]
+    }
+
+    #[inline]
+    fn advance(&mut self, mut cnt: usize) {
+        while cnt > 0 {
+            {
+                let front = &mut self.bufs[0];
+                let rem = front.remaining();
+                if rem > cnt {
+                    front.advance(cnt);
+                    return;
+                } else {
+                    front.advance(rem);
+                    cnt -= rem;
+                }
+            }
+            self.bufs.pop_front();
+        }
+    }
+
+    #[inline]
+    fn bytes_vectored<'t>(&'t self, dst: &mut [IoSlice<'t>]) -> usize {
+        if dst.is_empty() {
+            return 0;
+        }
+        let mut vecs = 0;
+        for buf in &self.bufs {
+            vecs += buf.bytes_vectored(&mut dst[vecs..]);
+            if vecs == dst.len() {
+                break;
+            }
+        }
+        vecs
+    }
+}

src/common/mod.rs

@@ -7,6 +7,7 @@ macro_rules! ready {
     };
 }
 
+pub(crate) mod buf;
 pub(crate) mod drain;
 pub(crate) mod exec;
 pub(crate) mod io;

src/proto/h1/io.rs

@@ -1,6 +1,5 @@
 use std::cell::Cell;
 use std::cmp;
-use std::collections::VecDeque;
 use std::fmt;
 use std::io::{self, IoSlice};
 
@@ -8,6 +7,7 @@ use bytes::{Buf, BufMut, Bytes, BytesMut};
 use tokio::io::{AsyncRead, AsyncWrite};
 
 use super::{Http1Transaction, ParseContext, ParsedMessage};
+use crate::common::buf::BufList;
 use crate::common::{task, Pin, Poll, Unpin};
 
 /// The initial buffer size allocated before trying to read from IO.
@@ -90,7 +90,7 @@ where
     pub fn set_write_strategy_flatten(&mut self) {
         // this should always be called only at construction time,
         // so this assert is here to catch myself
-        debug_assert!(self.write_buf.queue.bufs.is_empty());
+        debug_assert!(self.write_buf.queue.bufs_cnt() == 0);
         self.write_buf.set_strategy(WriteStrategy::Flatten);
     }
 
@@ -431,16 +431,16 @@ pub(super) struct WriteBuf<B> {
     headers: Cursor<Vec<u8>>,
     max_buf_size: usize,
     /// Deque of user buffers if strategy is Queue
-    queue: BufDeque<B>,
+    queue: BufList<B>,
     strategy: WriteStrategy,
 }
 
-impl<B> WriteBuf<B> {
+impl<B: Buf> WriteBuf<B> {
     fn new() -> WriteBuf<B> {
         WriteBuf {
             headers: Cursor::new(Vec::with_capacity(INIT_BUFFER_SIZE)),
             max_buf_size: DEFAULT_MAX_BUFFER_SIZE,
-            queue: BufDeque::new(),
+            queue: BufList::new(),
             strategy: WriteStrategy::Auto,
         }
     }
@@ -479,7 +479,7 @@ where
                 }
             }
             WriteStrategy::Auto | WriteStrategy::Queue => {
-                self.queue.bufs.push_back(buf.into());
+                self.queue.push(buf.into());
             }
         }
     }
@@ -488,7 +488,7 @@ where
         match self.strategy {
             WriteStrategy::Flatten => self.remaining() < self.max_buf_size,
             WriteStrategy::Auto | WriteStrategy::Queue => {
-                self.queue.bufs.len() < MAX_BUF_LIST_BUFFERS && self.remaining() < self.max_buf_size
+                self.queue.bufs_cnt() < MAX_BUF_LIST_BUFFERS && self.remaining() < self.max_buf_size
             }
         }
     }
@@ -608,66 +608,6 @@ enum WriteStrategy {
     Queue,
 }
 
-struct BufDeque<T> {
-    bufs: VecDeque<T>,
-}
-
-impl<T> BufDeque<T> {
-    fn new() -> BufDeque<T> {
-        BufDeque {
-            bufs: VecDeque::new(),
-        }
-    }
-}
-
-impl<T: Buf> Buf for BufDeque<T> {
-    #[inline]
-    fn remaining(&self) -> usize {
-        self.bufs.iter().map(|buf| buf.remaining()).sum()
-    }
-
-    #[inline]
-    fn bytes(&self) -> &[u8] {
-        for buf in &self.bufs {
-            return buf.bytes();
-        }
-        &[]
-    }
-
-    #[inline]
-    fn advance(&mut self, mut cnt: usize) {
-        while cnt > 0 {
-            {
-                let front = &mut self.bufs[0];
-                let rem = front.remaining();
-                if rem > cnt {
-                    front.advance(cnt);
-                    return;
-                } else {
-                    front.advance(rem);
-                    cnt -= rem;
-                }
-            }
-            self.bufs.pop_front();
-        }
-    }
-
-    #[inline]
-    fn bytes_vectored<'t>(&'t self, dst: &mut [IoSlice<'t>]) -> usize {
-        if dst.is_empty() {
-            return 0;
-        }
-        let mut vecs = 0;
-        for buf in &self.bufs {
-            vecs += buf.bytes_vectored(&mut dst[vecs..]);
-            if vecs == dst.len() {
-                break;
-            }
-        }
-        vecs
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -871,12 +811,12 @@ mod tests {
         buffered.buffer(Cursor::new(b"world, ".to_vec()));
         buffered.buffer(Cursor::new(b"it's ".to_vec()));
         buffered.buffer(Cursor::new(b"hyper!".to_vec()));
-        assert_eq!(buffered.write_buf.queue.bufs.len(), 3);
+        assert_eq!(buffered.write_buf.queue.bufs_cnt(), 3);
         buffered.flush().unwrap();
 
         assert_eq!(buffered.io, b"hello world, it's hyper!");
         assert_eq!(buffered.io.num_writes(), 1);
-        assert_eq!(buffered.write_buf.queue.bufs.len(), 0);
+        assert_eq!(buffered.write_buf.queue.bufs_cnt(), 0);
     }
     */
 
@@ -896,7 +836,7 @@ mod tests {
         buffered.buffer(Cursor::new(b"world, ".to_vec()));
         buffered.buffer(Cursor::new(b"it's ".to_vec()));
         buffered.buffer(Cursor::new(b"hyper!".to_vec()));
-        assert_eq!(buffered.write_buf.queue.bufs.len(), 0);
+        assert_eq!(buffered.write_buf.queue.bufs_cnt(), 0);
 
         buffered.flush().await.expect("flush");
     }
@@ -921,11 +861,11 @@ mod tests {
         buffered.buffer(Cursor::new(b"world, ".to_vec()));
         buffered.buffer(Cursor::new(b"it's ".to_vec()));
         buffered.buffer(Cursor::new(b"hyper!".to_vec()));
-        assert_eq!(buffered.write_buf.queue.bufs.len(), 3);
+        assert_eq!(buffered.write_buf.queue.bufs_cnt(), 3);
 
         buffered.flush().await.expect("flush");
 
-        assert_eq!(buffered.write_buf.queue.bufs.len(), 0);
+        assert_eq!(buffered.write_buf.queue.bufs_cnt(), 0);
     }
 
     #[tokio::test]
@@ -949,11 +889,11 @@ mod tests {
         buffered.buffer(Cursor::new(b"world, ".to_vec()));
         buffered.buffer(Cursor::new(b"it's ".to_vec()));
         buffered.buffer(Cursor::new(b"hyper!".to_vec()));
-        assert_eq!(buffered.write_buf.queue.bufs.len(), 3);
+        assert_eq!(buffered.write_buf.queue.bufs_cnt(), 3);
 
         buffered.flush().await.expect("flush");
 
-        assert_eq!(buffered.write_buf.queue.bufs.len(), 0);
+        assert_eq!(buffered.write_buf.queue.bufs_cnt(), 0);
     }
 
     #[cfg(feature = "nightly")]