Skip to main content

nautilus_network/
sink.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! Ordered availability-edge publication for socket transports.
17//!
18//! # Ordering
19//!
20//! A [`SocketStateSink`] reports transitions into and out of [`ConnectionMode::Active`] for one
21//! client, rather than every internal connection mode. Most paths perform the mode transition and
22//! callback under the same serialization lock, so concurrent loss and recovery attempts publish
23//! at most one ordered edge. WebSocket reconnect changes mode before publishing loss and uses the
24//! same lock for publication, while the controller prevents recovery until publication completes.
25
26use std::{
27    fmt::Debug,
28    sync::{
29        Arc,
30        atomic::{AtomicU8, Ordering},
31    },
32};
33
34use parking_lot::Mutex;
35
36use crate::mode::ConnectionMode;
37
38/// Receives ordered semantic state changes from a single socket client.
39///
40/// Clients must route every transition into or out of [`ConnectionMode::Active`] through the
41/// sink-backed transition methods so each edge is reported once.
42#[derive(Clone)]
43pub struct SocketStateSink {
44    callback: Arc<dyn Fn(SocketState) + Send + Sync>,
45    transition_lock: Arc<Mutex<()>>,
46}
47
48impl SocketStateSink {
49    /// Creates a new [`SocketStateSink`] instance.
50    ///
51    /// The callback runs synchronously with each successful state transition and should return
52    /// promptly. It must not initiate another transition using the same sink because callbacks are
53    /// serialized under a non-reentrant lock.
54    #[must_use]
55    pub fn new<F>(callback: F) -> Self
56    where
57        F: Fn(SocketState) + Send + Sync + 'static,
58    {
59        Self {
60            callback: Arc::new(callback),
61            transition_lock: Arc::new(Mutex::new(())),
62        }
63    }
64
65    /// Returns a sink that invokes `callback` before forwarding each state to this sink.
66    #[must_use]
67    pub fn with_callback<F>(self, callback: F) -> Self
68    where
69        F: Fn(SocketState) + Send + Sync + 'static,
70    {
71        let Self {
72            callback: forwarded,
73            transition_lock,
74        } = self;
75
76        Self {
77            callback: Arc::new(move |state| {
78                callback(state);
79                forwarded(state);
80            }),
81            transition_lock,
82        }
83    }
84
85    pub(crate) fn transition(
86        &self,
87        value: &AtomicU8,
88        current: ConnectionMode,
89        next: ConnectionMode,
90        state: SocketState,
91    ) -> bool {
92        self.transition_result(value, current, next, state).is_ok()
93    }
94
95    pub(crate) fn transition_result(
96        &self,
97        value: &AtomicU8,
98        current: ConnectionMode,
99        next: ConnectionMode,
100        state: SocketState,
101    ) -> Result<(), ConnectionMode> {
102        let _guard = self.transition_lock.lock();
103
104        if let Err(actual) = value.compare_exchange(
105            current.as_u8(),
106            next.as_u8(),
107            Ordering::SeqCst,
108            Ordering::SeqCst,
109        ) {
110            return Err(ConnectionMode::from_u8(actual));
111        }
112
113        self.notify(state);
114
115        Ok(())
116    }
117
118    pub(crate) fn publish_websocket(&self, state: SocketState) {
119        let _guard = self.transition_lock.lock();
120        self.notify(state);
121    }
122
123    pub(crate) fn close_on_loss(&self, value: &AtomicU8) -> bool {
124        let _guard = self.transition_lock.lock();
125        let current = ConnectionMode::from_atomic(value);
126
127        if !matches!(current, ConnectionMode::Active | ConnectionMode::Reconnect)
128            || value
129                .compare_exchange(
130                    current.as_u8(),
131                    ConnectionMode::Closed.as_u8(),
132                    Ordering::SeqCst,
133                    Ordering::SeqCst,
134                )
135                .is_err()
136        {
137            return false;
138        }
139
140        if current.is_active() {
141            self.notify(SocketState::Disconnected);
142        }
143
144        true
145    }
146
147    fn notify(&self, state: SocketState) {
148        if std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| (self.callback)(state)))
149            .is_err()
150        {
151            log::error!("Socket state sink panicked while handling {state:?}");
152        }
153    }
154}
155
156impl Debug for SocketStateSink {
157    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
158        f.debug_struct(stringify!(SocketStateSink))
159            .finish_non_exhaustive()
160    }
161}
162
163/// Represents the availability state reported by a socket transport.
164#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
165pub enum SocketState {
166    /// The transport is available.
167    Connected,
168    /// An active transport was lost.
169    Disconnected,
170}
171
172#[cfg(test)]
173mod tests {
174    use std::sync::{
175        Barrier,
176        atomic::{AtomicU8, AtomicUsize, Ordering as AtomicOrdering},
177    };
178
179    use parking_lot::Mutex;
180    use rstest::rstest;
181
182    use super::*;
183    use crate::mode::{ConnectionMode, ReconnectOutcome};
184
185    #[rstest]
186    fn state_sink_reports_only_successful_edges_in_order() {
187        let states = Arc::new(Mutex::new(Vec::new()));
188        let states_callback = Arc::clone(&states);
189        let sink = SocketStateSink::new(move |state| {
190            states_callback.lock().push(state);
191        });
192
193        let mode = AtomicU8::new(ConnectionMode::Reconnect.as_u8());
194
195        assert_eq!(
196            ConnectionMode::complete_reconnect_with_sink(&mode, Some(&sink)),
197            ReconnectOutcome::Reconnected
198        );
199        assert!(ConnectionMode::request_reconnect_with_sink(
200            &mode,
201            Some(&sink)
202        ));
203        assert!(!ConnectionMode::request_reconnect_with_sink(
204            &mode,
205            Some(&sink)
206        ));
207        assert_eq!(
208            ConnectionMode::complete_reconnect_with_sink(&mode, Some(&sink)),
209            ReconnectOutcome::Reconnected
210        );
211
212        assert_eq!(
213            *states.lock(),
214            vec![
215                SocketState::Connected,
216                SocketState::Disconnected,
217                SocketState::Connected,
218            ]
219        );
220    }
221
222    #[rstest]
223    fn state_sink_with_callback_runs_before_forwarded_sink() {
224        let calls = Arc::new(Mutex::new(Vec::new()));
225        let forwarded_calls = Arc::clone(&calls);
226        let sink = SocketStateSink::new(move |_| {
227            forwarded_calls.lock().push("forwarded");
228        });
229
230        let transition_lock = Arc::clone(&sink.transition_lock);
231        let callback_calls = Arc::clone(&calls);
232        let sink = sink.with_callback(move |_| {
233            callback_calls.lock().push("callback");
234        });
235
236        let mode = AtomicU8::new(ConnectionMode::Reconnect.as_u8());
237
238        assert_eq!(
239            ConnectionMode::complete_reconnect_with_sink(&mode, Some(&sink)),
240            ReconnectOutcome::Reconnected
241        );
242        assert!(Arc::ptr_eq(&sink.transition_lock, &transition_lock));
243        assert_eq!(*calls.lock(), vec!["callback", "forwarded"]);
244    }
245
246    #[rstest]
247    fn state_sink_reports_one_concurrent_loss() {
248        let states = Arc::new(Mutex::new(Vec::new()));
249        let states_callback = Arc::clone(&states);
250        let sink = SocketStateSink::new(move |state| {
251            states_callback.lock().push(state);
252        });
253
254        let mode = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
255        let barrier = Arc::new(Barrier::new(8));
256
257        let mut transitions = Vec::with_capacity(8);
258
259        for _ in 0..8 {
260            let mode = Arc::clone(&mode);
261            let sink = sink.clone();
262            let barrier = Arc::clone(&barrier);
263            transitions.push(std::thread::spawn(move || {
264                barrier.wait();
265                ConnectionMode::request_reconnect_with_sink(&mode, Some(&sink))
266            }));
267        }
268
269        let successful = transitions
270            .into_iter()
271            .map(|transition| transition.join().unwrap())
272            .filter(|successful| *successful)
273            .count();
274
275        assert_eq!(successful, 1);
276        assert_eq!(
277            ConnectionMode::from_atomic(&mode),
278            ConnectionMode::Reconnect
279        );
280        assert_eq!(*states.lock(), vec![SocketState::Disconnected]);
281    }
282
283    #[rstest]
284    fn state_sink_reports_one_mixed_concurrent_loss() {
285        let states = Arc::new(Mutex::new(Vec::new()));
286        let states_callback = Arc::clone(&states);
287        let sink = SocketStateSink::new(move |state| {
288            states_callback.lock().push(state);
289        });
290
291        let mode = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
292        let barrier = Arc::new(Barrier::new(2));
293
294        let reconnect = {
295            let mode = Arc::clone(&mode);
296            let sink = sink.clone();
297            let barrier = Arc::clone(&barrier);
298            std::thread::spawn(move || {
299                barrier.wait();
300                ConnectionMode::request_reconnect_with_sink(&mode, Some(&sink))
301            })
302        };
303
304        let close = std::thread::spawn({
305            let mode = Arc::clone(&mode);
306            move || {
307                barrier.wait();
308                ConnectionMode::close_websocket_on_loss(&mode, Some(&sink))
309            }
310        });
311
312        reconnect.join().unwrap();
313        let closed = close.join().unwrap();
314
315        assert!(closed);
316        assert_eq!(ConnectionMode::from_atomic(&mode), ConnectionMode::Closed);
317        assert_eq!(*states.lock(), vec![SocketState::Disconnected]);
318    }
319
320    #[rstest]
321    fn state_sink_continues_after_callback_panic() {
322        let calls = Arc::new(AtomicUsize::new(0));
323        let calls_callback = Arc::clone(&calls);
324        let states = Arc::new(Mutex::new(Vec::new()));
325        let states_callback = Arc::clone(&states);
326
327        let sink = SocketStateSink::new(move |state| {
328            assert_ne!(
329                calls_callback.fetch_add(1, AtomicOrdering::SeqCst),
330                0,
331                "test socket state callback panic"
332            );
333            states_callback.lock().push(state);
334        });
335
336        let mode = AtomicU8::new(ConnectionMode::Reconnect.as_u8());
337
338        assert_eq!(
339            ConnectionMode::complete_reconnect_with_sink(&mode, Some(&sink)),
340            ReconnectOutcome::Reconnected
341        );
342        assert!(ConnectionMode::request_reconnect_with_sink(
343            &mode,
344            Some(&sink)
345        ));
346
347        assert_eq!(calls.load(AtomicOrdering::SeqCst), 2);
348        assert_eq!(
349            ConnectionMode::from_atomic(&mode),
350            ConnectionMode::Reconnect
351        );
352        assert_eq!(*states.lock(), vec![SocketState::Disconnected]);
353    }
354
355    #[rstest]
356    fn state_sink_suppresses_deliberate_disconnect() {
357        let states = Arc::new(Mutex::new(Vec::new()));
358        let states_callback = Arc::clone(&states);
359        let sink = SocketStateSink::new(move |state| {
360            states_callback.lock().push(state);
361        });
362
363        let mode = AtomicU8::new(ConnectionMode::Active.as_u8());
364
365        assert!(ConnectionMode::request_disconnect(&mode));
366        assert!(!ConnectionMode::request_reconnect_with_sink(
367            &mode,
368            Some(&sink)
369        ));
370
371        assert_eq!(*states.lock(), Vec::new());
372    }
373
374    #[rstest]
375    fn state_sink_closes_after_reported_loss_without_another_event() {
376        let states = Arc::new(Mutex::new(Vec::new()));
377        let states_callback = Arc::clone(&states);
378        let sink = SocketStateSink::new(move |state| {
379            states_callback.lock().push(state);
380        });
381
382        let mode = AtomicU8::new(ConnectionMode::Reconnect.as_u8());
383
384        assert!(ConnectionMode::close_websocket_on_loss(&mode, Some(&sink)));
385
386        assert_eq!(ConnectionMode::from_atomic(&mode), ConnectionMode::Closed);
387        assert_eq!(*states.lock(), Vec::new());
388    }
389}