nuclei/io/util/
split.rs

1//! [`split`]: a stream that reads and writes, as a read half and a write half
2//! two tasks can own apart.
3//!
4//! The halves share the stream behind a flag each poll call holds for its
5//! length (a poll of the stream, never an await), so a contended call spins
6//! for at most one poll of the other half. A nuclei [`Handle`](crate::Handle)
7//! needs none of this: `&Handle` reads and writes on its own.
8
9use core::cell::UnsafeCell;
10use core::pin::Pin;
11use core::task::{Context, Poll};
12use std::io::{self, IoSlice};
13use std::sync::Arc;
14
15use futures_io::{AsyncRead, AsyncWrite};
16
17use crate::sync::internal::{AtomicBool, Ordering};
18
19struct Inner<T> {
20    locked: AtomicBool,
21    stream: UnsafeCell<T>,
22}
23
24// SAFETY: the stream is only reached through `with`, which holds the flag:
25// one half at a time, so `T: Send` is all sharing it needs.
26unsafe impl<T: Send> Send for Inner<T> {}
27// SAFETY: as for `Send`.
28unsafe impl<T: Send> Sync for Inner<T> {}
29
30impl<T> Inner<T> {
31    /// Runs `f` on the stream, holding the flag for its length.
32    fn with<R>(&self, f: impl FnOnce(Pin<&mut T>) -> R) -> R {
33        while self
34            .locked
35            .compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
36            .is_err()
37        {
38            core::hint::spin_loop();
39        }
40        struct Release<'a>(&'a AtomicBool);
41        impl Drop for Release<'_> {
42            fn drop(&mut self) {
43                self.0.store(false, Ordering::Release);
44            }
45        }
46        let _release = Release(&self.locked);
47        // SAFETY: the flag is held: no other reference to the stream exists;
48        // the stream never moves out of the `Arc` while both halves live, so
49        // pinning it is sound.
50        f(unsafe { Pin::new_unchecked(&mut *self.stream.get()) })
51    }
52}
53
54/// The read half of a stream [`split`] in two.
55pub struct ReadHalf<T> {
56    inner: Arc<Inner<T>>,
57}
58
59/// The write half of a stream [`split`] in two.
60pub struct WriteHalf<T> {
61    inner: Arc<Inner<T>>,
62}
63
64/// Splits `stream` into a read half and a write half; [`ReadHalf::unsplit`]
65/// puts them back together.
66pub fn split<T: AsyncRead + AsyncWrite>(stream: T) -> (ReadHalf<T>, WriteHalf<T>) {
67    let inner = Arc::new(Inner {
68        locked: AtomicBool::new(false),
69        stream: UnsafeCell::new(stream),
70    });
71    (
72        ReadHalf {
73            inner: Arc::clone(&inner),
74        },
75        WriteHalf { inner },
76    )
77}
78
79impl<T> ReadHalf<T> {
80    /// Whether `other` is this half's write half.
81    pub fn is_pair_of(&self, other: &WriteHalf<T>) -> bool {
82        Arc::ptr_eq(&self.inner, &other.inner)
83    }
84
85    /// The stream, back from its two halves.
86    ///
87    /// # Panics
88    ///
89    /// When `write` is not this half's write half.
90    #[expect(
91        clippy::unreachable,
92        reason = "the other half is dropped above, so this one holds the last reference"
93    )]
94    pub fn unsplit(self, write: WriteHalf<T>) -> T
95    where
96        T: Unpin,
97    {
98        assert!(
99            self.is_pair_of(&write),
100            "unsplit of halves of two different streams"
101        );
102        drop(write);
103        match Arc::try_unwrap(self.inner) {
104            Ok(inner) => inner.stream.into_inner(),
105            Err(_) => unreachable!("both halves were here, and are gone"),
106        }
107    }
108}
109
110impl<T> WriteHalf<T> {
111    /// Whether `other` is this half's read half.
112    pub fn is_pair_of(&self, other: &ReadHalf<T>) -> bool {
113        Arc::ptr_eq(&self.inner, &other.inner)
114    }
115}
116
117impl<T: AsyncRead> AsyncRead for ReadHalf<T> {
118    fn poll_read(
119        self: Pin<&mut Self>,
120        cx: &mut Context<'_>,
121        buf: &mut [u8],
122    ) -> Poll<io::Result<usize>> {
123        self.inner.with(|stream| stream.poll_read(cx, buf))
124    }
125}
126
127impl<T: AsyncWrite> AsyncWrite for WriteHalf<T> {
128    fn poll_write(
129        self: Pin<&mut Self>,
130        cx: &mut Context<'_>,
131        buf: &[u8],
132    ) -> Poll<io::Result<usize>> {
133        self.inner.with(|stream| stream.poll_write(cx, buf))
134    }
135
136    fn poll_write_vectored(
137        self: Pin<&mut Self>,
138        cx: &mut Context<'_>,
139        bufs: &[IoSlice<'_>],
140    ) -> Poll<io::Result<usize>> {
141        self.inner
142            .with(|stream| stream.poll_write_vectored(cx, bufs))
143    }
144
145    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
146        self.inner.with(|stream| stream.poll_flush(cx))
147    }
148
149    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
150        self.inner.with(|stream| stream.poll_close(cx))
151    }
152}
153
154impl<T> core::fmt::Debug for ReadHalf<T> {
155    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
156        f.debug_struct("ReadHalf").finish_non_exhaustive()
157    }
158}
159
160impl<T> core::fmt::Debug for WriteHalf<T> {
161    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
162        f.debug_struct("WriteHalf").finish_non_exhaustive()
163    }
164}
nuclei/io/util/split.rs

nuclei/io/util/
split.rs
