madsim_real_tokio/io/
split.rs

1//! Split a single value implementing `AsyncRead + AsyncWrite` into separate
2//! `AsyncRead` and `AsyncWrite` handles.
3//!
4//! To restore this read/write object from its `split::ReadHalf` and
5//! `split::WriteHalf` use `unsplit`.
6
7use crate::io::{AsyncRead, AsyncWrite, ReadBuf};
8
9use std::fmt;
10use std::io;
11use std::pin::Pin;
12use std::sync::Arc;
13use std::sync::Mutex;
14use std::task::{Context, Poll};
15
16cfg_io_util! {
17    /// The readable half of a value returned from [`split`](split()).
18    pub struct ReadHalf<T> {
19        inner: Arc<Inner<T>>,
20    }
21
22    /// The writable half of a value returned from [`split`](split()).
23    pub struct WriteHalf<T> {
24        inner: Arc<Inner<T>>,
25    }
26
27    /// Splits a single value implementing `AsyncRead + AsyncWrite` into separate
28    /// `AsyncRead` and `AsyncWrite` handles.
29    ///
30    /// To restore this read/write object from its `ReadHalf` and
31    /// `WriteHalf` use [`unsplit`](ReadHalf::unsplit()).
32    pub fn split<T>(stream: T) -> (ReadHalf<T>, WriteHalf<T>)
33    where
34        T: AsyncRead + AsyncWrite,
35    {
36        let is_write_vectored = stream.is_write_vectored();
37
38        let inner = Arc::new(Inner {
39            stream: Mutex::new(stream),
40            is_write_vectored,
41        });
42
43        let rd = ReadHalf {
44            inner: inner.clone(),
45        };
46
47        let wr = WriteHalf { inner };
48
49        (rd, wr)
50    }
51}
52
53struct Inner<T> {
54    stream: Mutex<T>,
55    is_write_vectored: bool,
56}
57
58impl<T> Inner<T> {
59    fn with_lock<R>(&self, f: impl FnOnce(Pin<&mut T>) -> R) -> R {
60        let mut guard = self.stream.lock().unwrap();
61
62        // safety: we do not move the stream.
63        let stream = unsafe { Pin::new_unchecked(&mut *guard) };
64
65        f(stream)
66    }
67}
68
69impl<T> ReadHalf<T> {
70    /// Checks if this `ReadHalf` and some `WriteHalf` were split from the same
71    /// stream.
72    pub fn is_pair_of(&self, other: &WriteHalf<T>) -> bool {
73        other.is_pair_of(self)
74    }
75
76    /// Reunites with a previously split `WriteHalf`.
77    ///
78    /// # Panics
79    ///
80    /// If this `ReadHalf` and the given `WriteHalf` do not originate from the
81    /// same `split` operation this method will panic.
82    /// This can be checked ahead of time by comparing the stream ID
83    /// of the two halves.
84    #[track_caller]
85    pub fn unsplit(self, wr: WriteHalf<T>) -> T
86    where
87        T: Unpin,
88    {
89        if self.is_pair_of(&wr) {
90            drop(wr);
91
92            let inner = Arc::try_unwrap(self.inner)
93                .ok()
94                .expect("`Arc::try_unwrap` failed");
95
96            inner.stream.into_inner().unwrap()
97        } else {
98            panic!("Unrelated `split::Write` passed to `split::Read::unsplit`.")
99        }
100    }
101}
102
103impl<T> WriteHalf<T> {
104    /// Checks if this `WriteHalf` and some `ReadHalf` were split from the same
105    /// stream.
106    pub fn is_pair_of(&self, other: &ReadHalf<T>) -> bool {
107        Arc::ptr_eq(&self.inner, &other.inner)
108    }
109}
110
111impl<T: AsyncRead> AsyncRead for ReadHalf<T> {
112    fn poll_read(
113        self: Pin<&mut Self>,
114        cx: &mut Context<'_>,
115        buf: &mut ReadBuf<'_>,
116    ) -> Poll<io::Result<()>> {
117        self.inner.with_lock(|stream| stream.poll_read(cx, buf))
118    }
119}
120
121impl<T: AsyncWrite> AsyncWrite for WriteHalf<T> {
122    fn poll_write(
123        self: Pin<&mut Self>,
124        cx: &mut Context<'_>,
125        buf: &[u8],
126    ) -> Poll<Result<usize, io::Error>> {
127        self.inner.with_lock(|stream| stream.poll_write(cx, buf))
128    }
129
130    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
131        self.inner.with_lock(|stream| stream.poll_flush(cx))
132    }
133
134    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
135        self.inner.with_lock(|stream| stream.poll_shutdown(cx))
136    }
137
138    fn poll_write_vectored(
139        self: Pin<&mut Self>,
140        cx: &mut Context<'_>,
141        bufs: &[io::IoSlice<'_>],
142    ) -> Poll<Result<usize, io::Error>> {
143        self.inner
144            .with_lock(|stream| stream.poll_write_vectored(cx, bufs))
145    }
146
147    fn is_write_vectored(&self) -> bool {
148        self.inner.is_write_vectored
149    }
150}
151
152unsafe impl<T: Send> Send for ReadHalf<T> {}
153unsafe impl<T: Send> Send for WriteHalf<T> {}
154unsafe impl<T: Sync> Sync for ReadHalf<T> {}
155unsafe impl<T: Sync> Sync for WriteHalf<T> {}
156
157impl<T: fmt::Debug> fmt::Debug for ReadHalf<T> {
158    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
159        fmt.debug_struct("split::ReadHalf").finish()
160    }
161}
162
163impl<T: fmt::Debug> fmt::Debug for WriteHalf<T> {
164    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
165        fmt.debug_struct("split::WriteHalf").finish()
166    }
167}
madsim_real_tokio/io/split.rs

madsim_real_tokio/io/
split.rs
