From eabb7ce61c2a38fe3bc95cba2c758b5028d6e5e9 Mon Sep 17 00:00:00 2001
From: David Pedersen <david.pdrsn@gmail.com>
Date: Tue, 6 Apr 2021 00:02:00 +0200
Subject: [PATCH] io: add `AsyncWriteExt::write_vectored` (#3678)

Fixes #3676
---
 tokio/src/io/util/async_write_ext.rs | 43 +++++++++++++++++++++++++
 tokio/src/io/util/mod.rs             |  1 +
 tokio/src/io/util/write_vectored.rs  | 47 ++++++++++++++++++++++++++++
 3 files changed, 91 insertions(+)
 create mode 100644 tokio/src/io/util/write_vectored.rs

diff --git a/tokio/src/io/util/async_write_ext.rs b/tokio/src/io/util/async_write_ext.rs
index 978062eb2..d011d82c1 100644
--- a/tokio/src/io/util/async_write_ext.rs
+++ b/tokio/src/io/util/async_write_ext.rs
@@ -11,7 +11,9 @@ use crate::io::util::write_int::{
     WriteU128, WriteU128Le, WriteU16, WriteU16Le, WriteU32, WriteU32Le, WriteU64, WriteU64Le,
     WriteU8,
 };
+use crate::io::util::write_vectored::{write_vectored, WriteVectored};
 use crate::io::AsyncWrite;
+use std::io::IoSlice;
 
 use bytes::Buf;
 
@@ -116,6 +118,47 @@ cfg_io_util! {
             write(self, src)
         }
 
+        /// Like [`write`], except that it writes from a slice of buffers.
+        ///
+        /// Equivalent to:
+        ///
+        /// ```ignore
+        /// async fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize>;
+        /// ```
+        ///
+        /// See [`AsyncWrite::poll_write_vectored`] for more details.
+        ///
+        /// # Examples
+        ///
+        /// ```no_run
+        /// use tokio::io::{self, AsyncWriteExt};
+        /// use tokio::fs::File;
+        /// use std::io::IoSlice;
+        ///
+        /// #[tokio::main]
+        /// async fn main() -> io::Result<()> {
+        ///     let mut file = File::create("foo.txt").await?;
+        ///
+        ///     let bufs: &[_] = &[
+        ///         IoSlice::new(b"hello"),
+        ///         IoSlice::new(b" "),
+        ///         IoSlice::new(b"world"),
+        ///     ];
+        ///
+        ///     file.write_vectored(&bufs).await?;
+        ///
+        ///     Ok(())
+        /// }
+        /// ```
+        ///
+        /// [`write`]: AsyncWriteExt::write
+        fn write_vectored<'a, 'b>(&'a mut self, bufs: &'a [IoSlice<'b>]) -> WriteVectored<'a, 'b, Self>
+        where
+            Self: Unpin,
+        {
+            write_vectored(self, bufs)
+        }
+
 
         /// Writes a buffer into this writer, advancing the buffer's internal
         /// cursor.
diff --git a/tokio/src/io/util/mod.rs b/tokio/src/io/util/mod.rs
index 9ddb7758d..3ac897311 100644
--- a/tokio/src/io/util/mod.rs
+++ b/tokio/src/io/util/mod.rs
@@ -71,6 +71,7 @@ cfg_io_util! {
     pub use take::Take;
 
     mod write;
+    mod write_vectored;
     mod write_all;
     mod write_buf;
     mod write_int;
diff --git a/tokio/src/io/util/write_vectored.rs b/tokio/src/io/util/write_vectored.rs
new file mode 100644
index 000000000..be4032294
--- /dev/null
+++ b/tokio/src/io/util/write_vectored.rs
@@ -0,0 +1,47 @@
+use crate::io::AsyncWrite;
+
+use pin_project_lite::pin_project;
+use std::io;
+use std::marker::PhantomPinned;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+use std::{future::Future, io::IoSlice};
+
+pin_project! {
+    /// A future to write a slice of buffers to an `AsyncWrite`.
+    #[derive(Debug)]
+    #[must_use = "futures do nothing unless you `.await` or poll them"]
+    pub struct WriteVectored<'a, 'b, W: ?Sized> {
+        writer: &'a mut W,
+        bufs: &'a [IoSlice<'b>],
+        // Make this future `!Unpin` for compatibility with async trait methods.
+        #[pin]
+        _pin: PhantomPinned,
+    }
+}
+
+pub(crate) fn write_vectored<'a, 'b, W>(
+    writer: &'a mut W,
+    bufs: &'a [IoSlice<'b>],
+) -> WriteVectored<'a, 'b, W>
+where
+    W: AsyncWrite + Unpin + ?Sized,
+{
+    WriteVectored {
+        writer,
+        bufs,
+        _pin: PhantomPinned,
+    }
+}
+
+impl<W> Future for WriteVectored<'_, '_, W>
+where
+    W: AsyncWrite + Unpin + ?Sized,
+{
+    type Output = io::Result<usize>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> {
+        let me = self.project();
+        Pin::new(&mut *me.writer).poll_write_vectored(cx, me.bufs)
+    }
+}