Simplify vectored write API

Provide a write_vec_all() API so that users don't have to loop
themselves.
This commit is contained in:
Yuchen Wu 2024-05-22 09:31:56 -07:00 committed by Yuchen Wu
parent ea1db2fb9d
commit 4019aa0819
3 changed files with 41 additions and 15 deletions

2
.bleep
View file

@ -1 +1 @@
8812c147a92f6577ce6c0c979f058dd6e16842b0
becf2775fceab572867708b00df1a23330acf187

View file

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

View file

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