nautilus_binance/spot/websocket/streams/
handler.rs1use std::{
23 collections::VecDeque,
24 sync::{
25 Arc,
26 atomic::{AtomicBool, AtomicU64, Ordering},
27 },
28};
29
30use nautilus_network::{
31 RECONNECTED,
32 websocket::{SubscriptionState, WebSocketClient},
33};
34use tokio_tungstenite::tungstenite::Message;
35use ustr::Ustr;
36
37pub use super::parse::{MarketDataMessage, decode_market_data};
38use super::{
39 messages::{
40 BinanceSpotServerShutdownMsg, BinanceSpotWsMessage, BinanceSpotWsStreamsCommand,
41 BinanceWsErrorMsg, BinanceWsErrorResponse, BinanceWsResponse, BinanceWsSubscription,
42 },
43 parse::decode_market_data as decode_sbe,
44};
45use crate::common::{
46 consts::BINANCE_RATE_LIMIT_KEY_SUBSCRIPTION,
47 websocket::{
48 PendingSubscriptionRequest, PendingSubscriptionRequests, reset_requests_after_reconnect,
49 },
50};
51
52pub(super) struct BinanceSpotWsFeedHandler {
58 #[allow(dead_code)]
59 signal: Arc<AtomicBool>,
60 inner: Option<WebSocketClient>,
61 cmd_rx: tokio::sync::mpsc::UnboundedReceiver<BinanceSpotWsStreamsCommand>,
62 raw_rx: tokio::sync::mpsc::UnboundedReceiver<Message>,
63 #[allow(dead_code)]
64 out_tx: tokio::sync::mpsc::UnboundedSender<BinanceSpotWsMessage>,
65 subscriptions: SubscriptionState,
66 request_id_counter: Arc<AtomicU64>,
67 pending_messages: VecDeque<BinanceSpotWsMessage>,
68 pending_requests: PendingSubscriptionRequests,
69}
70
71impl BinanceSpotWsFeedHandler {
72 pub(super) fn new(
74 signal: Arc<AtomicBool>,
75 cmd_rx: tokio::sync::mpsc::UnboundedReceiver<BinanceSpotWsStreamsCommand>,
76 raw_rx: tokio::sync::mpsc::UnboundedReceiver<Message>,
77 out_tx: tokio::sync::mpsc::UnboundedSender<BinanceSpotWsMessage>,
78 subscriptions: SubscriptionState,
79 request_id_counter: Arc<AtomicU64>,
80 ) -> Self {
81 Self {
82 signal,
83 inner: None,
84 cmd_rx,
85 raw_rx,
86 out_tx,
87 subscriptions,
88 request_id_counter,
89 pending_messages: VecDeque::new(),
90 pending_requests: PendingSubscriptionRequests::default(),
91 }
92 }
93
94 pub(super) async fn next(&mut self) -> Option<BinanceSpotWsMessage> {
98 if let Some(message) = self.pending_messages.pop_front() {
99 return Some(message);
100 }
101
102 loop {
103 tokio::select! {
104 Some(cmd) = self.cmd_rx.recv() => {
105 match cmd {
106 BinanceSpotWsStreamsCommand::SetClient(client) => {
107 log::debug!("Handler received WebSocket client");
108 self.inner = Some(client);
109 }
110 BinanceSpotWsStreamsCommand::Disconnect => {
111 log::debug!("Handler disconnecting WebSocket client");
112 self.inner = None;
113 return None;
114 }
115 BinanceSpotWsStreamsCommand::Subscribe { streams } => {
116 if let Err(e) = self.handle_subscribe(streams).await {
117 log::error!("Failed to handle subscribe command: {e}");
118 }
119 }
120 BinanceSpotWsStreamsCommand::Unsubscribe { streams } => {
121 if let Err(e) = self.handle_unsubscribe(streams).await {
122 log::error!("Failed to handle unsubscribe command: {e}");
123 }
124 }
125 }
126 }
127 Some(msg) = self.raw_rx.recv() => {
128 if let Message::Text(ref text) = msg
129 && text.as_str() == RECONNECTED
130 {
131 reset_requests_after_reconnect(
132 &mut self.pending_requests,
133 &self.subscriptions,
134 );
135 log::debug!("Handler received reconnection signal");
136 return Some(BinanceSpotWsMessage::Reconnected);
137 }
138
139 let messages = self.handle_message(msg);
140 if !messages.is_empty() {
141 let mut iter = messages.into_iter();
142 let first = iter.next();
143 self.pending_messages.extend(iter);
144
145 if let Some(msg) = first {
146 return Some(msg);
147 }
148 }
149 }
150 else => {
151 return None;
152 }
153 }
154 }
155 }
156
157 fn handle_message(&mut self, msg: Message) -> Vec<BinanceSpotWsMessage> {
158 match msg {
159 Message::Binary(data) => self.handle_binary_frame(&data),
160 Message::Text(text) => self.handle_text_frame(&text),
161 Message::Close(_) => {
162 log::debug!("Received close frame");
163 vec![]
164 }
165 Message::Ping(_) | Message::Pong(_) | Message::Frame(_) => vec![],
166 }
167 }
168
169 fn handle_binary_frame(&self, data: &[u8]) -> Vec<BinanceSpotWsMessage> {
170 match decode_sbe(data) {
171 Ok(MarketDataMessage::Trades(event)) => {
172 vec![BinanceSpotWsMessage::Trades(event)]
173 }
174 Ok(MarketDataMessage::BestBidAsk(event)) => {
175 vec![BinanceSpotWsMessage::BestBidAsk(event)]
176 }
177 Ok(MarketDataMessage::DepthSnapshot(event)) => {
178 vec![BinanceSpotWsMessage::DepthSnapshot(event)]
179 }
180 Ok(MarketDataMessage::DepthDiff(event)) => {
181 vec![BinanceSpotWsMessage::DepthDiff(event)]
182 }
183 Err(e) => {
184 log::error!("SBE decode error: {e}");
185 vec![BinanceSpotWsMessage::RawBinary(data.to_vec())]
186 }
187 }
188 }
189
190 fn handle_text_frame(&mut self, text: &str) -> Vec<BinanceSpotWsMessage> {
191 if let Ok(error) = serde_json::from_str::<BinanceWsErrorResponse>(text) {
192 if let Some(id) = error.id
193 && let Some(request) = self.pending_requests.take(id)
194 {
195 request.mark_failure(&self.subscriptions);
196 log::warn!(
197 "Subscription request failed: id={id}, request={request:?}, code={}, msg={}",
198 error.code,
199 error.msg
200 );
201 }
202 return vec![BinanceSpotWsMessage::Error(BinanceWsErrorMsg {
203 code: error.code,
204 msg: error.msg,
205 })];
206 }
207
208 if let Ok(response) = serde_json::from_str::<BinanceWsResponse>(text) {
209 self.handle_subscription_response(&response);
210 return vec![];
211 }
212
213 classify_unsolicited_json(text)
214 }
215
216 fn handle_subscription_response(&mut self, response: &BinanceWsResponse) {
217 if let Some(request) = self.pending_requests.take(response.id) {
218 if response.result.is_none() {
219 request.confirm(&self.subscriptions);
220 log::debug!("Subscription request confirmed: request={request:?}");
221 } else {
222 request.mark_failure(&self.subscriptions);
223 log::warn!(
224 "Subscription request failed: request={request:?}, result={:?}",
225 response.result
226 );
227 }
228 } else {
229 log::debug!("Received response for unknown request: id={}", response.id);
230 }
231 }
232
233 async fn handle_subscribe(&mut self, streams: Vec<String>) -> anyhow::Result<()> {
234 let request_id = self.request_id_counter.fetch_add(1, Ordering::SeqCst);
235 let request = BinanceWsSubscription::subscribe(streams.clone(), request_id);
236 let payload = serde_json::to_string(&request)?;
237
238 for stream in &streams {
239 self.subscriptions.mark_subscribe(stream);
240 }
241
242 self.pending_requests
243 .insert(request_id, PendingSubscriptionRequest::subscribe(streams));
244
245 if let Err(e) = self
246 .send_text(
247 payload,
248 Some(BINANCE_RATE_LIMIT_KEY_SUBSCRIPTION.as_slice()),
249 )
250 .await
251 {
252 if let Some(request) = self.pending_requests.take(request_id) {
253 request.mark_failure(&self.subscriptions);
254 }
255 return Err(e);
256 }
257 Ok(())
258 }
259
260 async fn handle_unsubscribe(&mut self, streams: Vec<String>) -> anyhow::Result<()> {
261 let request_id = self.request_id_counter.fetch_add(1, Ordering::SeqCst);
262 let request = BinanceWsSubscription::unsubscribe(streams.clone(), request_id);
263 let payload = serde_json::to_string(&request)?;
264
265 for stream in &streams {
266 self.subscriptions.mark_unsubscribe(stream);
267 }
268
269 self.pending_requests
270 .insert(request_id, PendingSubscriptionRequest::unsubscribe(streams));
271
272 if let Err(e) = self
273 .send_text(
274 payload,
275 Some(BINANCE_RATE_LIMIT_KEY_SUBSCRIPTION.as_slice()),
276 )
277 .await
278 {
279 self.pending_requests.take(request_id);
280 return Err(e);
281 }
282
283 Ok(())
284 }
285
286 async fn send_text(
287 &self,
288 payload: String,
289 rate_limit_keys: Option<&[Ustr]>,
290 ) -> anyhow::Result<()> {
291 let Some(client) = &self.inner else {
292 anyhow::bail!("No active WebSocket client");
293 };
294 client
295 .send_text(payload, rate_limit_keys)
296 .await
297 .map_err(|e| anyhow::anyhow!("Failed to send message: {e}"))?;
298 Ok(())
299 }
300}
301
302fn classify_unsolicited_json(text: &str) -> Vec<BinanceSpotWsMessage> {
306 let Ok(value) = serde_json::from_str::<serde_json::Value>(text) else {
307 log::warn!("Failed to parse JSON message: {text}");
308 return vec![];
309 };
310
311 if value.get("e").and_then(|v| v.as_str()) == Some("serverShutdown")
312 && let Ok(msg) = serde_json::from_value::<BinanceSpotServerShutdownMsg>(value.clone())
313 {
314 log::warn!(
315 "Binance server shutdown notice received (event_time={}); disconnect expected ~10 minutes from event",
316 msg.event_time,
317 );
318 return vec![BinanceSpotWsMessage::ServerShutdown(msg)];
319 }
320
321 vec![BinanceSpotWsMessage::RawJson(value)]
322}
323
324#[cfg(test)]
325mod tests {
326 use std::sync::{
327 Arc,
328 atomic::{AtomicBool, AtomicU64},
329 };
330
331 use nautilus_network::websocket::SubscriptionState;
332 use rstest::rstest;
333
334 use super::*;
335
336 #[rstest]
337 fn test_classify_unsolicited_json_server_shutdown_emits_variant() {
338 let text = r#"{"e":"serverShutdown","E":1700000000000}"#;
339 let out = classify_unsolicited_json(text);
340 assert_eq!(out.len(), 1);
341 match &out[0] {
342 BinanceSpotWsMessage::ServerShutdown(msg) => {
343 assert_eq!(msg.event_type, "serverShutdown");
344 assert_eq!(msg.event_time, 1_700_000_000_000);
345 }
346 other => panic!("expected ServerShutdown variant, was {other:?}"),
347 }
348 }
349
350 #[rstest]
351 fn test_classify_unsolicited_json_unrelated_emits_raw_json() {
352 let text = r#"{"e":"trade","p":"50000"}"#;
353 let out = classify_unsolicited_json(text);
354 assert_eq!(out.len(), 1);
355 assert!(matches!(out[0], BinanceSpotWsMessage::RawJson(_)));
356 }
357
358 #[rstest]
359 fn test_classify_unsolicited_json_invalid_returns_empty() {
360 let out = classify_unsolicited_json("not json");
361 assert!(out.is_empty());
362 }
363
364 #[rstest]
365 fn test_handle_text_frame_error_with_id_emits_error_and_clears_pending_request() {
366 let signal = Arc::new(AtomicBool::new(false));
367 let request_id_counter = Arc::new(AtomicU64::new(2));
368 let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
369 let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
370 let (out_tx, _out_rx) = tokio::sync::mpsc::unbounded_channel();
371 let subscriptions = SubscriptionState::new('@');
372
373 let mut handler = BinanceSpotWsFeedHandler::new(
374 signal,
375 cmd_rx,
376 raw_rx,
377 out_tx,
378 subscriptions,
379 request_id_counter,
380 );
381 handler.pending_requests.insert(
382 1,
383 PendingSubscriptionRequest::subscribe(vec!["btcusdt@trade".to_string()]),
384 );
385
386 let out = handler.handle_text_frame(r#"{"code":2,"msg":"Invalid request","id":1}"#);
387 assert_eq!(out.len(), 1);
388 match &out[0] {
389 BinanceSpotWsMessage::Error(err) => {
390 assert_eq!(err.code, 2);
391 assert_eq!(err.msg, "Invalid request");
392 }
393 other => panic!("expected Error variant, was {other:?}"),
394 }
395 assert_eq!(handler.pending_requests.len(), 0);
396 }
397}