diff --git a/tokio-tcp/src/split.rs b/tokio-tcp/src/split.rs index fa001a9b6..cafbd4926 100644 --- a/tokio-tcp/src/split.rs +++ b/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 for TcpStreamWriteHalf { } } +impl AsRef for TcpStreamReadHalfMut<'_> { + fn as_ref(&self) -> &TcpStream { + self.0 + } +} + +impl AsRef 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> { + self.0.poll_read_priv(cx, buf) + } + + fn poll_read_buf( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut B, + ) -> Poll> { + 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> { + self.0.poll_write_priv(cx, buf) + } + + #[inline] + fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + // 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> { + self.0.shutdown(Shutdown::Write).into() + } + + fn poll_write_buf( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut B, + ) -> Poll> { + self.0.poll_write_buf_priv(cx, buf) + } +} diff --git a/tokio-tcp/src/stream.rs b/tokio-tcp/src/stream.rs index df4de87e0..1c5d63d7a 100644 --- a/tokio-tcp/src/stream.rs +++ b/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`):