//! Split a single value implementing `AsyncRead + AsyncWrite` into separate //! `AsyncRead` and `AsyncWrite` handles. //! //! To restore this read/write object from its `split::ReadHalf` and //! `split::WriteHalf` use `unsplit`. use crate::{AsyncRead, AsyncWrite}; use bytes::{Buf, BufMut}; use futures_core::ready; use std::cell::UnsafeCell; use std::fmt; use std::io; use std::pin::Pin; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering::{Acquire, Release}; use std::sync::Arc; use std::task::{Context, Poll}; /// The readable half of a value returned from `split`. pub struct ReadHalf { inner: Arc>, } /// The writable half of a value returned from `split`. pub struct WriteHalf { inner: Arc>, } struct Inner { locked: AtomicBool, stream: UnsafeCell, } struct Guard<'a, T> { inner: &'a Inner, } /// Split a single value implementing `AsyncRead + AsyncWrite` into separate /// `AsyncRead` and `AsyncWrite` handles. /// /// To restore this read/write object from its `split::ReadHalf` and /// `split::WriteHalf` use `unsplit`. pub fn split(stream: T) -> (ReadHalf, WriteHalf) where T: AsyncRead + AsyncWrite, { let inner = Arc::new(Inner { locked: AtomicBool::new(false), stream: UnsafeCell::new(stream), }); let rd = ReadHalf { inner: inner.clone(), }; let wr = WriteHalf { inner }; (rd, wr) } impl ReadHalf { /// Reunite with a previously split `WriteHalf`. /// /// # Panics /// /// If this `ReadHalf` and the given `WriteHalf` do not originate from the /// same `split` operation this method will panic. pub fn unsplit(self, wr: WriteHalf) -> T { if Arc::ptr_eq(&self.inner, &wr.inner) { drop(wr); let inner = Arc::try_unwrap(self.inner) .ok() .expect("Arc::try_unwrap failed"); inner.stream.into_inner() } else { panic!("Unrelated `split::Write` passed to `split::Read::unsplit`.") } } } impl AsyncRead for ReadHalf { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut [u8], ) -> Poll> { let mut inner = ready!(self.inner.poll_lock(cx)); inner.stream_pin().poll_read(cx, buf) } fn poll_read_buf( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut B, ) -> Poll> { let mut inner = ready!(self.inner.poll_lock(cx)); inner.stream_pin().poll_read_buf(cx, buf) } } impl AsyncWrite for WriteHalf { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { let mut inner = ready!(self.inner.poll_lock(cx)); inner.stream_pin().poll_write(cx, buf) } fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let mut inner = ready!(self.inner.poll_lock(cx)); inner.stream_pin().poll_flush(cx) } fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let mut inner = ready!(self.inner.poll_lock(cx)); inner.stream_pin().poll_shutdown(cx) } fn poll_write_buf( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut B, ) -> Poll> { let mut inner = ready!(self.inner.poll_lock(cx)); inner.stream_pin().poll_write_buf(cx, buf) } } impl Inner { fn poll_lock(&self, cx: &mut Context<'_>) -> Poll> { if !self.locked.compare_and_swap(false, true, Acquire) { Poll::Ready(Guard { inner: self }) } else { // Spin... but investigate a better strategy ::std::thread::yield_now(); cx.waker().wake_by_ref(); Poll::Pending } } } impl Guard<'_, T> { fn stream_pin(&mut self) -> Pin<&mut T> { // safety: the stream is pinned in `Arc` and the `Guard` ensures mutual // exclusion. unsafe { Pin::new_unchecked(&mut *self.inner.stream.get()) } } } impl Drop for Guard<'_, T> { fn drop(&mut self) { self.inner.locked.store(false, Release); } } unsafe impl Send for ReadHalf {} unsafe impl Send for WriteHalf {} unsafe impl Sync for ReadHalf {} unsafe impl Sync for WriteHalf {} impl fmt::Debug for ReadHalf { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("split::ReadHalf").finish() } } impl fmt::Debug for WriteHalf { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("split::WriteHalf").finish() } }