mirror of
https://github.com/cloudflare/pingora.git
synced 2024-09-20 02:31:35 +02:00
Simplify vectored write API
Provide a write_vec_all() API so that users don't have to loop themselves.
This commit is contained in:
parent
ea1db2fb9d
commit
4019aa0819
3 changed files with 41 additions and 15 deletions
2
.bleep
2
.bleep
|
@ -1 +1 @@
|
|||
8812c147a92f6577ce6c0c979f058dd6e16842b0
|
||||
becf2775fceab572867708b00df1a23330acf187
|
|
@ -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))
|
||||
|
|
|
@ -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
|
||||
|
|
Loading…
Reference in a new issue