util: add writev-aware poll_write_buf (#3156)

## Motivation

In Tokio 0.2, `AsyncRead` and `AsyncWrite` had `poll_write_buf` and
`poll_read_buf` methods for reading and writing to implementers of
`bytes` `Buf` and `BufMut` traits. In 0.3, these were removed, but
`poll_read_buf` was added as a free function in `tokio-util`. However,
there is currently no `poll_write_buf`.

Now that `AsyncWrite` has regained support for vectored writes in #3149,
there's a lot of potential benefit in having a `poll_write_buf` that
uses vectored writes when supported and non-vectored writes when not
supported, so that users don't have to reimplement this.

## Solution

This PR adds a `poll_write_buf` function to `tokio_util::io`, analogous
to the existing `poll_read_buf` function.

This function writes from a `Buf` to an `AsyncWrite`, advancing the
`Buf`'s internal cursor. In addition, when the `AsyncWrite` supports
vectored writes (i.e. its `is_write_vectored` method returns `true`),
it will use vectored IO.

I copied the documentation for this functions from the docs from Tokio
0.2's `AsyncWrite::poll_write_buf` , with some minor modifications as
appropriate.

Finally, I fixed a minor issue in the existing docs for `poll_read_buf`
and `read_buf`, and updated `tokio_util::codec` to use `poll_write_buf`.

Signed-off-by: Eliza Weisman <[email protected]>
This commit is contained in:
Eliza Weisman
2020-12-03 11:19:16 -08:00
committed by GitHub
parent a6051a61ec
commit 647299866a
5 changed files with 77 additions and 14 deletions
+3 -5
View File
@@ -6,7 +6,7 @@ use tokio::{
stream::Stream,
};
use bytes::{Buf, BytesMut};
use bytes::BytesMut;
use futures_core::ready;
use futures_sink::Sink;
use log::trace;
@@ -189,6 +189,7 @@ where
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
use crate::util::poll_write_buf;
trace!("flushing framed transport");
let mut pinned = self.project();
@@ -196,8 +197,7 @@ where
let WriteFrame { buffer } = pinned.state.borrow_mut();
trace!("writing; remaining={}", buffer.len());
let buf = &buffer;
let n = ready!(pinned.inner.as_mut().poll_write(cx, &buf))?;
let n = ready!(poll_write_buf(pinned.inner.as_mut(), cx, buffer))?;
if n == 0 {
return Poll::Ready(Err(io::Error::new(
@@ -207,8 +207,6 @@ where
)
.into()));
}
pinned.state.borrow_mut().buffer.advance(n);
}
// Try flushing the underlying IO