1use 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#[derive(Clone)]
43pub struct SocketStateSink {
44 callback: Arc<dyn Fn(SocketState) + Send + Sync>,
45 transition_lock: Arc<Mutex<()>>,
46}
47
48impl SocketStateSink {
49 #[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 #[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#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
165pub enum SocketState {
166 Connected,
168 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}