brokapi_bybit_protocol/ws/
session.rs

1//! A background WebSocket session that reconnects and restores its subscriptions.
2//!
3//! One session is one connection: public data, or the private streams (logged in with the key). The session sends a
4//! heartbeat every 20 seconds as the exchange asks, reconnects when the link breaks or goes silent, logs in again,
5//! and sends its subscriptions again. With `announce_connects` it tells the reader about every connection with
6//! [`WsIncoming::Connected`] before the restored data, so the reader can tell a reconnect from the first connect.
7
8use std::sync::Arc;
9use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
10use std::time::Duration;
11
12use futures_util::{SinkExt, StreamExt};
13use kanal::{AsyncReceiver, AsyncSender};
14use tokio::sync::Mutex;
15use tokio::time::{self, Instant, sleep};
16use tokio_tungstenite::connect_async;
17use tokio_tungstenite::tungstenite::Message;
18use tracing::{debug, error, info, warn};
19
20use super::models::{WsIncoming, WsRequest, parse_incoming};
21use crate::error::BybitApiError;
22use crate::http::client::BybitHttpClient;
23
24/// The heartbeat the exchange asks for.
25pub const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(20);
26/// Without any frame (a heartbeat answer counts) for this long the link is taken for dead.
27pub const SILENCE_LIMIT: Duration = Duration::from_secs(50);
28/// The first pause before a reconnect.
29pub const DEFAULT_RECONNECT_MIN: Duration = Duration::from_millis(500);
30/// The longest pause before a reconnect.
31pub const DEFAULT_RECONNECT_MAX: Duration = Duration::from_secs(30);
32const HEALTHY_AFTER: Duration = Duration::from_secs(10);
33const LOGIN_TIMEOUT: Duration = Duration::from_secs(10);
34/// How long the login stays valid after it is sent, in milliseconds.
35const LOGIN_EXPIRES_MS: i64 = 10_000;
36
37/// Which stream the session connects to.
38#[derive(Debug, Clone, Copy, PartialEq, Eq)]
39pub enum WsStream {
40    /// Public market data of the category of the config. No login.
41    Public,
42    /// The private streams of the account (`order`, `execution`, `position`, `wallet`). Logs in with the key.
43    Private,
44}
45
46/// How the session behaves.
47#[derive(Debug, Clone, Copy)]
48pub struct WsOptions {
49    /// Tell the reader about every connection with [`WsIncoming::Connected`].
50    pub announce_connects: bool,
51    /// The first pause before a reconnect.
52    pub reconnect_min: Duration,
53    /// The longest pause before a reconnect.
54    pub reconnect_max: Duration,
55}
56
57impl Default for WsOptions {
58    fn default() -> Self {
59        Self {
60            announce_connects: false,
61            reconnect_min: DEFAULT_RECONNECT_MIN,
62            reconnect_max: DEFAULT_RECONNECT_MAX,
63        }
64    }
65}
66
67struct Queued {
68    seq: Option<u64>,
69    request: WsRequest,
70}
71
72/// One WebSocket connection to Bybit, kept alive in a task of its own.
73pub struct WsSession {
74    user_tx: AsyncSender<Queued>,
75    /// Subscriptions to restore after a reconnect: (number, topic, req_id).
76    subscriptions: Arc<Mutex<Vec<(u64, String, String)>>>,
77    next_seq: AtomicU64,
78    is_connected: Arc<AtomicBool>,
79    dropped_events: Arc<AtomicU64>,
80    handle: tokio::task::JoinHandle<()>,
81}
82
83impl Drop for WsSession {
84    fn drop(&mut self) {
85        self.handle.abort();
86    }
87}
88
89impl WsSession {
90    /// Starts the session. It connects at once, in the background; `event_tx` gets every frame.
91    pub fn spawn(
92        http: BybitHttpClient,
93        stream: WsStream,
94        event_tx: AsyncSender<WsIncoming>,
95        options: WsOptions,
96    ) -> Self {
97        let (user_tx, user_rx) = kanal::bounded_async::<Queued>(128);
98        let subscriptions = Arc::new(Mutex::new(Vec::new()));
99        let is_connected = Arc::new(AtomicBool::new(false));
100        let dropped_events = Arc::new(AtomicU64::new(0));
101        let handle = tokio::spawn(run_loop(
102            http,
103            stream,
104            event_tx,
105            options,
106            Arc::clone(&subscriptions),
107            user_rx,
108            Arc::clone(&is_connected),
109            Arc::clone(&dropped_events),
110        ));
111        Self {
112            user_tx,
113            subscriptions,
114            next_seq: AtomicU64::new(0),
115            is_connected,
116            dropped_events,
117            handle,
118        }
119    }
120
121    /// How many frames were dropped because the reader's channel was full.
122    pub fn dropped_events(&self) -> u64 {
123        self.dropped_events.load(Ordering::Relaxed)
124    }
125
126    /// Whether the link is up (and logged in, for the private stream).
127    pub fn is_connected(&self) -> bool {
128        self.is_connected.load(Ordering::Acquire)
129    }
130
131    /// Subscribes to `topic`. The request is sent when the link is up and again after every reconnect. The answer of
132    /// the exchange comes as a [`WsIncoming::Ack`] carrying `req_id`.
133    ///
134    /// # Errors
135    ///
136    /// [`BybitApiError::Subscription`] when the session is gone.
137    pub async fn subscribe(&self, req_id: &str, topic: &str) -> Result<(), BybitApiError> {
138        let seq = self.next_seq.fetch_add(1, Ordering::Relaxed);
139        self.subscriptions
140            .lock()
141            .await
142            .push((seq, topic.to_string(), req_id.to_string()));
143        self.user_tx
144            .send(Queued {
145                seq: Some(seq),
146                request: WsRequest::subscribe(req_id, topic),
147            })
148            .await
149            .map_err(|e| BybitApiError::Subscription {
150                message: format!("failed to send the subscription: {e}"),
151            })
152    }
153
154    /// Asks the exchange for `topic` again: it is cancelled and subscribed anew under `req_id`, so an order book
155    /// topic sends a fresh snapshot. A reconnect restores it under the new `req_id`.
156    ///
157    /// # Errors
158    ///
159    /// [`BybitApiError::Subscription`] when the session is gone.
160    pub async fn resubscribe(&self, req_id: &str, topic: &str) -> Result<(), BybitApiError> {
161        for (_, t, r) in self.subscriptions.lock().await.iter_mut() {
162            if t == topic {
163                *r = req_id.to_string();
164            }
165        }
166        for request in [
167            WsRequest::unsubscribe("", topic),
168            WsRequest::subscribe(req_id, topic),
169        ] {
170            self.user_tx
171                .send(Queued { seq: None, request })
172                .await
173                .map_err(|e| BybitApiError::Subscription {
174                    message: format!("failed to send the request: {e}"),
175                })?;
176        }
177        Ok(())
178    }
179
180    /// Cancels the subscription to `topic` and forgets it, so a reconnect does not restore it.
181    ///
182    /// # Errors
183    ///
184    /// [`BybitApiError::Subscription`] when the session is gone.
185    pub async fn unsubscribe(&self, topic: &str) -> Result<(), BybitApiError> {
186        self.subscriptions
187            .lock()
188            .await
189            .retain(|(_, t, _)| t != topic);
190        self.user_tx
191            .send(Queued {
192                seq: None,
193                request: WsRequest::unsubscribe("", topic),
194            })
195            .await
196            .map_err(|e| BybitApiError::Subscription {
197                message: format!("failed to send the unsubscription: {e}"),
198            })
199    }
200}
201
202#[allow(clippy::too_many_arguments)]
203async fn run_loop(
204    http: BybitHttpClient,
205    stream: WsStream,
206    event_tx: AsyncSender<WsIncoming>,
207    options: WsOptions,
208    subscriptions: Arc<Mutex<Vec<(u64, String, String)>>>,
209    user_rx: AsyncReceiver<Queued>,
210    is_connected: Arc<AtomicBool>,
211    dropped_events: Arc<AtomicU64>,
212) {
213    let url = match stream {
214        WsStream::Public => http.config().ws_public_url(),
215        WsStream::Private => http.config().ws_private_url(),
216    };
217    let backoff_min = options.reconnect_min.max(Duration::from_millis(10));
218    let backoff_max = options.reconnect_max.max(backoff_min);
219    let mut backoff = backoff_min;
220
221    'session: loop {
222        let ws_stream = match connect_async(url.as_str()).await {
223            Ok((s, res)) => {
224                info!(
225                    "WebSocket connected to {url} with status <{}>",
226                    res.status()
227                );
228                s
229            }
230            Err(e) => {
231                error!("WS connect failed: {e}");
232                sleep(backoff).await;
233                backoff = (backoff * 2).min(backoff_max);
234                continue;
235            }
236        };
237        let connected_at = Instant::now();
238        let (mut write, mut read) = ws_stream.split();
239
240        if stream == WsStream::Private {
241            let cfg = http.config();
242            let login = WsRequest::auth(
243                &cfg.api_key,
244                &cfg.api_secret,
245                http.timestamp_ms() + LOGIN_EXPIRES_MS,
246            );
247            if write
248                .send(Message::Text(login.to_json().into()))
249                .await
250                .is_err()
251            {
252                error!("WS: failed to send the login");
253                sleep(backoff).await;
254                backoff = (backoff * 2).min(backoff_max);
255                continue 'session;
256            }
257            match wait_login(&mut read).await {
258                Ok(()) => {}
259                Err(reason) => {
260                    error!("WS: login failed: {reason}");
261                    // The reader is told, so a subscription waiting for the stream fails with the reason.
262                    let _ = event_tx
263                        .send(WsIncoming::Ack {
264                            op: "auth".into(),
265                            req_id: None,
266                            success: false,
267                            message: reason,
268                        })
269                        .await;
270                    sleep(backoff).await;
271                    backoff = (backoff * 2).min(backoff_max);
272                    continue 'session;
273                }
274            }
275        }
276
277        // The reader is told first: the data of the restored subscriptions comes after this marker.
278        if options.announce_connects && event_tx.send(WsIncoming::Connected).await.is_err() {
279            break 'session;
280        }
281
282        // Restore. The numbers sent here are remembered: the same request still in the queue is not sent twice.
283        let mut restored: std::collections::HashSet<u64> = std::collections::HashSet::new();
284        {
285            let subs = subscriptions.lock().await;
286            for (seq, topic, req_id) in subs.iter() {
287                let text = WsRequest::subscribe(req_id, topic).to_json();
288                if write.send(Message::Text(text.into())).await.is_err() {
289                    error!("WS: failed to restore a subscription");
290                    continue 'session;
291                }
292                restored.insert(*seq);
293            }
294            if !subs.is_empty() {
295                info!("Restored {} subscriptions", subs.len());
296            }
297        }
298        is_connected.store(true, Ordering::Release);
299
300        let mut heartbeat = time::interval(HEARTBEAT_INTERVAL);
301        heartbeat.tick().await; // the first tick is at once
302        let mut last_rx = Instant::now();
303
304        'inner: loop {
305            tokio::select! {
306                frame = read.next() => match frame {
307                    Some(Ok(Message::Text(text))) => {
308                        last_rx = Instant::now();
309                        let incoming = parse_incoming(&text);
310                        if incoming == WsIncoming::Pong {
311                            continue 'inner;
312                        }
313                        match event_tx.try_send(incoming) {
314                            Ok(true) => {}
315                            Ok(false) => {
316                                dropped_events.fetch_add(1, Ordering::Relaxed);
317                                warn!("WS: the receiving channel is full, a frame is dropped");
318                            }
319                            Err(_) => break 'session,
320                        }
321                    }
322                    Some(Ok(Message::Ping(data))) => {
323                        last_rx = Instant::now();
324                        if write.send(Message::Pong(data)).await.is_err() {
325                            break 'inner;
326                        }
327                    }
328                    Some(Ok(Message::Close(_))) => {
329                        info!("WS closed by the server");
330                        break 'inner;
331                    }
332                    Some(Ok(_)) => {}
333                    Some(Err(e)) => {
334                        error!("WS read error: {e}");
335                        break 'inner;
336                    }
337                    None => break 'inner,
338                },
339                _ = heartbeat.tick() => {
340                    if last_rx.elapsed() > SILENCE_LIMIT {
341                        warn!("WS: nothing from the exchange for {SILENCE_LIMIT:?}, reconnecting");
342                        break 'inner;
343                    }
344                    if write.send(Message::Text(WsRequest::ping().to_json().into())).await.is_err() {
345                        break 'inner;
346                    }
347                    debug!("WS heartbeat sent");
348                }
349                msg = user_rx.recv() => match msg {
350                    Ok(Queued { seq, request }) => {
351                        if seq.is_some_and(|n| restored.remove(&n)) {
352                            continue 'inner;
353                        }
354                        if write.send(Message::Text(request.to_json().into())).await.is_err() {
355                            error!("WS write error");
356                            break 'inner;
357                        }
358                    }
359                    Err(_) => break 'session,
360                },
361            }
362        }
363
364        is_connected.store(false, Ordering::Release);
365        if connected_at.elapsed() >= HEALTHY_AFTER {
366            backoff = backoff_min;
367        }
368        warn!("WS disconnected, reconnecting in {backoff:?}...");
369        sleep(backoff).await;
370        backoff = (backoff * 2).min(backoff_max);
371    }
372    is_connected.store(false, Ordering::Release);
373}
374
375/// Reads frames until the answer to the login. Returns the reason when the exchange refused it or the link failed.
376async fn wait_login<S>(read: &mut S) -> Result<(), String>
377where
378    S: futures_util::Stream<Item = Result<Message, tokio_tungstenite::tungstenite::Error>> + Unpin,
379{
380    let deadline = Instant::now() + LOGIN_TIMEOUT;
381    loop {
382        let next = time::timeout_at(deadline, read.next())
383            .await
384            .map_err(|_| "no answer to the login".to_string())?;
385        match next {
386            Some(Ok(Message::Text(text))) => {
387                if let WsIncoming::Ack {
388                    op,
389                    success,
390                    message,
391                    ..
392                } = parse_incoming(&text)
393                    && op == "auth"
394                {
395                    return if success {
396                        Ok(())
397                    } else {
398                        Err(if message.is_empty() {
399                            "the exchange refused the login".to_string()
400                        } else {
401                            message
402                        })
403                    };
404                }
405            }
406            Some(Ok(Message::Close(_))) | None => {
407                return Err("the link closed during the login".into());
408            }
409            Some(Err(e)) => return Err(format!("the link failed during the login: {e}")),
410            Some(Ok(_)) => {}
411        }
412    }
413}
414
415#[cfg(test)]
416mod tests {
417    use super::*;
418
419    #[test]
420    fn the_pauses_are_the_range_of_the_other_connectors() {
421        let o = WsOptions::default();
422        assert_eq!(
423            (o.reconnect_min, o.reconnect_max),
424            (Duration::from_millis(500), Duration::from_secs(30))
425        );
426        assert!(!o.announce_connects);
427    }
428
429    #[test]
430    fn the_silence_limit_allows_two_missed_heartbeats() {
431        assert!(SILENCE_LIMIT > HEARTBEAT_INTERVAL * 2);
432    }
433}
brokapi_bybit_protocol/ws/session.rs

brokapi_bybit_protocol/ws/
session.rs
