Simplify vectored write API

Provide a write_vec_all() API so that users don't have to loop
themselves.

3 files changed, 41 insertions, 15 deletions

diff --git a/.bleep b/.bleep
index 0d35e95..917e9bf 100644
--- a/.bleep
+++ b/.bleep
@@ -1 +1 @@
-8812c147a92f6577ce6c0c979f058dd6e16842b0
\ No newline at end of file
+becf2775fceab572867708b00df1a23330acf187
\ No newline at end of file
diff --git a/pingora-core/src/protocols/http/v1/body.rs b/pingora-core/src/protocols/http/v1/body.rs
index c317bb7..0ddc90a 100644
--- a/pingora-core/src/protocols/http/v1/body.rs
+++ b/pingora-core/src/protocols/http/v1/body.rs
@@ -538,20 +538,10 @@ impl BodyWriter {
 
                 let chuck_size_buf = format!("{:X}\r\n", chunk_size);
                 let mut output_buf = Bytes::from(chuck_size_buf).chain(buf).chain(&b"\r\n"[..]);
-
-                while output_buf.has_remaining() {
-                    let res = stream.write_vec(&mut output_buf).await;
-                    match res {
-                        Ok(n) => {
-                            if n == 0 {
-                                return Error::e_explain(ConnectionClosed, "while writing body");
-                            }
-                        }
-                        Err(e) => {
-                            return Error::e_because(WriteError, "while writing body", e);
-                        }
-                    }
-                }
+                stream
+                    .write_vec_all(&mut output_buf)
+                    .await
+                    .or_err(WriteError, "while writing body")?;
                 stream.flush().await.or_err(WriteError, "flushing body")?;
                 self.body_mode = BM::ChunkedEncoding(written + chunk_size);
                 Ok(Some(chunk_size))
diff --git a/pingora-core/src/protocols/l4/stream.rs b/pingora-core/src/protocols/l4/stream.rs
index 7e32f2c..6a13c65 100644
--- a/pingora-core/src/protocols/l4/stream.rs
+++ b/pingora-core/src/protocols/l4/stream.rs
@@ -375,6 +375,12 @@ pub mod async_write_vec {
         buf: &'a mut B,
     }
 
+    #[must_use = "futures do nothing unless you `.await` or poll them"]
+    pub struct WriteVecAll<'a, W, B> {
+        writer: &'a mut W,
+        buf: &'a mut B,
+    }
+
     pub trait AsyncWriteVec {
         fn poll_write_vec<B: Buf>(
             self: Pin<&mut Self>,
@@ -392,6 +398,17 @@ pub mod async_write_vec {
                 buf: src,
             }
         }
+
+        fn write_vec_all<'a, B>(&'a mut self, src: &'a mut B) -> WriteVecAll<'a, Self, B>
+        where
+            Self: Sized,
+            B: Buf,
+        {
+            WriteVecAll {
+                writer: self,
+                buf: src,
+            }
+        }
     }
 
     impl<W, B> Future for WriteVec<'_, W, B>
@@ -407,6 +424,25 @@ pub mod async_write_vec {
         }
     }
 
+    impl<W, B> Future for WriteVecAll<'_, W, B>
+    where
+        W: AsyncWriteVec + Unpin,
+        B: Buf,
+    {
+        type Output = io::Result<()>;
+
+        fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<io::Result<()>> {
+            let me = &mut *self;
+            while me.buf.has_remaining() {
+                let n = ready!(Pin::new(&mut *me.writer).poll_write_vec(ctx, me.buf))?;
+                if n == 0 {
+                    return Poll::Ready(Err(io::ErrorKind::WriteZero.into()));
+                }
+            }
+            Poll::Ready(Ok(()))
+        }
+    }
+
     /* from https://github.com/tokio-rs/tokio/blob/master/tokio-util/src/lib.rs#L177 */
     impl<T> AsyncWriteVec for T
     where