tcp: implement "split_mut" for TcpStream (#1289)

tokio-tcp/src/split.rs

@@ -42,6 +42,26 @@ pub(crate) fn split(stream: TcpStream) -> (TcpStreamReadHalf, TcpStreamWriteHalf
     )
 }
 
+/// Read half of a `TcpStream`.
+#[derive(Debug)]
+pub struct TcpStreamReadHalfMut<'a>(&'a TcpStream);
+
+/// Write half of a `TcpStream`.
+///
+/// Note that in the `AsyncWrite` implemenation of `TcpStreamWriteHalf`,
+/// `poll_shutdown` actually shuts down the TCP stream in the write direction.
+#[derive(Debug)]
+pub struct TcpStreamWriteHalfMut<'a>(&'a TcpStream);
+
+pub(crate) fn split_mut<'a>(
+    stream: &'a mut TcpStream,
+) -> (TcpStreamReadHalfMut<'a>, TcpStreamWriteHalfMut<'a>) {
+    (
+        TcpStreamReadHalfMut(&*stream),
+        TcpStreamWriteHalfMut(&*stream),
+    )
+}
+
 /// Error indicating two halves were not from the same stream, and thus could
 /// not be `reunite`d.
 #[derive(Debug)]
@@ -97,6 +117,18 @@ impl AsRef<TcpStream> for TcpStreamWriteHalf {
     }
 }
 
+impl AsRef<TcpStream> for TcpStreamReadHalfMut<'_> {
+    fn as_ref(&self) -> &TcpStream {
+        self.0
+    }
+}
+
+impl AsRef<TcpStream> for TcpStreamWriteHalfMut<'_> {
+    fn as_ref(&self) -> &TcpStream {
+        self.0
+    }
+}
+
 impl AsyncRead for TcpStreamReadHalf {
     unsafe fn prepare_uninitialized_buffer(&self, _: &mut [u8]) -> bool {
         false
@@ -147,3 +179,54 @@ impl AsyncWrite for TcpStreamWriteHalf {
         self.0.poll_write_buf_priv(cx, buf)
     }
 }
+
+impl AsyncRead for TcpStreamReadHalfMut<'_> {
+    unsafe fn prepare_uninitialized_buffer(&self, _: &mut [u8]) -> bool {
+        false
+    }
+
+    fn poll_read(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut [u8],
+    ) -> Poll<io::Result<usize>> {
+        self.0.poll_read_priv(cx, buf)
+    }
+
+    fn poll_read_buf<B: BufMut>(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut B,
+    ) -> Poll<io::Result<usize>> {
+        self.0.poll_read_buf_priv(cx, buf)
+    }
+}
+
+impl AsyncWrite for TcpStreamWriteHalfMut<'_> {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        self.0.poll_write_priv(cx, buf)
+    }
+
+    #[inline]
+    fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
+        // tcp flush is a no-op
+        Poll::Ready(Ok(()))
+    }
+
+    // `poll_shutdown` on a write half shutdowns the stream in the "write" direction.
+    fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
+        self.0.shutdown(Shutdown::Write).into()
+    }
+
+    fn poll_write_buf<B: Buf>(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut B,
+    ) -> Poll<io::Result<usize>> {
+        self.0.poll_write_buf_priv(cx, buf)
+    }
+}

tokio-tcp/src/stream.rs

@@ -1,4 +1,7 @@
-use crate::split::{split, TcpStreamReadHalf, TcpStreamWriteHalf};
+use crate::split::{
+    split, split_mut, TcpStreamReadHalf, TcpStreamReadHalfMut, TcpStreamWriteHalf,
+    TcpStreamWriteHalfMut,
+};
 use bytes::{Buf, BufMut};
 use futures_core::ready;
 use iovec::IoVec;
@@ -724,6 +727,15 @@ impl TcpStream {
         split(self)
     }
 
+    /// Split a `TcpStream` into a read half and a write half, which can be used
+    /// to read and write the stream concurrently.
+    ///
+    /// See the module level documenation of [`split`](super::split) for more
+    /// details.
+    pub fn split_mut<'a>(&'a mut self) -> (TcpStreamReadHalfMut<'a>, TcpStreamWriteHalfMut<'a>) {
+        split_mut(self)
+    }
+
     // == Poll IO functions that takes `&self` ==
     //
     // They are not public because (taken from the doc of `PollEvented`):