bothan_bybit/api/
websocket.rs

1//! Bybit WebSocket API client implementation.
2//!
3//! This module provides the [`WebSocketConnector`] and [`WebSocketConnection`] for interacting
4//! with the Bybit WebSocket API. It enables real-time streaming of market data, such as ticker
5//! updates, and is used internally to implement the [`AssetInfoProvider`] trait for asset workers.
6//!
7//! This module provides:
8//!
9//! - Establishes WebSocket connections to Bybit servers
10//! - Subscribes and unsubscribes to ticker streams for specified symbols
11//! - Processes incoming WebSocket messages, including ticker updates and ping events
12//! - Transforms WebSocket messages into [`AssetInfo`] for use in workers
13//! - Handles connection management, including closing connections gracefully
14
15use bothan_lib::types::AssetInfo;
16use bothan_lib::worker::websocket::{AssetInfoProvider, AssetInfoProviderConnector, Data};
17use futures_util::{SinkExt, StreamExt};
18use rust_decimal::Decimal;
19use serde_json::json;
20use tokio::net::TcpStream;
21use tokio_tungstenite::tungstenite::Message;
22use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async, tungstenite};
23
24use crate::api::error::{Error, ListeningError};
25use crate::api::types::{MAX_ARGS, PublicTickerResponse, Response};
26
27/// A connector for establishing WebSocket connections to the Bybit WebSocket API.
28pub struct WebSocketConnector {
29    url: String,
30}
31
32/// This struct provides methods to create a new connector and connect to the WebSocket server.
33///
34/// # Examples
35///
36/// ```rust
37/// use bothan_bybit::WebSocketConnector;
38///
39/// let connector = WebSocketConnector::new("wss://example.com/socket");
40/// let connection = connector.connect();
41/// ```
42impl WebSocketConnector {
43    /// Creates a new `BybitWebSocketConnector` with the given URL.
44    pub fn new(url: impl Into<String>) -> Self {
45        Self { url: url.into() }
46    }
47
48    /// Establishes a WebSocket connection to the Bybit server.
49    ///
50    /// # Errors
51    ///
52    /// Returns a [`tungstenite::Error`] if the WebSocket connection fails.
53    pub async fn connect(&self) -> Result<WebSocketConnection, tungstenite::Error> {
54        let (wss, _) = connect_async(self.url.clone()).await?;
55
56        Ok(WebSocketConnection::new(wss))
57    }
58}
59
60#[async_trait::async_trait]
61impl AssetInfoProviderConnector for WebSocketConnector {
62    type Provider = WebSocketConnection;
63    type Error = tungstenite::Error;
64
65    async fn connect(&self) -> Result<WebSocketConnection, Self::Error> {
66        WebSocketConnector::connect(self).await
67    }
68}
69
70/// Represents an active WebSocket connection to Bybit.
71///
72/// This struct encapsulates the WebSocket stream and provides methods for subscribing to
73/// ticker streams, receiving messages, and closing the connection.
74pub struct WebSocketConnection {
75    ws_stream: WebSocketStream<MaybeTlsStream<TcpStream>>,
76}
77
78impl WebSocketConnection {
79    /// Creates a new `BybitWebSocketConnection`.
80    pub fn new(ws_stream: WebSocketStream<MaybeTlsStream<TcpStream>>) -> Self {
81        Self { ws_stream }
82    }
83
84    /// Subscribes to the ticker stream for the specified symbol IDs.
85    ///
86    /// This method sends a subscription request to the Bybit WebSocket API for the specified symbol IDs.
87    /// Each symbol ID is transformed into a ticker stream identifier before being sent.
88    ///
89    /// # Parameters
90    ///
91    /// - `symbols`: A slice of symbol IDs to subscribe to.
92    ///
93    /// # Errors
94    ///
95    /// Returns a [`tungstenite::Error`] if the WebSocket subscription request fails.
96    pub async fn subscribe_ticker<T: AsRef<str>>(
97        &mut self,
98        symbols: &[T],
99    ) -> Result<(), tungstenite::Error> {
100        let payload = json!({
101            "op": "subscribe",
102            "args": symbols.iter().map(|s| format!("tickers.{}", s.as_ref())).collect::<Vec<String>>(),
103        });
104
105        // Send the subscription message.
106        let message = Message::Text(payload.to_string());
107        self.ws_stream.send(message).await
108    }
109
110    /// Unsubscribes from the ticker stream for the specified symbol IDs.
111    ///
112    /// This method sends an unsubscription request to the Bybit WebSocket API for the specified symbol IDs.
113    /// Each symbol ID is transformed into a ticker stream identifier before being sent.
114    ///
115    /// # Parameters
116    ///
117    /// - `symbols`: A slice of symbol IDs to unsubscribe from.
118    ///
119    /// # Errors
120    ///
121    /// Returns a [`tungstenite::Error`] if the WebSocket unsubscription request fails.
122    pub async fn unsubscribe_ticker<T: AsRef<str>>(
123        &mut self,
124        symbols: &[T],
125    ) -> Result<(), tungstenite::Error> {
126        let payload = json!({
127            "op": "unsubscribe",
128            "args": symbols.iter().map(|s| format!("tickers.{}", s.as_ref())).collect::<Vec<String>>(),
129        });
130
131        // Send the unsubscription message.
132        let message = Message::Text(payload.to_string());
133        self.ws_stream.send(message).await
134    }
135
136    /// Retrieves the next message from the WebSocket stream.
137    ///
138    /// This method listens for incoming WebSocket messages and processes them.
139    /// Supported message types include text messages (parsed as `Response`), ping messages, and close messages.
140    pub async fn next(&mut self) -> Option<Result<Response, Error>> {
141        match self.ws_stream.next().await {
142            Some(Ok(Message::Text(msg))) => Some(parse_msg(msg)),
143            Some(Ok(Message::Ping(_))) => Some(Ok(Response::Ping)),
144            Some(Ok(Message::Close(_))) => None,
145            Some(Ok(_)) => Some(Err(Error::UnsupportedWebsocketMessageType)),
146            Some(Err(_)) => None, // Consider the connection closed if error detected
147            None => None,
148        }
149    }
150
151    /// Closes the WebSocket connection gracefully.
152    ///
153    /// This method sends a close frame to the WebSocket server and waits for the connection to close.
154    pub async fn close(&mut self) -> Result<(), tungstenite::Error> {
155        self.ws_stream.close(None).await?;
156        Ok(())
157    }
158}
159
160fn parse_msg(msg: String) -> Result<Response, Error> {
161    Ok(serde_json::from_str::<Response>(&msg)?)
162}
163
164#[async_trait::async_trait]
165impl AssetInfoProvider for WebSocketConnection {
166    type SubscriptionError = tungstenite::Error;
167    type ListeningError = ListeningError;
168
169    async fn subscribe(&mut self, ids: &[String]) -> Result<(), Self::SubscriptionError> {
170        for chunk in ids.chunks(MAX_ARGS) {
171            self.subscribe_ticker(chunk).await?;
172        }
173        Ok(())
174    }
175
176    async fn next(&mut self) -> Option<Result<Data, Self::ListeningError>> {
177        WebSocketConnection::next(self).await.map(|r| {
178            Ok(match r? {
179                Response::PublicTicker(t) => parse_public_ticker(t)?,
180                Response::Ping => Data::Ping,
181                _ => Data::Unused,
182            })
183        })
184    }
185
186    async fn try_close(mut self) {
187        tokio::spawn(async move { self.close().await });
188    }
189}
190
191fn parse_public_ticker(ticker: PublicTickerResponse) -> Result<Data, rust_decimal::Error> {
192    let asset_info = AssetInfo::new(
193        ticker.data.symbol,
194        Decimal::from_str_exact(&ticker.data.last_price)?,
195        ticker.ts / 1000, // convert from millisecond to second
196    );
197    Ok(Data::AssetInfo(vec![asset_info]))
198}
199
200#[cfg(test)]
201pub(crate) mod test {
202    use tokio::sync::mpsc;
203    use ws_mock::ws_mock_server::{WsMock, WsMockServer};
204
205    use super::*;
206    use crate::api::types::{PublicMessageResponse, PublicTickerResponse, Response, Ticker};
207
208    pub(crate) async fn setup_mock_server() -> WsMockServer {
209        WsMockServer::start().await
210    }
211
212    #[tokio::test]
213    async fn test_recv_public_ticker() {
214        // Set up the mock server and the WebSocket connector.
215        let server = setup_mock_server().await;
216        let connector = WebSocketConnector::new(server.uri().await);
217        let (mpsc_send, mpsc_recv) = mpsc::channel::<Message>(32);
218
219        // Create a mock ticker response.
220        let mock_ticker = Ticker {
221            symbol: "BTCUSDT".to_string(),
222            last_price: "42000.00".to_string(),
223            high_price24h: "44000.00".to_string(),
224            low_price24h: "40000.00".to_string(),
225            prev_price24h: "40000.00".to_string(),
226            volume24h: "100000.00".to_string(),
227            turnover24h: "4200000000.00".to_string(),
228            price24h_pcnt: "0.05".to_string(),
229            usd_index_price: "42000.00".to_string(),
230        };
231
232        // Create the mock PublicTickerResponse.
233        let mock_resp = Response::PublicTicker(PublicTickerResponse {
234            topic: "tickers.BTCUSDT".to_string(),
235            ts: 1673853746003,
236            ticker_type: "snapshot".to_string(),
237            cs: 2588407389,
238            data: mock_ticker,
239        });
240
241        // Mount the mock WebSocket server and send the mock response.
242        WsMock::new()
243            .forward_from_channel(mpsc_recv)
244            .mount(&server)
245            .await;
246        mpsc_send
247            .send(Message::Text(serde_json::to_string(&mock_resp).unwrap()))
248            .await
249            .unwrap();
250
251        // Connect to the mock WebSocket server and retrieve the response.
252        let mut connection = connector.connect().await.unwrap();
253        let resp = connection.next().await.unwrap().unwrap();
254
255        // Assert that the received response matches the mock response.
256        assert_eq!(resp, mock_resp);
257    }
258
259    /// Test for receiving a public message response from the WebSocket.
260    #[tokio::test]
261    async fn test_recv_public_message() {
262        // Set up the mock server and the WebSocket connector.
263        let server = setup_mock_server().await;
264        let connector = WebSocketConnector::new(server.uri().await);
265        let (mpsc_send, mpsc_recv) = mpsc::channel::<Message>(32);
266
267        // Create a mock public message response.
268        let mock_message = PublicMessageResponse {
269            success: true,
270            ret_msg: "subscribe".to_string(),
271            conn_id: "2324d924-aa4d-45b0-a858-7b8be29ab52b".to_string(),
272            req_id: Some("10001".to_string()),
273            op: "subscribe".to_string(),
274        };
275
276        // Create the mock BybitResponse with the PublicMessageResponse.
277        let mock_resp = Response::PublicMessage(mock_message);
278
279        // Mount the mock WebSocket server and send the mock response.
280        WsMock::new()
281            .forward_from_channel(mpsc_recv)
282            .mount(&server)
283            .await;
284        mpsc_send
285            .send(Message::Text(serde_json::to_string(&mock_resp).unwrap()))
286            .await
287            .unwrap();
288
289        // Connect to the mock WebSocket server and retrieve the response.
290        let mut connection = connector.connect().await.unwrap();
291        let resp = connection.next().await.unwrap().unwrap();
292
293        // Assert that the received response matches the mock response.
294        assert_eq!(resp, mock_resp);
295    }
296
297    #[tokio::test]
298    async fn test_recv_close() {
299        // Set up the mock server and the WebSocket connector.
300        let server = setup_mock_server().await;
301        let connector = WebSocketConnector::new(server.uri().await);
302        let (mpsc_send, mpsc_recv) = mpsc::channel::<Message>(32);
303
304        // Mount the mock WebSocket server and send a close message.
305        WsMock::new()
306            .forward_from_channel(mpsc_recv)
307            .mount(&server)
308            .await;
309        mpsc_send.send(Message::Close(None)).await.unwrap();
310
311        // Connect to the mock WebSocket server and verify the connection closure.
312        let mut connection = connector.connect().await.unwrap();
313        let resp = connection.next().await;
314        assert!(resp.is_none());
315    }
316}
bothan_bybit/api/websocket.rs

bothan_bybit/api/
websocket.rs
