brokapi_alor_protocol/ws/
session.rs

1//! The background session of the data channel `/ws`: connect, keep alive, reconnect and restore subscriptions.
2
3use std::collections::HashMap;
4use std::sync::Arc;
5use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
6use std::time::Duration;
7
8use futures_util::{SinkExt, StreamExt};
9use kanal::{AsyncReceiver, AsyncSender};
10use tokio::sync::Mutex;
11use tokio::time::{self, sleep};
12use tokio_tungstenite::connect_async;
13use tokio_tungstenite::tungstenite::Message;
14use tracing::{debug, error, info, warn};
15
16use crate::error::AlorApiError;
17use crate::http::client::AlorHttpClient;
18use crate::ws::models::{WS_PATH, WsIncoming, WsSubscribe, parse_incoming};
19
20/// How often a ping frame is sent to keep the connection alive.
21const WS_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(30);
22/// Default pause before the first reconnect attempt (the same as in the other connectors): 500 milliseconds.
23pub const DEFAULT_RECONNECT_MIN: Duration = Duration::from_millis(500);
24/// Default upper bound of the pause between reconnect attempts: 30 seconds.
25pub const DEFAULT_RECONNECT_MAX: Duration = Duration::from_secs(30);
26/// A connection that lived this long is healthy: the pause goes back to the minimum.
27const WS_HEALTHY_AFTER: Duration = Duration::from_secs(10);
28/// The most subscriptions one connection may hold (a limit of the broker).
29pub const WS_MAX_SUBSCRIPTIONS: usize = 5000;
30
31/// Options of a [`WsSession`].
32#[derive(Debug, Clone, Copy)]
33pub struct WsOptions {
34    /// Put [`WsIncoming::Reconnected`] into the stream at every connect, before the data of the restored
35    /// subscriptions. The connector uses it to see a lost connection.
36    pub announce_connects: bool,
37    /// Pause before the first reconnect, doubling up to `reconnect_max`.
38    pub reconnect_min: Duration,
39    /// Upper bound of the pause.
40    pub reconnect_max: Duration,
41}
42
43impl Default for WsOptions {
44    fn default() -> Self {
45        Self {
46            announce_connects: false,
47            reconnect_min: DEFAULT_RECONNECT_MIN,
48            reconnect_max: DEFAULT_RECONNECT_MAX,
49        }
50    }
51}
52
53/// A request on its way from the caller to the session task.
54struct Queued {
55    /// The number of a subscribe request in the restore list (it is sent with the list when the connection is
56    /// new, so the queued copy is not sent a second time).
57    seq: Option<u64>,
58    sub: WsSubscribe,
59}
60
61/// A WebSocket session of the data channel `/ws`.
62///
63/// It connects in the background, keeps the connection alive with ping frames, and when the connection drops it
64/// reconnects with growing pauses and sends the subscriptions again (each with a fresh access token). Frames
65/// arrive parsed on the channel given to [`WsSession::spawn`]. Dropping the session closes it.
66///
67/// # Examples
68///
69/// ```
70/// # #[tokio::main(flavor = "current_thread")]
71/// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
72/// # let server = brokapi_mocks::alor::MockServer::start(Default::default()).await;
73/// # let cfg = brokapi_alor_protocol::config::AlorConfig::for_mock("mock-refresh-token", &server.http_base);
74/// use brokapi_alor_protocol::http::client::AlorHttpClient;
75/// use brokapi_alor_protocol::ws::models::{WsIncoming, WsSubscribe};
76/// use brokapi_alor_protocol::ws::session::WsSession;
77///
78/// let http = AlorHttpClient::new(cfg)?;
79/// let (tx, rx) = kanal::bounded_async(64);
80/// let session = WsSession::spawn(http, tx);
81/// session.subscribe(WsSubscribe::quotes("q1", "MOEX", "SBER", Some("TQBR"))).await?;
82/// loop {
83///     if let WsIncoming::Data { guid, .. } = rx.recv().await? {
84///         assert_eq!(guid, "q1");
85///         break;
86///     }
87/// }
88/// # Ok(())
89/// # }
90/// ```
91pub struct WsSession {
92    user_tx: AsyncSender<Queued>,
93    subscriptions: Arc<Mutex<Vec<(u64, WsSubscribe)>>>,
94    next_seq: AtomicU64,
95    is_connected: Arc<AtomicBool>,
96    dropped_events: Arc<AtomicU64>,
97    handle: tokio::task::JoinHandle<()>,
98}
99
100impl Drop for WsSession {
101    fn drop(&mut self) {
102        self.handle.abort();
103    }
104}
105
106impl WsSession {
107    /// Starts a session with the default options. It returns at once and connects in the background.
108    pub fn spawn(http: AlorHttpClient, event_tx: AsyncSender<WsIncoming>) -> Self {
109        Self::spawn_with(http, event_tx, WsOptions::default())
110    }
111
112    /// Starts a session with the given options.
113    pub fn spawn_with(
114        http: AlorHttpClient,
115        event_tx: AsyncSender<WsIncoming>,
116        options: WsOptions,
117    ) -> Self {
118        let (user_tx, user_rx) = kanal::bounded_async::<Queued>(128);
119        let subscriptions = Arc::new(Mutex::new(Vec::new()));
120        let is_connected = Arc::new(AtomicBool::new(false));
121        let dropped_events = Arc::new(AtomicU64::new(0));
122        let handle = tokio::spawn(run_loop(
123            http,
124            event_tx,
125            options,
126            Arc::clone(&subscriptions),
127            user_rx,
128            Arc::clone(&is_connected),
129            Arc::clone(&dropped_events),
130        ));
131        Self {
132            user_tx,
133            subscriptions,
134            next_seq: AtomicU64::new(0),
135            is_connected,
136            dropped_events,
137            handle,
138        }
139    }
140
141    /// How many parsed frames were dropped because the receiving channel was full.
142    pub fn dropped_events(&self) -> u64 {
143        self.dropped_events.load(Ordering::Relaxed)
144    }
145
146    /// `true` while the connection is up and the subscriptions are sent.
147    pub fn is_connected(&self) -> bool {
148        self.is_connected.load(Ordering::Acquire)
149    }
150
151    /// Subscribes. The request is remembered and sent again after every reconnect until [`WsSession::unsubscribe`].
152    ///
153    /// # Errors
154    ///
155    /// [`AlorApiError::Subscription`] at 5000 subscriptions or when the session is closed.
156    pub async fn subscribe(&self, sub: WsSubscribe) -> Result<(), AlorApiError> {
157        let seq = self.next_seq.fetch_add(1, Ordering::Relaxed);
158        {
159            let mut subs = self.subscriptions.lock().await;
160            if subs.len() >= WS_MAX_SUBSCRIPTIONS {
161                return Err(AlorApiError::Subscription {
162                    message: format!(
163                        "subscription limit reached ({WS_MAX_SUBSCRIPTIONS} per connection)"
164                    ),
165                });
166            }
167            subs.push((seq, sub.clone()));
168        }
169        self.user_tx
170            .send(Queued {
171                seq: Some(seq),
172                sub,
173            })
174            .await
175            .map_err(|e| AlorApiError::Subscription {
176                message: format!("failed to send the subscription: {e}"),
177            })
178    }
179
180    /// Cancels the subscription named by `guid`: forgets it and sends `unsubscribe`.
181    ///
182    /// # Errors
183    ///
184    /// [`AlorApiError::Subscription`] when the session is closed.
185    pub async fn unsubscribe(&self, guid: &str) -> Result<(), AlorApiError> {
186        self.subscriptions
187            .lock()
188            .await
189            .retain(|(_, s)| s.guid != guid);
190        self.user_tx
191            .send(Queued {
192                seq: None,
193                sub: WsSubscribe::unsubscribe(guid),
194            })
195            .await
196            .map_err(|e| AlorApiError::Subscription {
197                message: format!("failed to send the unsubscription: {e}"),
198            })
199    }
200}
201
202async fn run_loop(
203    http: AlorHttpClient,
204    event_tx: AsyncSender<WsIncoming>,
205    options: WsOptions,
206    subscriptions: Arc<Mutex<Vec<(u64, WsSubscribe)>>>,
207    user_rx: AsyncReceiver<Queued>,
208    is_connected: Arc<AtomicBool>,
209    dropped_events: Arc<AtomicU64>,
210) {
211    let url = format!("{}{}", http.ws_base_url(), WS_PATH);
212    let backoff_min = options.reconnect_min.max(Duration::from_millis(10));
213    let backoff_max = options.reconnect_max.max(backoff_min);
214    let mut backoff = backoff_min;
215
216    'session: loop {
217        let ws_stream = match connect_async(url.as_str()).await {
218            Ok((s, res)) => {
219                info!(
220                    "WebSocket connected to {url} with status <{}>",
221                    res.status()
222                );
223                s
224            }
225            Err(e) => {
226                error!("WS connect failed: {e}");
227                sleep(backoff).await;
228                backoff = (backoff * 2).min(backoff_max);
229                continue;
230            }
231        };
232        let connected_at = tokio::time::Instant::now();
233        let (mut write, mut read) = ws_stream.split();
234
235        // The reader is told first: the data of the restored subscriptions comes after this marker.
236        if options.announce_connects && event_tx.send(WsIncoming::Reconnected).await.is_err() {
237            break 'session;
238        }
239
240        // Restore. The numbers sent here are remembered: the same request still in the queue is not sent twice.
241        let mut restored: HashMap<u64, ()> = HashMap::new();
242        {
243            let token = match http.get_access_token().await {
244                Ok(t) => t,
245                Err(e) => {
246                    error!("WS: no access token: {e}");
247                    sleep(backoff).await;
248                    backoff = (backoff * 2).min(backoff_max);
249                    continue 'session;
250                }
251            };
252            let subs = subscriptions.lock().await;
253            for (seq, sub) in subs.iter() {
254                if write
255                    .send(Message::Text(sub.to_json(&token).into()))
256                    .await
257                    .is_err()
258                {
259                    error!("WS: failed to restore a subscription");
260                    continue 'session;
261                }
262                restored.insert(*seq, ());
263            }
264            if !subs.is_empty() {
265                info!("Restored {} subscriptions", subs.len());
266            }
267        }
268        is_connected.store(true, Ordering::Release);
269
270        let mut heartbeat = time::interval(WS_HEARTBEAT_INTERVAL);
271        heartbeat.tick().await; // the first tick is at once
272
273        'inner: loop {
274            tokio::select! {
275                frame = read.next() => match frame {
276                    Some(Ok(Message::Text(text))) => {
277                        let incoming = parse_incoming(&text);
278                        // The broker closes the connection after a 401 and the cached token is no good: get a
279                        // new one now, the reconnect needs it.
280                        if matches!(incoming, WsIncoming::Status { code: 401, .. }) {
281                            warn!("WS: the broker refused the token, renewing it");
282                            let _ = http.refresh_access_token().await;
283                        }
284                        match event_tx.try_send(incoming) {
285                            Ok(true) => {}
286                            Ok(false) => {
287                                dropped_events.fetch_add(1, Ordering::Relaxed);
288                                warn!("WS: the receiving channel is full, a frame is dropped");
289                            }
290                            Err(_) => break 'session,
291                        }
292                    }
293                    Some(Ok(Message::Ping(data))) => {
294                        if write.send(Message::Pong(data)).await.is_err() {
295                            break 'inner;
296                        }
297                    }
298                    Some(Ok(Message::Close(_))) => {
299                        info!("WS closed by the server");
300                        break 'inner;
301                    }
302                    Some(Ok(_)) => {}
303                    Some(Err(e)) => {
304                        error!("WS read error: {e}");
305                        break 'inner;
306                    }
307                    None => break 'inner,
308                },
309                _ = heartbeat.tick() => {
310                    if write.send(Message::Ping(Vec::new().into())).await.is_err() {
311                        break 'inner;
312                    }
313                    debug!("WS ping sent");
314                }
315                msg = user_rx.recv() => match msg {
316                    Ok(Queued { seq, sub }) => {
317                        if seq.is_some_and(|n| restored.remove(&n).is_some()) {
318                            continue 'inner;
319                        }
320                        let token = match http.get_access_token().await {
321                            Ok(t) => t,
322                            Err(e) => {
323                                // The subscribe stays in the list, the reconnect sends it with a token.
324                                error!("WS: no access token for a request: {e}");
325                                break 'inner;
326                            }
327                        };
328                        if write.send(Message::Text(sub.to_json(&token).into())).await.is_err() {
329                            error!("WS write error");
330                            break 'inner;
331                        }
332                    }
333                    Err(_) => break 'session,
334                },
335            }
336        }
337
338        is_connected.store(false, Ordering::Release);
339        if connected_at.elapsed() >= WS_HEALTHY_AFTER {
340            backoff = backoff_min;
341        }
342        warn!("WS disconnected, reconnecting in {backoff:?}...");
343        sleep(backoff).await;
344        backoff = (backoff * 2).min(backoff_max);
345    }
346    is_connected.store(false, Ordering::Release);
347}
348
349#[cfg(test)]
350mod tests {
351    use super::*;
352
353    #[test]
354    fn default_pause_is_the_range_of_the_other_connectors() {
355        let o = WsOptions::default();
356        assert_eq!(
357            (o.reconnect_min, o.reconnect_max),
358            (Duration::from_millis(500), Duration::from_secs(30))
359        );
360        assert!(!o.announce_connects);
361    }
362}
brokapi_alor_protocol/ws/session.rs

brokapi_alor_protocol/ws/
session.rs
