From 5d29bdfb0cb9e2562c662c76a362bb823f88ffd5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jakub=20Kub=C3=ADk?= <jakub.kubik.it@protonmail.com>
Date: Thu, 10 Aug 2023 10:11:00 +0200
Subject: [PATCH] io: delegate `WriteHalf::poll_write_vectored` (#5914)

---
 tokio/src/io/split.rs   | 17 +++++++++++++++++
 tokio/tests/io_split.rs | 39 ++++++++++++++++++++++++++++++++++++++-
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/tokio/src/io/split.rs b/tokio/src/io/split.rs
index f067b65a826..63f0960e4f3 100644
--- a/tokio/src/io/split.rs
+++ b/tokio/src/io/split.rs
@@ -35,9 +35,12 @@ cfg_io_util! {
     where
         T: AsyncRead + AsyncWrite,
     {
+        let is_write_vectored = stream.is_write_vectored();
+
         let inner = Arc::new(Inner {
             locked: AtomicBool::new(false),
             stream: UnsafeCell::new(stream),
+            is_write_vectored,
         });
 
         let rd = ReadHalf {
@@ -53,6 +56,7 @@ cfg_io_util! {
 struct Inner<T> {
     locked: AtomicBool,
     stream: UnsafeCell<T>,
+    is_write_vectored: bool,
 }
 
 struct Guard<'a, T> {
@@ -131,6 +135,19 @@ impl<T: AsyncWrite> AsyncWrite for WriteHalf<T> {
         let mut inner = ready!(self.inner.poll_lock(cx));
         inner.stream_pin().poll_shutdown(cx)
     }
+
+    fn poll_write_vectored(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        bufs: &[io::IoSlice<'_>],
+    ) -> Poll<Result<usize, io::Error>> {
+        let mut inner = ready!(self.inner.poll_lock(cx));
+        inner.stream_pin().poll_write_vectored(cx, bufs)
+    }
+
+    fn is_write_vectored(&self) -> bool {
+        self.inner.is_write_vectored
+    }
 }
 
 impl<T> Inner<T> {
diff --git a/tokio/tests/io_split.rs b/tokio/tests/io_split.rs
index 77b77a3a04c..9f17c9eb14e 100644
--- a/tokio/tests/io_split.rs
+++ b/tokio/tests/io_split.rs
@@ -1,7 +1,9 @@
 #![warn(rust_2018_idioms)]
 #![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi does not support panic recovery
 
-use tokio::io::{split, AsyncRead, AsyncWrite, ReadBuf, ReadHalf, WriteHalf};
+use tokio::io::{
+    split, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf, ReadHalf, WriteHalf,
+};
 
 use std::io;
 use std::pin::Pin;
@@ -36,6 +38,18 @@ impl AsyncWrite for RW {
     fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
         Poll::Ready(Ok(()))
     }
+
+    fn poll_write_vectored(
+        self: Pin<&mut Self>,
+        _cx: &mut Context<'_>,
+        _bufs: &[io::IoSlice<'_>],
+    ) -> Poll<Result<usize, io::Error>> {
+        Poll::Ready(Ok(2))
+    }
+
+    fn is_write_vectored(&self) -> bool {
+        true
+    }
 }
 
 #[test]
@@ -77,3 +91,26 @@ fn unsplit_err2() {
     let (r, _) = split(RW);
     r.unsplit(w);
 }
+
+#[test]
+fn method_delegation() {
+    let (mut r, mut w) = split(RW);
+    let mut buf = [0; 1];
+
+    tokio_test::block_on(async move {
+        assert_eq!(1, r.read(&mut buf).await.unwrap());
+        assert_eq!(b'z', buf[0]);
+
+        assert_eq!(1, w.write(&[b'x']).await.unwrap());
+        assert_eq!(
+            2,
+            w.write_vectored(&[io::IoSlice::new(&[b'x'])])
+                .await
+                .unwrap()
+        );
+        assert!(w.is_write_vectored());
+
+        assert!(w.flush().await.is_ok());
+        assert!(w.shutdown().await.is_ok());
+    });
+}