1use std::{
19 collections::VecDeque,
20 fmt::Debug,
21 sync::{
22 Arc,
23 atomic::{AtomicBool, AtomicU64, Ordering},
24 },
25};
26
27use nautilus_network::{
28 RECONNECTED,
29 websocket::{SubscriptionState, WebSocketClient},
30};
31use ustr::Ustr;
32
33use super::messages::{
34 BinanceCombinedStreamEvent, BinanceSpotBookTickerMsg, BinanceSpotDepthDiffMsg,
35 BinanceSpotKlineMsg, BinanceSpotPartialDepthMsg, BinanceSpotPartialDepthPayload,
36 BinanceSpotPublicWsCommand, BinanceSpotPublicWsMessage, BinanceSpotServerShutdownMsg,
37 BinanceSpotTickerMsg, BinanceSpotTradeMsg, BinanceSpotWsErrorResponse, BinanceSpotWsResponse,
38 BinanceWsSubscription,
39};
40use crate::common::{
41 consts::BINANCE_RATE_LIMIT_KEY_SUBSCRIPTION,
42 enums::BinanceWsEventType,
43 websocket::{
44 PendingSubscriptionRequest, PendingSubscriptionRequests, reset_requests_after_reconnect,
45 },
46};
47
48pub(super) struct BinanceSpotPublicWsHandler {
50 signal: Arc<AtomicBool>,
51 cmd_rx: tokio::sync::mpsc::UnboundedReceiver<BinanceSpotPublicWsCommand>,
52 raw_rx: tokio::sync::mpsc::UnboundedReceiver<Vec<u8>>,
53 inner: Option<WebSocketClient>,
54 pending_messages: VecDeque<BinanceSpotPublicWsMessage>,
55 subscriptions: SubscriptionState,
56 request_id_counter: Arc<AtomicU64>,
57 pending_requests: PendingSubscriptionRequests,
58}
59
60impl Debug for BinanceSpotPublicWsHandler {
61 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
62 f.debug_struct(stringify!(BinanceSpotPublicWsHandler))
63 .field("pending_requests", &self.pending_requests.len())
64 .finish_non_exhaustive()
65 }
66}
67
68impl BinanceSpotPublicWsHandler {
69 pub(super) fn new(
70 signal: Arc<AtomicBool>,
71 cmd_rx: tokio::sync::mpsc::UnboundedReceiver<BinanceSpotPublicWsCommand>,
72 raw_rx: tokio::sync::mpsc::UnboundedReceiver<Vec<u8>>,
73 subscriptions: SubscriptionState,
74 request_id_counter: Arc<AtomicU64>,
75 ) -> Self {
76 Self {
77 signal,
78 cmd_rx,
79 raw_rx,
80 inner: None,
81 pending_messages: VecDeque::new(),
82 subscriptions,
83 request_id_counter,
84 pending_requests: PendingSubscriptionRequests::default(),
85 }
86 }
87
88 pub(super) async fn next(&mut self) -> Option<BinanceSpotPublicWsMessage> {
89 loop {
90 if let Some(msg) = self.pending_messages.pop_front() {
91 return Some(msg);
92 }
93
94 if self.signal.load(Ordering::Relaxed) {
95 return None;
96 }
97
98 tokio::select! {
99 Some(cmd) = self.cmd_rx.recv() => {
100 self.handle_command(cmd).await;
101 }
102 Some(raw) = self.raw_rx.recv() => {
103 let out = self.handle_raw_message(raw).await;
104 if !out.is_empty() {
105 let mut iter = out.into_iter();
106 let first = iter.next();
107 self.pending_messages.extend(iter);
108
109 if let Some(msg) = first {
110 return Some(msg);
111 }
112 }
113 }
114 else => {
115 return None;
116 }
117 }
118 }
119 }
120
121 async fn handle_command(&mut self, cmd: BinanceSpotPublicWsCommand) {
122 match cmd {
123 BinanceSpotPublicWsCommand::SetClient(client) => {
124 self.inner = Some(client);
125 }
126 BinanceSpotPublicWsCommand::Disconnect => {
127 if let Some(client) = &self.inner {
128 let () = client.disconnect().await;
129 }
130 self.inner = None;
131 }
132 BinanceSpotPublicWsCommand::Subscribe { streams } => {
133 self.send_subscribe(streams).await;
134 }
135 BinanceSpotPublicWsCommand::Unsubscribe { streams } => {
136 self.send_unsubscribe(streams).await;
137 }
138 }
139 }
140
141 async fn send_subscribe(&mut self, streams: Vec<String>) {
142 for stream in &streams {
143 self.subscriptions.mark_subscribe(stream);
144 }
145
146 let Some(client) = &self.inner else {
147 log::warn!("Cannot subscribe: no client connected");
148 return;
149 };
150
151 let request_id = self.request_id_counter.fetch_add(1, Ordering::Relaxed);
152 let request = BinanceWsSubscription::subscribe(streams.clone(), request_id);
153 let json = match serde_json::to_string(&request) {
154 Ok(j) => j,
155 Err(e) => {
156 log::error!("Failed to serialize subscribe request: {e}");
157 return;
158 }
159 };
160
161 self.pending_requests
162 .insert(request_id, PendingSubscriptionRequest::subscribe(streams));
163
164 if let Err(e) = client
165 .send_text(json, Some(BINANCE_RATE_LIMIT_KEY_SUBSCRIPTION.as_slice()))
166 .await
167 {
168 if let Some(request) = self.pending_requests.take(request_id) {
169 request.mark_failure(&self.subscriptions);
170 }
171 log::error!("Failed to send subscribe request: {e}");
172 }
173 }
174
175 async fn send_unsubscribe(&mut self, streams: Vec<String>) {
176 for stream in &streams {
177 self.subscriptions.mark_unsubscribe(stream);
178 }
179
180 let Some(client) = &self.inner else {
181 log::warn!("Cannot unsubscribe: no client connected");
182 return;
183 };
184
185 let request_id = self.request_id_counter.fetch_add(1, Ordering::Relaxed);
186
187 let request = BinanceWsSubscription::unsubscribe(streams.clone(), request_id);
188 let json = match serde_json::to_string(&request) {
189 Ok(j) => j,
190 Err(e) => {
191 log::error!("Failed to serialize unsubscribe request: {e}");
192 return;
193 }
194 };
195
196 self.pending_requests
197 .insert(request_id, PendingSubscriptionRequest::unsubscribe(streams));
198
199 if let Err(e) = client
200 .send_text(json, Some(BINANCE_RATE_LIMIT_KEY_SUBSCRIPTION.as_slice()))
201 .await
202 {
203 self.pending_requests.take(request_id);
204 log::error!("Failed to send unsubscribe request: {e}");
205 }
206 }
207
208 async fn handle_raw_message(&mut self, raw: Vec<u8>) -> Vec<BinanceSpotPublicWsMessage> {
209 if let Ok(text) = std::str::from_utf8(&raw)
210 && text == RECONNECTED
211 {
212 reset_requests_after_reconnect(&mut self.pending_requests, &self.subscriptions);
213 log::debug!("WebSocket reconnected signal received");
214 return vec![BinanceSpotPublicWsMessage::Reconnected];
215 }
216
217 let json: serde_json::Value = match serde_json::from_slice(&raw) {
218 Ok(j) => j,
219 Err(e) => {
220 log::warn!("Failed to parse Spot public JSON message: {e}");
221 return vec![];
222 }
223 };
224
225 if let Some(code) = json.get("code")
226 && let Some(code) = code.as_i64()
227 {
228 self.handle_subscription_response(&json);
229 let msg = json
230 .get("msg")
231 .and_then(|m| m.as_str())
232 .unwrap_or("Unknown error")
233 .to_string();
234 return vec![BinanceSpotPublicWsMessage::Error(
235 crate::spot::websocket::streams::messages::BinanceWsErrorMsg {
236 code: code as i32,
237 msg,
238 },
239 )];
240 }
241
242 if json.get("result").is_some() || json.get("id").is_some() {
243 self.handle_subscription_response(&json);
244 return vec![];
245 }
246
247 self.handle_stream_data(&json)
248 }
249
250 fn handle_subscription_response(&mut self, json: &serde_json::Value) {
251 if json.get("result").is_some()
252 && let Ok(response) = serde_json::from_value::<BinanceSpotWsResponse>(json.clone())
253 {
254 if let Some(request) = self.pending_requests.take(response.id) {
255 if response.result.is_none() {
256 request.confirm(&self.subscriptions);
257 log::debug!("Subscription request confirmed: request={request:?}");
258 } else {
259 request.mark_failure(&self.subscriptions);
260 log::warn!(
261 "Subscription request failed: request={request:?}, result={:?}",
262 response.result
263 );
264 }
265 }
266 } else if let Ok(error) = serde_json::from_value::<BinanceSpotWsErrorResponse>(json.clone())
267 {
268 if let Some(id) = error.id
269 && let Some(request) = self.pending_requests.take(id)
270 {
271 request.mark_failure(&self.subscriptions);
272 }
273 log::warn!(
274 "WebSocket error response: code={}, msg={}",
275 error.code,
276 error.msg
277 );
278 }
279 }
280
281 fn handle_stream_data(&self, json: &serde_json::Value) -> Vec<BinanceSpotPublicWsMessage> {
282 let (stream_name, payload) = split_combined_payload(json);
283
284 if let Some(depth) = parse_partial_depth_with_symbol(&payload, stream_name.as_deref()) {
285 return vec![BinanceSpotPublicWsMessage::DepthSnapshot(depth)];
286 }
287
288 if let Some(stream_name) = stream_name.as_deref()
289 && stream_name.ends_with("@bookTicker")
290 {
291 return serde_json::from_value::<BinanceSpotBookTickerMsg>(payload)
292 .map(BinanceSpotPublicWsMessage::BookTicker)
293 .map_err(|e| log::warn!("Failed to parse Spot bookTicker: {e}"))
294 .ok()
295 .into_iter()
296 .collect();
297 }
298
299 if payload.get("e").and_then(|v| v.as_str()) == Some("serverShutdown") {
303 return serde_json::from_value::<BinanceSpotServerShutdownMsg>(payload)
304 .map(BinanceSpotPublicWsMessage::ServerShutdown)
305 .map_err(|e| log::warn!("Failed to parse Spot server shutdown event: {e}"))
306 .ok()
307 .into_iter()
308 .collect();
309 }
310
311 let Some(event_type) = extract_event_type(&payload) else {
312 return vec![BinanceSpotPublicWsMessage::RawJson(payload)];
313 };
314
315 match event_type {
316 BinanceWsEventType::Trade => serde_json::from_value::<BinanceSpotTradeMsg>(payload)
317 .map(BinanceSpotPublicWsMessage::Trade)
318 .map_err(|e| log::warn!("Failed to parse Spot trade: {e}"))
319 .ok()
320 .into_iter()
321 .collect(),
322 BinanceWsEventType::BookTicker => {
323 serde_json::from_value::<BinanceSpotBookTickerMsg>(payload)
324 .map(BinanceSpotPublicWsMessage::BookTicker)
325 .map_err(|e| log::warn!("Failed to parse Spot bookTicker: {e}"))
326 .ok()
327 .into_iter()
328 .collect()
329 }
330 BinanceWsEventType::DepthUpdate => {
331 serde_json::from_value::<BinanceSpotDepthDiffMsg>(payload)
332 .map(BinanceSpotPublicWsMessage::DepthDiff)
333 .map_err(|e| log::warn!("Failed to parse Spot depthUpdate: {e}"))
334 .ok()
335 .into_iter()
336 .collect()
337 }
338 BinanceWsEventType::Kline => serde_json::from_value::<BinanceSpotKlineMsg>(payload)
339 .map(BinanceSpotPublicWsMessage::Kline)
340 .map_err(|e| log::warn!("Failed to parse Spot kline: {e}"))
341 .ok()
342 .into_iter()
343 .collect(),
344 BinanceWsEventType::Ticker24Hr => {
345 serde_json::from_value::<BinanceSpotTickerMsg>(payload)
346 .map(BinanceSpotPublicWsMessage::Ticker)
347 .map_err(|e| log::warn!("Failed to parse Spot ticker: {e}"))
348 .ok()
349 .into_iter()
350 .collect()
351 }
352 _ => vec![BinanceSpotPublicWsMessage::RawJson(payload)],
353 }
354 }
355}
356
357fn split_combined_payload(json: &serde_json::Value) -> (Option<String>, serde_json::Value) {
358 if let Ok(wrapper) = serde_json::from_value::<BinanceCombinedStreamEvent>(json.clone()) {
359 (Some(wrapper.stream), wrapper.data)
360 } else {
361 (None, json.clone())
362 }
363}
364
365fn extract_event_type(json: &serde_json::Value) -> Option<BinanceWsEventType> {
366 json.get("e")
367 .and_then(|v| serde_json::from_value(v.clone()).ok())
368}
369
370fn parse_partial_depth_with_symbol(
371 payload: &serde_json::Value,
372 stream_name: Option<&str>,
373) -> Option<BinanceSpotPartialDepthMsg> {
374 let parsed = serde_json::from_value::<BinanceSpotPartialDepthPayload>(payload.clone()).ok()?;
375
376 let symbol = stream_name
377 .and_then(|stream| stream.split('@').next())
378 .map(|s| Ustr::from(s.to_uppercase().as_str()))?;
379
380 Some(BinanceSpotPartialDepthMsg {
381 symbol,
382 last_update_id: parsed.last_update_id,
383 bids: parsed.bids,
384 asks: parsed.asks,
385 })
386}
387
388#[cfg(test)]
389mod tests {
390 use std::sync::{
391 Arc,
392 atomic::{AtomicBool, AtomicU64},
393 };
394
395 use nautilus_network::{RECONNECTED, websocket::SubscriptionState};
396 use rstest::rstest;
397 use serde_json::json;
398 use ustr::Ustr;
399
400 use super::*;
401
402 #[rstest]
403 fn test_parse_partial_depth_with_symbol_uppercases_symbol_from_stream_name() {
404 let payload = json!({
405 "lastUpdateId": 12345,
406 "bids": [["42000.1", "0.5"]],
407 "asks": [["42000.2", "0.8"]]
408 });
409
410 let parsed = parse_partial_depth_with_symbol(&payload, Some("btcusdt@depth20"))
411 .expect("depth payload should parse");
412
413 assert_eq!(parsed.symbol, Ustr::from("BTCUSDT"));
414 assert_eq!(parsed.last_update_id, 12345);
415 assert_eq!(parsed.bids.len(), 1);
416 assert_eq!(parsed.asks.len(), 1);
417 }
418
419 #[tokio::test]
420 async fn test_handle_raw_message_emits_reconnected_signal() {
421 let signal = Arc::new(AtomicBool::new(false));
422 let request_id_counter = Arc::new(AtomicU64::new(1));
423 let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
424 let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
425 let subscriptions = SubscriptionState::new('@');
426
427 let mut handler = BinanceSpotPublicWsHandler::new(
428 signal,
429 cmd_rx,
430 raw_rx,
431 subscriptions,
432 request_id_counter,
433 );
434
435 let out = handler
436 .handle_raw_message(RECONNECTED.as_bytes().to_vec())
437 .await;
438 assert_eq!(out.len(), 1);
439 assert!(matches!(out[0], BinanceSpotPublicWsMessage::Reconnected));
440 }
441
442 #[rstest]
443 #[tokio::test]
444 async fn test_subscription_intent_is_preserved_without_active_client() {
445 let signal = Arc::new(AtomicBool::new(false));
446 let request_id_counter = Arc::new(AtomicU64::new(1));
447 let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
448 let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
449 let subscriptions = SubscriptionState::new('@');
450 let subscribe_topic = "btcusdt@trade";
451 let unsubscribe_topic = "ethusdt@trade";
452 subscriptions.mark_subscribe(unsubscribe_topic);
453 subscriptions.confirm_subscribe(unsubscribe_topic);
454
455 let mut handler = BinanceSpotPublicWsHandler::new(
456 signal,
457 cmd_rx,
458 raw_rx,
459 subscriptions.clone(),
460 request_id_counter,
461 );
462
463 handler
464 .send_subscribe(vec![subscribe_topic.to_string()])
465 .await;
466 handler
467 .send_unsubscribe(vec![unsubscribe_topic.to_string()])
468 .await;
469
470 assert_eq!(
471 subscriptions.pending_subscribe_topics(),
472 [subscribe_topic.to_string()]
473 );
474 assert_eq!(
475 subscriptions.pending_unsubscribe_topics(),
476 [unsubscribe_topic.to_string()]
477 );
478 assert_eq!(subscriptions.len(), 0);
479 assert_eq!(handler.pending_requests.len(), 0);
480 }
481
482 #[tokio::test]
483 async fn test_handle_raw_message_error_with_id_emits_error() {
484 let signal = Arc::new(AtomicBool::new(false));
485 let request_id_counter = Arc::new(AtomicU64::new(2));
486 let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
487 let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
488 let subscriptions = SubscriptionState::new('@');
489
490 let mut handler = BinanceSpotPublicWsHandler::new(
491 signal,
492 cmd_rx,
493 raw_rx,
494 subscriptions,
495 request_id_counter,
496 );
497 handler.pending_requests.insert(
498 1,
499 PendingSubscriptionRequest::subscribe(vec!["btcusdt@trade".to_string()]),
500 );
501
502 let payload = json!({
503 "code": 2,
504 "msg": "Invalid request",
505 "id": 1
506 });
507
508 let out = handler
509 .handle_raw_message(payload.to_string().into_bytes())
510 .await;
511 assert_eq!(out.len(), 1);
512 match &out[0] {
513 BinanceSpotPublicWsMessage::Error(err) => {
514 assert_eq!(err.code, 2);
515 assert_eq!(err.msg, "Invalid request");
516 }
517 other => panic!("expected Error variant, was {other:?}"),
518 }
519 assert_eq!(handler.pending_requests.len(), 0);
520 }
521
522 #[rstest]
523 fn test_subscription_responses_preserve_unsubscribe_intent_until_acknowledged() {
524 let signal = Arc::new(AtomicBool::new(false));
525 let request_id_counter = Arc::new(AtomicU64::new(3));
526 let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
527 let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
528 let subscriptions = SubscriptionState::new('@');
529 let topic = "btcusdt@trade";
530
531 let mut handler = BinanceSpotPublicWsHandler::new(
532 signal,
533 cmd_rx,
534 raw_rx,
535 subscriptions.clone(),
536 request_id_counter,
537 );
538 subscriptions.mark_subscribe(topic);
539 handler.pending_requests.insert(
540 1,
541 PendingSubscriptionRequest::subscribe(vec![topic.to_string()]),
542 );
543 subscriptions.mark_unsubscribe(topic);
544 handler.pending_requests.insert(
545 2,
546 PendingSubscriptionRequest::unsubscribe(vec![topic.to_string()]),
547 );
548
549 handler.handle_subscription_response(&json!({"result": null, "id": 1}));
550
551 assert_eq!(
552 subscriptions.pending_unsubscribe_topics(),
553 [topic.to_string()]
554 );
555 assert_eq!(subscriptions.len(), 0);
556 assert_eq!(handler.pending_requests.len(), 1);
557
558 handler.handle_subscription_response(&json!({"result": null, "id": 2}));
559
560 assert!(subscriptions.is_empty());
561 assert_eq!(handler.pending_requests.len(), 0);
562 }
563
564 #[rstest]
565 fn test_stale_responses_do_not_mutate_latest_subscription_request() {
566 let signal = Arc::new(AtomicBool::new(false));
567 let request_id_counter = Arc::new(AtomicU64::new(4));
568 let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
569 let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
570 let subscriptions = SubscriptionState::new('@');
571 let topic = "btcusdt@trade";
572
573 let mut handler = BinanceSpotPublicWsHandler::new(
574 signal,
575 cmd_rx,
576 raw_rx,
577 subscriptions.clone(),
578 request_id_counter,
579 );
580 subscriptions.mark_subscribe(topic);
581 handler.pending_requests.insert(
582 1,
583 PendingSubscriptionRequest::subscribe(vec![topic.to_string()]),
584 );
585 subscriptions.mark_unsubscribe(topic);
586 handler.pending_requests.insert(
587 2,
588 PendingSubscriptionRequest::unsubscribe(vec![topic.to_string()]),
589 );
590 subscriptions.mark_subscribe(topic);
591 handler.pending_requests.insert(
592 3,
593 PendingSubscriptionRequest::subscribe(vec![topic.to_string()]),
594 );
595
596 handler.handle_subscription_response(&json!({"result": null, "id": 3}));
597 handler.handle_subscription_response(&json!({"code": 2, "msg": "Stale failure", "id": 1}));
598 handler.handle_subscription_response(&json!({"result": null, "id": 2}));
599
600 assert_eq!(subscriptions.len(), 1);
601 assert_eq!(subscriptions.all_topics(), [topic]);
602 assert!(subscriptions.pending_subscribe_topics().is_empty());
603 assert!(subscriptions.pending_unsubscribe_topics().is_empty());
604 assert_eq!(handler.pending_requests.len(), 0);
605 }
606
607 #[rstest]
608 fn test_handle_stream_data_parses_book_ticker_without_event_type() {
609 let signal = Arc::new(AtomicBool::new(false));
610 let request_id_counter = Arc::new(AtomicU64::new(1));
611 let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
612 let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
613 let subscriptions = SubscriptionState::new('@');
614
615 let handler = BinanceSpotPublicWsHandler::new(
616 signal,
617 cmd_rx,
618 raw_rx,
619 subscriptions,
620 request_id_counter,
621 );
622
623 let payload = json!({
624 "stream": "btcusdt@bookTicker",
625 "data": {
626 "u": 94528182161_u64,
627 "s": "BTCUSDT",
628 "b": "73650.51000000",
629 "B": "2.95126000",
630 "a": "73650.52000000",
631 "A": "1.38108000"
632 }
633 });
634
635 let out = handler.handle_stream_data(&payload);
636 assert_eq!(out.len(), 1);
637 assert!(matches!(out[0], BinanceSpotPublicWsMessage::BookTicker(_)));
638 }
639
640 #[rstest]
641 fn test_handle_stream_data_emits_server_shutdown() {
642 let signal = Arc::new(AtomicBool::new(false));
643 let request_id_counter = Arc::new(AtomicU64::new(1));
644 let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
645 let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
646 let subscriptions = SubscriptionState::new('@');
647
648 let handler = BinanceSpotPublicWsHandler::new(
649 signal,
650 cmd_rx,
651 raw_rx,
652 subscriptions,
653 request_id_counter,
654 );
655
656 let payload = json!({"e": "serverShutdown", "E": 1_700_000_000_000_i64});
659
660 let out = handler.handle_stream_data(&payload);
661 assert_eq!(out.len(), 1);
662 assert!(matches!(
663 out[0],
664 BinanceSpotPublicWsMessage::ServerShutdown(_)
665 ));
666 }
667}