nautilus_network/
mode.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//! Atomic connection state and controller lifecycle coordination for socket transports.
17//!
18//! # Transition contract
19//!
20//! [`ConnectionMode`] is shared across transport tasks. Reconnect transitions use atomic
21//! compare-and-exchange operations so late reconnect work cannot overwrite a concurrent
22//! `Disconnect` or `Closed` state. Sink-backed transitions pair each successful mode change with
23//! its semantic availability edge.
24//!
25//! # Session and controller fencing
26//!
27//! `ReadSessionFence` marks a reader as retired so its dispatch checks drop old-transport messages
28//! after observing invalidation. `ControllerLifecycle` prevents retained reconnect handles from
29//! accepting work after shutdown and defers aborting the controller until each in-flight request
30//! reaches its handoff boundary.
31
32use std::sync::{
33    Arc, OnceLock,
34    atomic::{AtomicBool, AtomicU8, AtomicUsize, Ordering},
35};
36
37use strum::{AsRefStr, Display, EnumString};
38
39use crate::sink::{SocketState, SocketStateSink};
40
41/// The lifecycle state of a socket client.
42///
43/// Clients store the active, reconnecting, disconnecting, or closed state in an atomic flag so
44/// transport tasks can coordinate lifecycle transitions across threads.
45#[derive(Clone, Copy, Debug, Default, Display, Hash, PartialEq, Eq, AsRefStr, EnumString)]
46#[repr(u8)]
47#[strum(serialize_all = "UPPERCASE")]
48pub enum ConnectionMode {
49    #[default]
50    /// The client is fully connected and operational.
51    /// All tasks (reading, writing, heartbeat) are running normally.
52    Active = 0,
53    /// The client has been disconnected or has been explicitly signaled to reconnect.
54    /// In this state, active tasks are paused until a new connection is established.
55    Reconnect = 1,
56    /// The client has been explicitly signaled to disconnect.
57    /// No further reconnection attempts will be made, and cleanup procedures are initiated.
58    Disconnect = 2,
59    /// The client is permanently closed.
60    /// All associated tasks have been terminated and the connection is no longer available.
61    Closed = 3,
62}
63
64impl ConnectionMode {
65    /// Converts a `u8` loaded from an [`AtomicU8`] into a [`ConnectionMode`].
66    ///
67    /// # Panics
68    ///
69    /// Panics if `value` is not a valid `ConnectionMode` discriminant (must be between 0 and 3 inclusive).
70    #[inline]
71    #[must_use]
72    pub fn from_u8(value: u8) -> Self {
73        match value {
74            0 => Self::Active,
75            1 => Self::Reconnect,
76            2 => Self::Disconnect,
77            3 => Self::Closed,
78            _ => panic!("Invalid `ConnectionMode` value: {value}"),
79        }
80    }
81
82    /// Loads a [`ConnectionMode`] from an [`AtomicU8`] using sequential consistency.
83    #[inline]
84    #[must_use]
85    pub fn from_atomic(value: &AtomicU8) -> Self {
86        Self::from_u8(value.load(Ordering::SeqCst))
87    }
88
89    /// Atomically transitions to `Reconnect`, but only from `Active`.
90    ///
91    /// Returns `true` if this call performed the transition. A concurrent
92    /// `Disconnect`/`Closed` (or an in-flight `Reconnect`) is left untouched,
93    /// so a writer detecting a dead connection cannot resurrect a client that
94    /// is being torn down.
95    pub fn request_reconnect(value: &AtomicU8) -> bool {
96        Self::request_reconnect_outcome(value) == ReconnectRequestOutcome::Accepted
97    }
98
99    /// Atomically requests reconnect and reports the observed state on rejection.
100    pub(crate) fn request_reconnect_outcome(value: &AtomicU8) -> ReconnectRequestOutcome {
101        match value.compare_exchange(
102            Self::Active.as_u8(),
103            Self::Reconnect.as_u8(),
104            Ordering::SeqCst,
105            Ordering::SeqCst,
106        ) {
107            Ok(_) => ReconnectRequestOutcome::Accepted,
108            Err(actual) => ReconnectRequestOutcome::from_rejected(Self::from_u8(actual)),
109        }
110    }
111
112    /// Atomically transitions from `Active` to `Reconnect` and reports the loss.
113    pub(crate) fn request_reconnect_with_sink(
114        value: &AtomicU8,
115        sink: Option<&SocketStateSink>,
116    ) -> bool {
117        Self::request_reconnect_outcome_with_sink(value, sink) == ReconnectRequestOutcome::Accepted
118    }
119
120    pub(crate) fn request_reconnect_outcome_with_sink(
121        value: &AtomicU8,
122        sink: Option<&SocketStateSink>,
123    ) -> ReconnectRequestOutcome {
124        sink.map_or_else(
125            || Self::request_reconnect_outcome(value),
126            |sink| {
127                sink.transition_result(
128                    value,
129                    Self::Active,
130                    Self::Reconnect,
131                    SocketState::Disconnected,
132                )
133                .map_or_else(ReconnectRequestOutcome::from_rejected, |()| {
134                    ReconnectRequestOutcome::Accepted
135                })
136            },
137        )
138    }
139
140    /// Atomically transitions from `Active` or `Reconnect` to `Closed` using WebSocket callback
141    /// serialization.
142    pub(crate) fn close_websocket_on_loss(
143        value: &AtomicU8,
144        sink: Option<&SocketStateSink>,
145    ) -> bool {
146        sink.map_or_else(
147            || {
148                value
149                    .try_update(Ordering::SeqCst, Ordering::SeqCst, |mode| {
150                        matches!(Self::from_u8(mode), Self::Active | Self::Reconnect)
151                            .then_some(Self::Closed.as_u8())
152                    })
153                    .is_ok()
154            },
155            |sink| sink.close_on_loss(value),
156        )
157    }
158
159    /// Atomically transitions to `Disconnect` from any non-`Closed` state.
160    ///
161    /// Returns `true` if the mode is now `Disconnect`; `false` if the
162    /// connection was already `Closed` (terminal state is preserved so status
163    /// queries keep reporting `Closed`).
164    pub fn request_disconnect(value: &AtomicU8) -> bool {
165        value
166            .try_update(Ordering::SeqCst, Ordering::SeqCst, |mode| {
167                (!Self::from_u8(mode).is_closed()).then_some(Self::Disconnect.as_u8())
168            })
169            .is_ok()
170    }
171
172    /// Atomically completes a reconnection by transitioning `Reconnect` to `Active`.
173    ///
174    /// Returns [`ReconnectOutcome::Reconnected`] if this call performed the
175    /// transition, and [`ReconnectOutcome::Aborted`] if the mode had already moved
176    /// on - which a concurrent [`Self::request_disconnect`] or a terminal `Closed`
177    /// store can do while the reconnect is in flight. Aborting here is not an
178    /// error: it means a teardown won the race and the replacement connection must
179    /// not be adopted.
180    pub(crate) fn complete_reconnect(value: &AtomicU8) -> ReconnectOutcome {
181        if value
182            .compare_exchange(
183                Self::Reconnect.as_u8(),
184                Self::Active.as_u8(),
185                Ordering::SeqCst,
186                Ordering::SeqCst,
187            )
188            .is_ok()
189        {
190            ReconnectOutcome::Reconnected
191        } else {
192            ReconnectOutcome::Aborted
193        }
194    }
195
196    /// Atomically transitions from `Reconnect` to `Active` and reports availability.
197    pub(crate) fn complete_reconnect_with_sink(
198        value: &AtomicU8,
199        sink: Option<&SocketStateSink>,
200    ) -> ReconnectOutcome {
201        let reconnected = sink.map_or_else(
202            || Self::complete_reconnect(value) == ReconnectOutcome::Reconnected,
203            |sink| sink.transition(value, Self::Reconnect, Self::Active, SocketState::Connected),
204        );
205
206        if reconnected {
207            ReconnectOutcome::Reconnected
208        } else {
209            ReconnectOutcome::Aborted
210        }
211    }
212
213    /// Converts a [`ConnectionMode`] to its `u8` representation.
214    #[inline]
215    #[must_use]
216    pub const fn as_u8(self) -> u8 {
217        self as u8
218    }
219
220    /// Returns true if the client is in an active state.
221    #[inline]
222    #[must_use]
223    pub const fn is_active(&self) -> bool {
224        matches!(self, Self::Active)
225    }
226
227    /// Returns true if the client is attempting to reconnect.
228    #[inline]
229    #[must_use]
230    pub const fn is_reconnect(&self) -> bool {
231        matches!(self, Self::Reconnect)
232    }
233
234    /// Returns true if the client is attempting to disconnect.
235    #[inline]
236    #[must_use]
237    pub const fn is_disconnect(&self) -> bool {
238        matches!(self, Self::Disconnect)
239    }
240
241    /// Returns true if the client connection is closed.
242    #[inline]
243    #[must_use]
244    pub const fn is_closed(&self) -> bool {
245        matches!(self, Self::Closed)
246    }
247}
248
249/// Outcome of a controller-owned reconnect request.
250#[derive(Clone, Copy, Debug, Eq, PartialEq)]
251pub enum ReconnectRequestOutcome {
252    /// The active transport entered reconnect mode.
253    Accepted,
254    /// The transport is already reconnecting.
255    AlreadyReconnecting,
256    /// The transport is disconnecting.
257    Disconnected,
258    /// The transport is permanently closed.
259    Closed,
260    /// The client uses stream mode and cannot replace its caller-owned reader.
261    Unsupported,
262}
263
264impl ReconnectRequestOutcome {
265    fn from_rejected(mode: ConnectionMode) -> Self {
266        match mode {
267            ConnectionMode::Active | ConnectionMode::Reconnect => Self::AlreadyReconnecting,
268            ConnectionMode::Disconnect => Self::Disconnected,
269            ConnectionMode::Closed => Self::Closed,
270        }
271    }
272}
273
274/// Result of a reconnection attempt that did not fail outright.
275///
276/// A reconnect can finish without reconnecting: a teardown may be requested while
277/// it is in flight, at which point it unwinds and leaves the mode terminal. That is
278/// normal control flow rather than an error, so it needs to be distinguishable from
279/// a completed reconnection by the caller.
280#[derive(Clone, Copy, Debug, Eq, PartialEq)]
281pub(crate) enum ReconnectOutcome {
282    /// A replacement connection was established and the mode is now `Active`.
283    Reconnected,
284    /// The attempt unwound without reconnecting; the mode was left unchanged.
285    Aborted,
286}
287
288/// Irreversible validity token for a single connection's read task.
289#[derive(Clone, Debug)]
290pub(crate) struct ReadSessionFence {
291    valid: Arc<AtomicBool>,
292}
293
294impl ReadSessionFence {
295    /// Creates a valid fence for a newly spawned read task.
296    #[must_use]
297    pub(crate) fn new() -> Self {
298        Self {
299            valid: Arc::new(AtomicBool::new(true)),
300        }
301    }
302
303    /// Invalidates the associated read session.
304    pub(crate) fn invalidate(&self) {
305        self.valid.store(false, Ordering::SeqCst);
306    }
307
308    /// Returns whether the associated read session is still current.
309    #[must_use]
310    pub(crate) fn is_valid(&self) -> bool {
311        self.valid.load(Ordering::SeqCst)
312    }
313}
314
315const CONTROLLER_CLOSED: usize = 1 << (usize::BITS - 1);
316const CONTROLLER_REQUEST_MASK: usize = CONTROLLER_CLOSED - 1;
317
318pub(crate) struct ControllerLifecycle {
319    state: AtomicUsize,
320    abort_handle: OnceLock<tokio::task::AbortHandle>,
321}
322
323impl ControllerLifecycle {
324    pub(crate) const fn new() -> Self {
325        Self {
326            state: AtomicUsize::new(0),
327            abort_handle: OnceLock::new(),
328        }
329    }
330
331    pub(crate) fn enter_request(&self) -> Option<ControllerRequest<'_>> {
332        self.state
333            .try_update(Ordering::SeqCst, Ordering::SeqCst, |state| {
334                if state & CONTROLLER_CLOSED != 0 {
335                    None
336                } else {
337                    assert_ne!(
338                        state, CONTROLLER_REQUEST_MASK,
339                        "too many reconnect requests"
340                    );
341                    Some(state + 1)
342                }
343            })
344            .ok()
345            .map(|_| ControllerRequest(self))
346    }
347
348    pub(crate) fn set_abort_handle(&self, abort_handle: tokio::task::AbortHandle) {
349        assert!(
350            self.abort_handle.set(abort_handle).is_ok(),
351            "controller abort handle already set"
352        );
353    }
354
355    pub(crate) fn close_and_abort(&self) {
356        let previous = self.state.fetch_or(CONTROLLER_CLOSED, Ordering::SeqCst);
357        if previous & CONTROLLER_REQUEST_MASK == 0 {
358            self.abort();
359        }
360    }
361
362    pub(crate) fn activity(self: &Arc<Self>) -> ControllerActivity {
363        ControllerActivity(Arc::clone(self))
364    }
365
366    fn close(&self) {
367        self.state.fetch_or(CONTROLLER_CLOSED, Ordering::SeqCst);
368    }
369
370    fn abort(&self) {
371        if let Some(abort_handle) = self.abort_handle.get() {
372            abort_handle.abort();
373        }
374    }
375}
376
377pub(crate) struct ControllerRequest<'a>(&'a ControllerLifecycle);
378
379impl Drop for ControllerRequest<'_> {
380    fn drop(&mut self) {
381        let previous = self.0.state.fetch_sub(1, Ordering::SeqCst);
382        if previous == CONTROLLER_CLOSED | 1 {
383            self.0.abort();
384        }
385    }
386}
387
388pub(crate) struct ControllerActivity(Arc<ControllerLifecycle>);
389
390impl Drop for ControllerActivity {
391    fn drop(&mut self) {
392        self.0.close();
393    }
394}
395
396#[cfg(test)]
397mod tests {
398    use rstest::rstest;
399
400    use super::*;
401
402    #[rstest]
403    #[case(ConnectionMode::Active, true, ConnectionMode::Reconnect)]
404    #[case(ConnectionMode::Reconnect, false, ConnectionMode::Reconnect)]
405    #[case(ConnectionMode::Disconnect, false, ConnectionMode::Disconnect)]
406    #[case(ConnectionMode::Closed, false, ConnectionMode::Closed)]
407    fn request_reconnect_transitions(
408        #[case] start: ConnectionMode,
409        #[case] expected_result: bool,
410        #[case] expected_mode: ConnectionMode,
411    ) {
412        let mode = AtomicU8::new(start.as_u8());
413
414        assert_eq!(ConnectionMode::request_reconnect(&mode), expected_result);
415        assert_eq!(ConnectionMode::from_atomic(&mode), expected_mode);
416    }
417
418    // Only a mode still in `Reconnect` may be adopted as `Active`. A teardown that won
419    // the race - `Disconnect` or `Closed` - must report `Aborted` and leave the terminal
420    // mode intact, so the caller can tell a completed reconnection from an unwound one.
421    #[rstest]
422    #[case(
423        ConnectionMode::Reconnect,
424        ReconnectOutcome::Reconnected,
425        ConnectionMode::Active
426    )]
427    #[case(
428        ConnectionMode::Disconnect,
429        ReconnectOutcome::Aborted,
430        ConnectionMode::Disconnect
431    )]
432    #[case(
433        ConnectionMode::Closed,
434        ReconnectOutcome::Aborted,
435        ConnectionMode::Closed
436    )]
437    #[case(
438        ConnectionMode::Active,
439        ReconnectOutcome::Aborted,
440        ConnectionMode::Active
441    )]
442    fn complete_reconnect_transitions(
443        #[case] start: ConnectionMode,
444        #[case] expected_outcome: ReconnectOutcome,
445        #[case] expected_mode: ConnectionMode,
446    ) {
447        let mode = AtomicU8::new(start.as_u8());
448
449        assert_eq!(ConnectionMode::complete_reconnect(&mode), expected_outcome);
450        assert_eq!(ConnectionMode::from_atomic(&mode), expected_mode);
451    }
452
453    #[rstest]
454    #[case(ConnectionMode::Active, true, ConnectionMode::Disconnect)]
455    #[case(ConnectionMode::Reconnect, true, ConnectionMode::Disconnect)]
456    #[case(ConnectionMode::Disconnect, true, ConnectionMode::Disconnect)]
457    #[case(ConnectionMode::Closed, false, ConnectionMode::Closed)]
458    fn request_disconnect_transitions(
459        #[case] start: ConnectionMode,
460        #[case] expected_result: bool,
461        #[case] expected_mode: ConnectionMode,
462    ) {
463        let mode = AtomicU8::new(start.as_u8());
464
465        assert_eq!(ConnectionMode::request_disconnect(&mode), expected_result);
466        assert_eq!(ConnectionMode::from_atomic(&mode), expected_mode);
467    }
468}
nautilus_network/mode.rs

nautilus_network/
mode.rs
