mirror of
https://github.com/tokio-rs/tokio.git
synced 2026-08-27 00:00:12 +02:00
io: bring back split utility (#1521)
Bring back `split` utility as a free fn instead of a method on `AsyncRead`. This utility wraps the `stream` in an `Arc` and uses mutual exclusion to ensure correct access. Additionally, the specialized `split_mut` fn on TcpStream and UdsStream is promoted to `split`.
This commit is contained in:
Carl Lerche
+17
-154
@@ -1,137 +1,39 @@
|
||||
//! `TcpStream` split support.
|
||||
//!
|
||||
//! A `TcpStream` can be split into a `TcpStreamReadHalf` and a
|
||||
//! `TcpStreamWriteHalf` with the `TcpStream::split` method. `TcpStreamReadHalf`
|
||||
//! implements `AsyncRead` while `TcpStreamWriteHalf` implements `AsyncWrite`.
|
||||
//! The two halves can be used concurrently, even from multiple tasks.
|
||||
//! A `TcpStream` can be split into a `ReadHalf` and a
|
||||
//! `WriteHalf` with the `TcpStream::split` method. `ReadHalf`
|
||||
//! implements `AsyncRead` while `WriteHalf` implements `AsyncWrite`.
|
||||
//!
|
||||
//! Compared to the generic split of `AsyncRead + AsyncWrite`, this specialized
|
||||
//! split gives read and write halves that are faster and smaller, because they
|
||||
//! do not use locks. They also provide access to the underlying `TcpStream`
|
||||
//! after split, implementing `AsRef<TcpStream>`. This allows you to call
|
||||
//! `TcpStream` methods that takes `&self`, e.g., to get local and peer
|
||||
//! addresses, to get and set socket options, and to shutdown the sockets.
|
||||
//! split has no associated overhead and enforces all invariants at the type
|
||||
//! level.
|
||||
|
||||
use super::TcpStream;
|
||||
|
||||
use tokio_io::{AsyncRead, AsyncWrite};
|
||||
|
||||
use bytes::{Buf, BufMut};
|
||||
use std::error::Error;
|
||||
use std::fmt;
|
||||
use std::io;
|
||||
use std::net::Shutdown;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
/// Read half of a `TcpStream`.
|
||||
#[derive(Debug)]
|
||||
pub struct TcpStreamReadHalf(Arc<TcpStream>);
|
||||
pub struct ReadHalf<'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 TcpStreamWriteHalf(Arc<TcpStream>);
|
||||
pub struct WriteHalf<'a>(&'a TcpStream);
|
||||
|
||||
pub(crate) fn split(stream: TcpStream) -> (TcpStreamReadHalf, TcpStreamWriteHalf) {
|
||||
let shared = Arc::new(stream);
|
||||
(
|
||||
TcpStreamReadHalf(shared.clone()),
|
||||
TcpStreamWriteHalf(shared),
|
||||
)
|
||||
pub(crate) fn split(stream: &mut TcpStream) -> (ReadHalf<'_>, WriteHalf<'_>) {
|
||||
(ReadHalf(&*stream), WriteHalf(&*stream))
|
||||
}
|
||||
|
||||
/// 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(
|
||||
stream: &mut TcpStream,
|
||||
) -> (TcpStreamReadHalfMut<'_>, TcpStreamWriteHalfMut<'_>) {
|
||||
(
|
||||
TcpStreamReadHalfMut(&*stream),
|
||||
TcpStreamWriteHalfMut(&*stream),
|
||||
)
|
||||
}
|
||||
|
||||
/// Error indicating two halves were not from the same stream, and thus could
|
||||
/// not be `reunite`d.
|
||||
#[derive(Debug)]
|
||||
pub struct ReuniteError(pub TcpStreamReadHalf, pub TcpStreamWriteHalf);
|
||||
|
||||
impl fmt::Display for ReuniteError {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(
|
||||
f,
|
||||
"tried to reunite halves that are not from the same stream"
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl Error for ReuniteError {}
|
||||
|
||||
impl TcpStreamReadHalf {
|
||||
/// Attempts to put the two "halves" of a `TcpStream` back together and
|
||||
/// recover the original stream. Succeeds only if the two "halves"
|
||||
/// originated from the same call to `TcpStream::split`.
|
||||
pub fn reunite(self, other: TcpStreamWriteHalf) -> Result<TcpStream, ReuniteError> {
|
||||
if Arc::ptr_eq(&self.0, &other.0) {
|
||||
drop(other);
|
||||
// Only two instances of the `Arc` are ever created, one for the
|
||||
// reader and one for the writer, and those `Arc`s are never exposed
|
||||
// externally. And so when we drop one here, the other one must be
|
||||
// the only remaining one.
|
||||
Ok(Arc::try_unwrap(self.0).expect("tcp: try_unwrap failed in reunite"))
|
||||
} else {
|
||||
Err(ReuniteError(self, other))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TcpStreamWriteHalf {
|
||||
/// Attempts to put the two "halves" of a `TcpStream` back together and
|
||||
/// recover the original stream. Succeeds only if the two "halves"
|
||||
/// originated from the same call to `TcpStream::split`.
|
||||
pub fn reunite(self, other: TcpStreamReadHalf) -> Result<TcpStream, ReuniteError> {
|
||||
other.reunite(self)
|
||||
}
|
||||
}
|
||||
|
||||
impl AsRef<TcpStream> for TcpStreamReadHalf {
|
||||
fn as_ref(&self) -> &TcpStream {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
|
||||
impl AsRef<TcpStream> for TcpStreamWriteHalf {
|
||||
fn as_ref(&self) -> &TcpStream {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
|
||||
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 {
|
||||
impl AsyncRead for ReadHalf<'_> {
|
||||
unsafe fn prepare_uninitialized_buffer(&self, _: &mut [u8]) -> bool {
|
||||
false
|
||||
}
|
||||
@@ -153,7 +55,7 @@ impl AsyncRead for TcpStreamReadHalf {
|
||||
}
|
||||
}
|
||||
|
||||
impl AsyncWrite for TcpStreamWriteHalf {
|
||||
impl AsyncWrite for WriteHalf<'_> {
|
||||
fn poll_write(
|
||||
self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
@@ -182,53 +84,14 @@ impl AsyncWrite for TcpStreamWriteHalf {
|
||||
}
|
||||
}
|
||||
|
||||
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 AsRef<TcpStream> for ReadHalf<'_> {
|
||||
fn as_ref(&self) -> &TcpStream {
|
||||
self.0
|
||||
}
|
||||
}
|
||||
|
||||
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)
|
||||
impl AsRef<TcpStream> for WriteHalf<'_> {
|
||||
fn as_ref(&self) -> &TcpStream {
|
||||
self.0
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,7 +1,4 @@
|
||||
use super::split::{
|
||||
split, split_mut, TcpStreamReadHalf, TcpStreamReadHalfMut, TcpStreamWriteHalf,
|
||||
TcpStreamWriteHalfMut,
|
||||
};
|
||||
use super::split::{split, ReadHalf, WriteHalf};
|
||||
use crate::driver::Handle;
|
||||
use crate::util::PollEvented;
|
||||
use crate::ToSocketAddrs;
|
||||
@@ -584,19 +581,10 @@ impl TcpStream {
|
||||
///
|
||||
/// See the module level documenation of [`split`](super::split) for more
|
||||
/// details.
|
||||
pub fn split(self) -> (TcpStreamReadHalf, TcpStreamWriteHalf) {
|
||||
pub fn split(&mut self) -> (ReadHalf<'_>, WriteHalf<'_>) {
|
||||
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(&mut self) -> (TcpStreamReadHalfMut<'_>, TcpStreamWriteHalfMut<'_>) {
|
||||
split_mut(self)
|
||||
}
|
||||
|
||||
// == Poll IO functions that takes `&self` ==
|
||||
//
|
||||
// They are not public because (taken from the doc of `PollEvented`):
|
||||
|
||||
+20
-111
@@ -1,15 +1,12 @@
|
||||
//! `UnixStream` split support.
|
||||
//!
|
||||
//! A `UnixStream` can be split into a read half and a write half with `UnixStream::split`
|
||||
//! and `UnixStream::split_mut` methods. The read half implements `AsyncRead` while
|
||||
//! the write half implements `AsyncWrite`. The two halves can be used concurrently.
|
||||
//! A `UnixStream` can be split into a read half and a write half with
|
||||
//! `UnixStream::split`. The read half implements `AsyncRead` while the write
|
||||
//! half implements `AsyncWrite`.
|
||||
//!
|
||||
//! Compared to the generic split of `AsyncRead + AsyncWrite`, this specialized
|
||||
//! split gives read and write halves that are faster and smaller, because they
|
||||
//! do not use locks. They also provide access to the underlying `UnixStream`
|
||||
//! after split, implementing `AsRef<UnixStream>`. This allows you to call
|
||||
//! `UnixStream` methods that takes `&self`, e.g., to get local and peer
|
||||
//! addresses, to get and set socket options, and to shutdown the sockets.
|
||||
//! split has no associated overhead and enforces all invariants at the type
|
||||
//! level.
|
||||
|
||||
use super::UnixStream;
|
||||
|
||||
@@ -19,73 +16,21 @@ use bytes::{Buf, BufMut};
|
||||
use std::io;
|
||||
use std::net::Shutdown;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
/// Read half of a `UnixStream`.
|
||||
#[derive(Debug)]
|
||||
pub struct UnixStreamReadHalf(Arc<UnixStream>);
|
||||
pub struct ReadHalf<'a>(&'a UnixStream);
|
||||
|
||||
/// Write half of a `UnixStream`.
|
||||
///
|
||||
/// Note that in the `AsyncWrite` implementation of `UnixStreamWriteHalf`,
|
||||
/// `poll_shutdown` actually shuts down the stream in the write direction.
|
||||
#[derive(Debug)]
|
||||
pub struct UnixStreamWriteHalf(Arc<UnixStream>);
|
||||
pub struct WriteHalf<'a>(&'a UnixStream);
|
||||
|
||||
/// Read half of a `UnixStream`.
|
||||
#[derive(Debug)]
|
||||
pub struct UnixStreamReadHalfMut<'a>(&'a UnixStream);
|
||||
|
||||
/// Write half of a `UnixStream`.
|
||||
///
|
||||
/// Note that in the `AsyncWrite` implementation of `UnixStreamWriteHalfMut`,
|
||||
/// `poll_shutdown` actually shuts down the stream in the write direction.
|
||||
#[derive(Debug)]
|
||||
pub struct UnixStreamWriteHalfMut<'a>(&'a UnixStream);
|
||||
|
||||
pub(crate) fn split(stream: UnixStream) -> (UnixStreamReadHalf, UnixStreamWriteHalf) {
|
||||
let shared = Arc::new(stream);
|
||||
(
|
||||
UnixStreamReadHalf(shared.clone()),
|
||||
UnixStreamWriteHalf(shared),
|
||||
)
|
||||
pub(crate) fn split(stream: &mut UnixStream) -> (ReadHalf<'_>, WriteHalf<'_>) {
|
||||
(ReadHalf(stream), WriteHalf(stream))
|
||||
}
|
||||
|
||||
pub(crate) fn split_mut(
|
||||
stream: &mut UnixStream,
|
||||
) -> (UnixStreamReadHalfMut<'_>, UnixStreamWriteHalfMut<'_>) {
|
||||
(
|
||||
UnixStreamReadHalfMut(stream),
|
||||
UnixStreamWriteHalfMut(stream),
|
||||
)
|
||||
}
|
||||
|
||||
impl AsRef<UnixStream> for UnixStreamReadHalf {
|
||||
fn as_ref(&self) -> &UnixStream {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
|
||||
impl AsRef<UnixStream> for UnixStreamWriteHalf {
|
||||
fn as_ref(&self) -> &UnixStream {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
|
||||
impl AsRef<UnixStream> for UnixStreamReadHalfMut<'_> {
|
||||
fn as_ref(&self) -> &UnixStream {
|
||||
self.0
|
||||
}
|
||||
}
|
||||
|
||||
impl AsRef<UnixStream> for UnixStreamWriteHalfMut<'_> {
|
||||
fn as_ref(&self) -> &UnixStream {
|
||||
self.0
|
||||
}
|
||||
}
|
||||
|
||||
impl AsyncRead for UnixStreamReadHalf {
|
||||
impl AsyncRead for ReadHalf<'_> {
|
||||
unsafe fn prepare_uninitialized_buffer(&self, _: &mut [u8]) -> bool {
|
||||
false
|
||||
}
|
||||
@@ -107,29 +52,7 @@ impl AsyncRead for UnixStreamReadHalf {
|
||||
}
|
||||
}
|
||||
|
||||
impl AsyncRead for UnixStreamReadHalfMut<'_> {
|
||||
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 UnixStreamWriteHalf {
|
||||
impl AsyncWrite for WriteHalf<'_> {
|
||||
fn poll_write(
|
||||
self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
@@ -155,28 +78,14 @@ impl AsyncWrite for UnixStreamWriteHalf {
|
||||
}
|
||||
}
|
||||
|
||||
impl AsyncWrite for UnixStreamWriteHalfMut<'_> {
|
||||
fn poll_write(
|
||||
self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
buf: &[u8],
|
||||
) -> Poll<io::Result<usize>> {
|
||||
self.0.poll_write_priv(cx, buf)
|
||||
}
|
||||
|
||||
fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
|
||||
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)
|
||||
impl AsRef<UnixStream> for ReadHalf<'_> {
|
||||
fn as_ref(&self) -> &UnixStream {
|
||||
self.0
|
||||
}
|
||||
}
|
||||
|
||||
impl AsRef<UnixStream> for WriteHalf<'_> {
|
||||
fn as_ref(&self) -> &UnixStream {
|
||||
self.0
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,7 +1,4 @@
|
||||
use super::split::{
|
||||
split, split_mut, UnixStreamReadHalf, UnixStreamReadHalfMut, UnixStreamWriteHalf,
|
||||
UnixStreamWriteHalfMut,
|
||||
};
|
||||
use super::split::{split, ReadHalf, WriteHalf};
|
||||
use super::ucred::{self, UCred};
|
||||
use crate::driver::Handle;
|
||||
use crate::util::PollEvented;
|
||||
@@ -112,18 +109,9 @@ impl UnixStream {
|
||||
///
|
||||
/// See the module level documenation of [`split`](super::split) for more
|
||||
/// details.
|
||||
pub fn split(self) -> (UnixStreamReadHalf, UnixStreamWriteHalf) {
|
||||
pub fn split(&mut self) -> (ReadHalf<'_>, WriteHalf<'_>) {
|
||||
split(self)
|
||||
}
|
||||
|
||||
/// Split a `UnixStream` 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(&mut self) -> (UnixStreamReadHalfMut<'_>, UnixStreamWriteHalfMut<'_>) {
|
||||
split_mut(self)
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<UnixStream> for mio_uds::UnixStream {
|
||||
|
||||
Reference in New Issue
Block a user