From 14bb2f624f1ac3d440dd302219a48e9e8e926b50 Mon Sep 17 00:00:00 2001 From: Russell Cohen Date: Sat, 1 May 2021 14:54:38 -0400 Subject: [PATCH] io: add write_all_buf to AsyncWriteExt (#3737) --- tokio/src/io/util/async_write_ext.rs | 58 ++++++++++++++++- tokio/src/io/util/mod.rs | 1 + tokio/src/io/util/write_all_buf.rs | 56 ++++++++++++++++ tokio/tests/io_write_all_buf.rs | 96 ++++++++++++++++++++++++++++ 4 files changed, 209 insertions(+), 2 deletions(-) create mode 100644 tokio/src/io/util/write_all_buf.rs create mode 100644 tokio/tests/io_write_all_buf.rs diff --git a/tokio/src/io/util/async_write_ext.rs b/tokio/src/io/util/async_write_ext.rs index d011d82c1..2510ccdea 100644 --- a/tokio/src/io/util/async_write_ext.rs +++ b/tokio/src/io/util/async_write_ext.rs @@ -2,6 +2,7 @@ use crate::io::util::flush::{flush, Flush}; use crate::io::util::shutdown::{shutdown, Shutdown}; use crate::io::util::write::{write, Write}; use crate::io::util::write_all::{write_all, WriteAll}; +use crate::io::util::write_all_buf::{write_all_buf, WriteAllBuf}; use crate::io::util::write_buf::{write_buf, WriteBuf}; use crate::io::util::write_int::{ WriteI128, WriteI128Le, WriteI16, WriteI16Le, WriteI32, WriteI32Le, WriteI64, WriteI64Le, @@ -159,7 +160,6 @@ cfg_io_util! { write_vectored(self, bufs) } - /// Writes a buffer into this writer, advancing the buffer's internal /// cursor. /// @@ -197,10 +197,11 @@ cfg_io_util! { /// /// # Examples /// - /// [`File`] implements `Read` and [`Cursor<&[u8]>`] implements [`Buf`]: + /// [`File`] implements [`AsyncWrite`] and [`Cursor`]`<&[u8]>` implements [`Buf`]: /// /// [`File`]: crate::fs::File /// [`Buf`]: bytes::Buf + /// [`Cursor`]: std::io::Cursor /// /// ```no_run /// use tokio::io::{self, AsyncWriteExt}; @@ -233,6 +234,59 @@ cfg_io_util! { write_buf(self, src) } + /// Attempts to write an entire buffer into this writer + /// + /// Equivalent to: + /// + /// ```ignore + /// async fn write_all_buf(&mut self, buf: impl Buf) -> Result<(), io::Error> { + /// while buf.has_remaining() { + /// self.write_buf(&mut buf).await?; + /// } + /// } + /// ``` + /// + /// This method will continuously call [`write`] until + /// [`buf.has_remaining()`](bytes::Buf::has_remaining) returns false. This method will not + /// return until the entire buffer has been successfully written or an error occurs. The + /// first error generated will be returned. + /// + /// The buffer is advanced after each chunk is successfully written. After failure, + /// `src.chunk()` will return the chunk that failed to write. + /// + /// # Examples + /// + /// [`File`] implements [`AsyncWrite`] and [`Cursor`]`<&[u8]>` implements [`Buf`]: + /// + /// [`File`]: crate::fs::File + /// [`Buf`]: bytes::Buf + /// [`Cursor`]: std::io::Cursor + /// + /// ```no_run + /// use tokio::io::{self, AsyncWriteExt}; + /// use tokio::fs::File; + /// + /// use std::io::Cursor; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut file = File::create("foo.txt").await?; + /// let mut buffer = Cursor::new(b"data to write"); + /// + /// file.write_all_buf(&mut buffer).await?; + /// Ok(()) + /// } + /// ``` + /// + /// [`write`]: AsyncWriteExt::write + fn write_all_buf<'a, B>(&'a mut self, src: &'a mut B) -> WriteAllBuf<'a, Self, B> + where + Self: Sized + Unpin, + B: Buf, + { + write_all_buf(self, src) + } + /// Attempts to write an entire buffer into this writer. /// /// Equivalent to: diff --git a/tokio/src/io/util/mod.rs b/tokio/src/io/util/mod.rs index ab3866437..fd3dd0dfc 100644 --- a/tokio/src/io/util/mod.rs +++ b/tokio/src/io/util/mod.rs @@ -77,6 +77,7 @@ cfg_io_util! { mod write_vectored; mod write_all; mod write_buf; + mod write_all_buf; mod write_int; diff --git a/tokio/src/io/util/write_all_buf.rs b/tokio/src/io/util/write_all_buf.rs new file mode 100644 index 000000000..05af7fe99 --- /dev/null +++ b/tokio/src/io/util/write_all_buf.rs @@ -0,0 +1,56 @@ +use crate::io::AsyncWrite; + +use bytes::Buf; +use pin_project_lite::pin_project; +use std::future::Future; +use std::io; +use std::marker::PhantomPinned; +use std::pin::Pin; +use std::task::{Context, Poll}; + +pin_project! { + /// A future to write some of the buffer to an `AsyncWrite`. + #[derive(Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct WriteAllBuf<'a, W, B> { + writer: &'a mut W, + buf: &'a mut B, + #[pin] + _pin: PhantomPinned, + } +} + +/// Tries to write some bytes from the given `buf` to the writer in an +/// asynchronous manner, returning a future. +pub(crate) fn write_all_buf<'a, W, B>(writer: &'a mut W, buf: &'a mut B) -> WriteAllBuf<'a, W, B> +where + W: AsyncWrite + Unpin, + B: Buf, +{ + WriteAllBuf { + writer, + buf, + _pin: PhantomPinned, + } +} + +impl Future for WriteAllBuf<'_, W, B> +where + W: AsyncWrite + Unpin, + B: Buf, +{ + type Output = io::Result<()>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let me = self.project(); + while me.buf.has_remaining() { + let n = ready!(Pin::new(&mut *me.writer).poll_write(cx, me.buf.chunk())?); + me.buf.advance(n); + if n == 0 { + return Poll::Ready(Err(io::ErrorKind::WriteZero.into())); + } + } + + Poll::Ready(Ok(())) + } +} diff --git a/tokio/tests/io_write_all_buf.rs b/tokio/tests/io_write_all_buf.rs new file mode 100644 index 000000000..b49a58ebe --- /dev/null +++ b/tokio/tests/io_write_all_buf.rs @@ -0,0 +1,96 @@ +#![warn(rust_2018_idioms)] +#![cfg(feature = "full")] + +use tokio::io::{AsyncWrite, AsyncWriteExt}; +use tokio_test::{assert_err, assert_ok}; + +use bytes::{Buf, Bytes, BytesMut}; +use std::cmp; +use std::io; +use std::pin::Pin; +use std::task::{Context, Poll}; + +#[tokio::test] +async fn write_all_buf() { + struct Wr { + buf: BytesMut, + cnt: usize, + } + + impl AsyncWrite for Wr { + fn poll_write( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let n = cmp::min(4, buf.len()); + dbg!(buf); + let buf = &buf[0..n]; + + self.cnt += 1; + self.buf.extend(buf); + Ok(buf.len()).into() + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Ok(()).into() + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Ok(()).into() + } + } + + let mut wr = Wr { + buf: BytesMut::with_capacity(64), + cnt: 0, + }; + + let mut buf = Bytes::from_static(b"hello").chain(Bytes::from_static(b"world")); + + assert_ok!(wr.write_all_buf(&mut buf).await); + assert_eq!(wr.buf, b"helloworld"[..]); + // expect 4 writes, [hell],[o],[worl],[d] + assert_eq!(wr.cnt, 4); + assert_eq!(buf.has_remaining(), false); +} + +#[tokio::test] +async fn write_buf_err() { + /// Error out after writing the first 4 bytes + struct Wr { + cnt: usize, + } + + impl AsyncWrite for Wr { + fn poll_write( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + _buf: &[u8], + ) -> Poll> { + self.cnt += 1; + if self.cnt == 2 { + return Poll::Ready(Err(io::Error::new(io::ErrorKind::Other, "whoops"))); + } + Poll::Ready(Ok(4)) + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Ok(()).into() + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Ok(()).into() + } + } + + let mut wr = Wr { cnt: 0 }; + + let mut buf = Bytes::from_static(b"hello").chain(Bytes::from_static(b"world")); + + assert_err!(wr.write_all_buf(&mut buf).await); + assert_eq!( + buf.copy_to_bytes(buf.remaining()), + Bytes::from_static(b"oworld") + ); +}