unifly_api/websocket/
runtime.rs

1use std::sync::Arc;
2use std::time::Duration;
3
4use futures_util::StreamExt;
5use tokio::sync::broadcast;
6use tokio_tungstenite::Connector;
7use tokio_tungstenite::tungstenite::{self, ClientRequestBuilder};
8use tokio_util::sync::CancellationToken;
9use url::Url;
10
11use crate::error::Error;
12use crate::transport::TlsMode;
13
14use super::parser::{UnifiEvent, parse_and_broadcast};
15use super::tls::build_tls_connector;
16
17const EVENT_CHANNEL_CAPACITY: usize = 1024;
18
19/// Exponential backoff configuration for WebSocket reconnection.
20#[derive(Debug, Clone)]
21pub struct ReconnectConfig {
22    /// Delay before the first reconnection attempt. Default: 1s.
23    pub initial_delay: Duration,
24
25    /// Upper bound on backoff delay. Default: 30s.
26    pub max_delay: Duration,
27
28    /// Maximum reconnection attempts before giving up.
29    /// `None` means retry forever.
30    pub max_retries: Option<u32>,
31}
32
33impl Default for ReconnectConfig {
34    fn default() -> Self {
35        Self {
36            initial_delay: Duration::from_secs(1),
37            max_delay: Duration::from_secs(30),
38            max_retries: None,
39        }
40    }
41}
42
43/// Handle to a running WebSocket event stream.
44///
45/// Cheaply cloneable via the inner broadcast sender. Drop all handles
46/// and call [`shutdown`](Self::shutdown) to tear down the background task.
47pub struct WebSocketHandle {
48    event_rx: broadcast::Receiver<Arc<UnifiEvent>>,
49    cancel: CancellationToken,
50}
51
52impl WebSocketHandle {
53    /// Connect to the controller WebSocket and spawn the reconnection loop.
54    ///
55    /// Returns immediately once the background task is spawned.
56    /// The first connection attempt happens asynchronously -- subscribe to
57    /// the event receiver to start consuming events.
58    pub fn connect(
59        ws_url: Url,
60        reconnect: ReconnectConfig,
61        cancel: CancellationToken,
62        cookie: Option<String>,
63        tls_mode: TlsMode,
64    ) -> Result<Self, Error> {
65        let (event_tx, event_rx) = broadcast::channel(EVENT_CHANNEL_CAPACITY);
66
67        let task_cancel = cancel.clone();
68        tokio::spawn(async move {
69            ws_loop(ws_url, event_tx, reconnect, task_cancel, cookie, tls_mode).await;
70        });
71
72        Ok(Self { event_rx, cancel })
73    }
74
75    /// Get a new broadcast receiver for the event stream.
76    ///
77    /// Multiple consumers can subscribe concurrently. If a consumer falls
78    /// behind, it receives [`broadcast::error::RecvError::Lagged`].
79    pub fn subscribe(&self) -> broadcast::Receiver<Arc<UnifiEvent>> {
80        self.event_rx.resubscribe()
81    }
82
83    /// Signal the background task to shut down gracefully.
84    pub fn shutdown(&self) {
85        self.cancel.cancel();
86    }
87}
88
89async fn ws_loop(
90    ws_url: Url,
91    event_tx: broadcast::Sender<Arc<UnifiEvent>>,
92    reconnect: ReconnectConfig,
93    cancel: CancellationToken,
94    cookie: Option<String>,
95    tls_mode: TlsMode,
96) {
97    let mut attempt: u32 = 0;
98
99    loop {
100        tokio::select! {
101            biased;
102            () = cancel.cancelled() => break,
103            result = connect_and_read(&ws_url, &event_tx, &cancel, cookie.as_deref(), &tls_mode) => {
104                match result {
105                    Ok(()) => {
106                        tracing::info!("WebSocket disconnected cleanly, reconnecting");
107                        attempt = 0;
108                    }
109                    Err(error) => {
110                        tracing::warn!(error = %error, attempt, "WebSocket error");
111
112                        if let Some(max) = reconnect.max_retries
113                            && attempt >= max {
114                                tracing::error!(
115                                    max_retries = max,
116                                    "WebSocket reconnection limit reached, giving up"
117                                );
118                                break;
119                            }
120
121                        let delay = calculate_backoff(attempt, &reconnect);
122                        let delay_ms = u64::try_from(delay.as_millis()).unwrap_or(u64::MAX);
123                        tracing::info!(delay_ms, attempt, "Waiting before reconnect");
124
125                        tokio::select! {
126                            biased;
127                            () = cancel.cancelled() => break,
128                            () = tokio::time::sleep(delay) => {}
129                        }
130
131                        attempt += 1;
132                    }
133                }
134            }
135        }
136    }
137
138    #[allow(unreachable_code)]
139    {
140        tracing::debug!("WebSocket loop exiting");
141    }
142}
143
144async fn connect_and_read(
145    url: &Url,
146    event_tx: &broadcast::Sender<Arc<UnifiEvent>>,
147    cancel: &CancellationToken,
148    cookie: Option<&str>,
149    tls_mode: &TlsMode,
150) -> Result<(), Error> {
151    tracing::info!(url = %url, "Connecting to WebSocket");
152
153    let uri: tungstenite::http::Uri =
154        url.as_str()
155            .parse()
156            .map_err(|error: tungstenite::http::uri::InvalidUri| {
157                Error::WebSocketConnect(error.to_string())
158            })?;
159
160    let mut request = ClientRequestBuilder::new(uri);
161    if let Some(cookie_val) = cookie {
162        request = request.with_header("Cookie", cookie_val);
163    }
164
165    let connector = if url.scheme() == "wss" {
166        build_tls_connector(tls_mode)?
167    } else {
168        Some(Connector::Plain)
169    };
170
171    let (ws_stream, _response) =
172        tokio_tungstenite::connect_async_tls_with_config(request, None, false, connector)
173            .await
174            .map_err(|error| Error::WebSocketConnect(error.to_string()))?;
175
176    tracing::info!("WebSocket connected");
177
178    let (_write, mut read) = ws_stream.split();
179
180    loop {
181        tokio::select! {
182            biased;
183            () = cancel.cancelled() => return Ok(()),
184            frame = read.next() => {
185                match frame {
186                    Some(Ok(tungstenite::Message::Text(text))) => {
187                        parse_and_broadcast(&text, event_tx);
188                    }
189                    Some(Ok(tungstenite::Message::Ping(_))) => {
190                        tracing::trace!("WebSocket ping");
191                    }
192                    Some(Ok(tungstenite::Message::Close(frame))) => {
193                        if let Some(ref cf) = frame {
194                            tracing::info!(
195                                code = %cf.code,
196                                reason = %cf.reason,
197                                "WebSocket close frame received"
198                            );
199                        } else {
200                            tracing::info!("WebSocket close frame received (no payload)");
201                        }
202                        return Ok(());
203                    }
204                    Some(Err(error)) => {
205                        return Err(Error::WebSocketConnect(error.to_string()));
206                    }
207                    None => {
208                        tracing::info!("WebSocket stream ended");
209                        return Ok(());
210                    }
211                    _ => {}
212                }
213            }
214        }
215    }
216}
217
218fn calculate_backoff(attempt: u32, config: &ReconnectConfig) -> Duration {
219    let base = config.initial_delay.as_secs_f64()
220        * 2.0_f64.powi(i32::try_from(attempt).unwrap_or(i32::MAX));
221    let capped = base.min(config.max_delay.as_secs_f64());
222
223    let jitter_factor = 1.0 + 0.25 * ((f64::from(attempt) * 7.3).sin());
224    let with_jitter = (capped * jitter_factor).max(0.0);
225
226    Duration::from_secs_f64(with_jitter)
227}
228
229#[cfg(test)]
230#[allow(clippy::unwrap_used)]
231mod tests {
232    use super::*;
233
234    #[test]
235    fn default_reconnect_config() {
236        let config = ReconnectConfig::default();
237        assert_eq!(config.initial_delay, Duration::from_secs(1));
238        assert_eq!(config.max_delay, Duration::from_secs(30));
239        assert!(config.max_retries.is_none());
240    }
241
242    #[test]
243    fn backoff_increases_exponentially() {
244        let config = ReconnectConfig::default();
245
246        let d0 = calculate_backoff(0, &config);
247        let d1 = calculate_backoff(1, &config);
248        let d2 = calculate_backoff(2, &config);
249
250        assert!(d1 > d0, "d1 ({d1:?}) should be greater than d0 ({d0:?})");
251        assert!(d2 > d1, "d2 ({d2:?}) should be greater than d1 ({d1:?})");
252    }
253
254    #[test]
255    fn backoff_caps_at_max_delay() {
256        let config = ReconnectConfig {
257            initial_delay: Duration::from_secs(1),
258            max_delay: Duration::from_secs(10),
259            max_retries: None,
260        };
261
262        let d10 = calculate_backoff(10, &config);
263        assert!(
264            d10 <= Duration::from_secs(13),
265            "delay at attempt 10 ({d10:?}) should be capped near max_delay"
266        );
267    }
268}
unifly_api/websocket/runtime.rs

unifly_api/websocket/
runtime.rs
