titan_client/tcp/
connection_status.rs

1use std::sync::{Arc, RwLock};
2use tracing::error;
3
4/// Represents the current state of the TCP connection
5#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6pub enum ConnectionStatus {
7    /// Initial state or deliberately disconnected
8    Disconnected,
9    /// Currently attempting to connect
10    Connecting,
11    /// Successfully connected to the server
12    Connected,
13    /// Connection was lost, attempting to reconnect
14    Reconnecting,
15}
16
17/// Trait for any type that can receive connection status updates
18pub trait StatusListener: Send + Sync {
19    fn notify(&self, status: ConnectionStatus) -> bool;
20}
21
22// Implement for std mpsc::Sender
23impl StatusListener for std::sync::mpsc::Sender<ConnectionStatus> {
24    fn notify(&self, status: ConnectionStatus) -> bool {
25        self.send(status).is_ok()
26    }
27}
28
29// Implement for tokio mpsc::Sender
30impl StatusListener for tokio::sync::mpsc::Sender<ConnectionStatus> {
31    fn notify(&self, status: ConnectionStatus) -> bool {
32        self.try_send(status).is_ok()
33    }
34}
35
36/// Thread-safe wrapper for tracking and updating connection status
37#[derive(Clone)]
38pub struct ConnectionStatusTracker {
39    status: Arc<RwLock<ConnectionStatus>>,
40    listeners: Arc<RwLock<Vec<Box<dyn StatusListener>>>>,
41}
42
43impl ConnectionStatusTracker {
44    /// Creates a new ConnectionStatusTracker with the initial status set to Disconnected
45    pub fn new() -> Self {
46        Self {
47            status: Arc::new(RwLock::new(ConnectionStatus::Disconnected)),
48            listeners: Arc::new(RwLock::new(Vec::new())),
49        }
50    }
51
52    /// Creates a new ConnectionStatusTracker with the specified initial status
53    pub fn with_status(initial_status: ConnectionStatus) -> Self {
54        Self {
55            status: Arc::new(RwLock::new(initial_status)),
56            listeners: Arc::new(RwLock::new(Vec::new())),
57        }
58    }
59
60    /// Get the current connection status
61    pub fn get_status(&self) -> ConnectionStatus {
62        match self.status.read() {
63            Ok(status) => *status,
64            Err(_) => {
65                // If the lock is poisoned, default to Disconnected
66                error!("Failed to read connection status due to poisoned lock");
67                ConnectionStatus::Disconnected
68            }
69        }
70    }
71
72    /// Register a listener to receive status updates
73    pub fn register_listener<L: StatusListener + 'static>(&self, listener: L) {
74        if let Ok(mut listeners) = self.listeners.write() {
75            listeners.push(Box::new(listener));
76        } else {
77            error!("Failed to register status listener due to poisoned lock");
78        }
79    }
80
81    /// Update the connection status and notify all listeners
82    pub fn update_status(&self, new_status: ConnectionStatus) {
83        if let Ok(mut status_guard) = self.status.write() {
84            *status_guard = new_status;
85
86            // Notify all listeners
87            if let Ok(mut listeners) = self.listeners.write() {
88                // Keep only listeners that successfully received the notification
89                listeners.retain(|listener| listener.notify(new_status));
90            }
91        } else {
92            error!("Failed to update connection status due to poisoned lock");
93        }
94    }
95
96    /// Create a helper closure that can be used to update the status
97    pub fn create_updater<'a>(&'a self) -> impl Fn(ConnectionStatus) + 'a {
98        let status = Arc::clone(&self.status);
99        let listeners = Arc::<RwLock<Vec<Box<dyn StatusListener>>>>::clone(&self.listeners);
100
101        move |new_status| {
102            if let Ok(mut status_guard) = status.write() {
103                *status_guard = new_status;
104
105                // Notify all listeners
106                if let Ok(mut listeners_guard) = listeners.write() {
107                    listeners_guard.retain(|listener| listener.notify(new_status));
108                }
109            } else {
110                error!("Failed to update connection status due to poisoned lock");
111            }
112        }
113    }
114
115    /// Get a clone of the inner Arc<RwLock<ConnectionStatus>>
116    /// This is useful when you need to share the status across threads
117    pub fn get_inner(&self) -> Arc<RwLock<ConnectionStatus>> {
118        Arc::clone(&self.status)
119    }
120}
121
122impl Default for ConnectionStatusTracker {
123    fn default() -> Self {
124        Self::new()
125    }
126}
titan_client/tcp/connection_status.rs

titan_client/tcp/
connection_status.rs
