trillium_websockets/websocket_connection/
split_compat.rs

1// Retained only so that callers of the deprecated `take_inbound_stream` and `inbound_stream` keep
2// working; remove it at the next breaking release. Removal is this module, the `mod` line and
3// `DriverSlot` field type in `websocket_connection.rs` (the field becomes `Interrupt<Driver>`,
4// polled with `driver::poll_inbound`), and the `inbound_taken` branch in `direction.rs`.
5//
6// An unsplit conn owns its driver outright. Splitting moves the driver behind a mutex shared by the
7// conn, which keeps sending, and the taken stream, which reads.
8
9use super::WebSocketConn;
10use crate::driver::{Driver, MessageResult, poll_inbound, write_out_unless_message};
11use futures_lite::Stream;
12use std::{
13    pin::Pin,
14    sync::{Arc, Mutex, MutexGuard, PoisonError},
15    task::{Context, Poll, Wake, Waker},
16};
17use swansong::Interrupt;
18use tungstenite::{Error, Message};
19
20impl WebSocketConn {
21    /// Take the inbound message stream from this conn.
22    ///
23    /// The conn keeps sending, and yields no inbound messages of its own afterwards. Returns
24    /// `None` if the stream was already taken.
25    #[deprecated(
26        note = "a WebSocketConn is itself the stream of inbound messages; poll it directly"
27    )]
28    pub fn take_inbound_stream(&mut self) -> Option<impl Stream<Item = MessageResult> + use<>> {
29        self.driver.split()
30    }
31
32    /// Borrow the inbound message stream from this conn, unless it was taken with
33    /// [`take_inbound_stream`][Self::take_inbound_stream].
34    #[deprecated(
35        note = "a WebSocketConn is itself the stream of inbound messages; poll it directly"
36    )]
37    pub fn inbound_stream(&mut self) -> Option<impl Stream<Item = MessageResult> + '_> {
38        if self.driver.is_split() {
39            None
40        } else {
41            Some(self)
42        }
43    }
44
45    pub(crate) fn inbound_taken(&self) -> bool {
46        self.driver.is_split()
47    }
48
49    pub(crate) fn poll_write_out(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
50        self.driver.poll_write_out(cx)
51    }
52}
53
54// The large variant is the usual one. Boxing it would add an indirection to every conn to save
55// space only on split ones.
56#[allow(clippy::large_enum_variant)]
57pub(crate) enum DriverSlot {
58    Owned(Interrupt<Driver>),
59    Shared(Arc<SharedDriver>),
60}
61
62impl DriverSlot {
63    pub(crate) fn is_split(&self) -> bool {
64        matches!(self, Self::Shared(_))
65    }
66
67    fn split(&mut self) -> Option<TakenInbound> {
68        if self.is_split() {
69            return None;
70        }
71
72        let shared = Arc::new(SharedDriver {
73            driver: Mutex::new(None),
74            write_wakers: Arc::default(),
75        });
76
77        match std::mem::replace(self, Self::Shared(Arc::clone(&shared))) {
78            Self::Owned(driver) => {
79                *lock(&shared.driver) = Some(driver);
80                Some(TakenInbound(shared))
81            }
82            already_split => {
83                *self = already_split;
84                None
85            }
86        }
87    }
88
89    pub(crate) fn poll_ready_to_write(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
90        match self {
91            Self::Owned(driver) => driver.poll_ready_to_write(cx),
92            Self::Shared(shared) => shared.poll_write(Half::Conn, cx, Driver::poll_ready_to_write),
93        }
94    }
95
96    pub(crate) fn write(&mut self, message: Message) -> Result<(), Error> {
97        match self {
98            Self::Owned(driver) => driver.write(message),
99            Self::Shared(shared) => shared
100                .with_driver(|driver| driver.write(message))
101                .unwrap_or(Err(Error::AlreadyClosed)),
102        }
103    }
104
105    pub(crate) fn close(&mut self) -> Result<(), Error> {
106        match self {
107            Self::Owned(driver) => driver.close(),
108            Self::Shared(shared) => shared
109                .with_driver(|driver| driver.close())
110                .unwrap_or(Err(Error::AlreadyClosed)),
111        }
112    }
113
114    pub(crate) fn poll_write_out(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
115        match self {
116            Self::Owned(driver) => driver.poll_write_out(cx),
117            Self::Shared(shared) => shared.poll_write(Half::Conn, cx, Driver::poll_write_out),
118        }
119    }
120
121    pub(crate) fn poll_flush(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
122        match self {
123            Self::Owned(driver) => driver.poll_flush(cx),
124            Self::Shared(shared) => shared.poll_write(Half::Conn, cx, Driver::poll_flush),
125        }
126    }
127
128    pub(crate) fn poll_inbound(&mut self, cx: &mut Context<'_>) -> Poll<Option<MessageResult>> {
129        match self {
130            Self::Owned(driver) => poll_inbound(driver, cx),
131            Self::Shared(_) => Poll::Ready(None),
132        }
133    }
134}
135
136pub(crate) struct SharedDriver {
137    // `None` only while `DriverSlot::split` moves the driver in.
138    driver: Mutex<Option<Interrupt<Driver>>>,
139    write_wakers: Arc<WriteWakers>,
140}
141
142impl SharedDriver {
143    fn with_driver<T>(&self, f: impl FnOnce(&mut Interrupt<Driver>) -> T) -> Option<T> {
144        lock(&self.driver).as_mut().map(f)
145    }
146
147    /// Either half can be left waiting on the socket to write, possibly from separate tasks, and a
148    /// socket keeps only the most recent waker it was polled with. Writing with a waker that wakes
149    /// both halves keeps one half's wait from replacing the other's.
150    fn poll_write(
151        &self,
152        half: Half,
153        cx: &mut Context<'_>,
154        poll: impl FnOnce(&mut Driver, &mut Context<'_>) -> Poll<Result<(), Error>>,
155    ) -> Poll<Result<(), Error>> {
156        self.write_wakers.register(half, cx.waker());
157        let waker = Waker::from(Arc::clone(&self.write_wakers));
158        let mut cx = Context::from_waker(&waker);
159        self.with_driver(|driver| poll(driver, &mut cx))
160            .unwrap_or(Poll::Ready(Err(Error::AlreadyClosed)))
161    }
162}
163
164struct TakenInbound(Arc<SharedDriver>);
165
166impl Stream for TakenInbound {
167    type Item = MessageResult;
168
169    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
170        let shared = &self.0;
171        let poll = shared
172            .with_driver(|driver| Pin::new(driver).poll_next(cx))
173            .unwrap_or(Poll::Ready(None));
174        write_out_unless_message(poll, || {
175            shared.poll_write(Half::Inbound, cx, Driver::poll_write_out)
176        })
177    }
178}
179
180#[derive(Clone, Copy)]
181enum Half {
182    Inbound,
183    Conn,
184}
185
186#[derive(Default)]
187struct WriteWakers {
188    inbound: Mutex<Option<Waker>>,
189    conn: Mutex<Option<Waker>>,
190}
191
192impl WriteWakers {
193    fn register(&self, half: Half, waker: &Waker) {
194        let mut slot = lock(match half {
195            Half::Inbound => &self.inbound,
196            Half::Conn => &self.conn,
197        });
198        if !slot
199            .as_ref()
200            .is_some_and(|registered| registered.will_wake(waker))
201        {
202            *slot = Some(waker.clone());
203        }
204    }
205}
206
207impl Wake for WriteWakers {
208    fn wake(self: Arc<Self>) {
209        self.wake_by_ref();
210    }
211
212    fn wake_by_ref(self: &Arc<Self>) {
213        for slot in [&self.inbound, &self.conn] {
214            if let Some(waker) = lock(slot).take() {
215                waker.wake();
216            }
217        }
218    }
219}
220
221fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
222    mutex.lock().unwrap_or_else(PoisonError::into_inner)
223}
224
225#[cfg(test)]
226mod tests {
227    use super::*;
228    use crate::driver::{
229        default_config,
230        test_support::{pipe, run_to_completion},
231    };
232    use futures_lite::{StreamExt, future::poll_fn};
233    use swansong::Swansong;
234    use tungstenite::protocol::Role;
235
236    /// While the conn waits to write, the taken stream answers pings, so both halves wait on the
237    /// same full socket. Both must be woken when it drains.
238    #[test]
239    fn both_halves_are_woken_to_write() {
240        const COUNT: usize = 32;
241        const LEN: usize = 1000;
242
243        let swansong = Swansong::new();
244        let (server, client) = pipe(512);
245        let mut conn = DriverSlot::Owned(swansong.interrupt(Driver::new(
246            Box::new(server),
247            Vec::new(),
248            Role::Server,
249            default_config().write_buffer_size(2048),
250        )));
251        let mut inbound = conn.split().unwrap();
252        let mut peer = swansong.interrupt(Driver::new(
253            Box::new(client),
254            Vec::new(),
255            Role::Client,
256            default_config(),
257        ));
258
259        let sending = async {
260            for i in 0..COUNT {
261                poll_fn(|cx| conn.poll_ready_to_write(cx)).await.unwrap();
262                conn.write(Message::binary(vec![i as u8; LEN])).unwrap();
263            }
264            poll_fn(|cx| conn.poll_flush(cx)).await.unwrap();
265        };
266
267        let reading = async {
268            while let Some(message) = inbound.next().await {
269                if message.unwrap() == Message::text("done") {
270                    break;
271                }
272            }
273        };
274
275        let peer = async {
276            for _ in 0..8 {
277                peer.write(Message::Ping("ping".into())).unwrap();
278            }
279            poll_fn(|cx| peer.poll_flush(cx)).await.unwrap();
280
281            let mut received = 0;
282            while received < COUNT {
283                let message = poll_fn(|cx| poll_inbound(&mut peer, cx)).await;
284                if message.unwrap().unwrap().is_binary() {
285                    received += 1;
286                }
287            }
288
289            peer.write(Message::text("done")).unwrap();
290            poll_fn(|cx| peer.poll_flush(cx)).await.unwrap();
291        };
292
293        assert!(
294            run_to_completion(vec![Box::pin(sending), Box::pin(reading), Box::pin(peer)]),
295            "stalled"
296        );
297    }
298}
trillium_websockets/websocket_connection/split_compat.rs

trillium_websockets/websocket_connection/
split_compat.rs
