1use std::{
23 fmt::Debug,
24 sync::{
25 Arc,
26 atomic::{AtomicBool, AtomicU64, Ordering},
27 },
28};
29
30use ahash::AHashMap;
31use nautilus_network::{
32 RECONNECTED,
33 websocket::{SubscriptionState, WebSocketClient},
34};
35
36use super::{
37 messages::{
38 BinanceFuturesAccountConfigMsg, BinanceFuturesAccountUpdateMsg, BinanceFuturesAggTradeMsg,
39 BinanceFuturesAlgoUpdateMsg, BinanceFuturesBookTickerMsg, BinanceFuturesDepthUpdateMsg,
40 BinanceFuturesKlineMsg, BinanceFuturesLiquidationMsg, BinanceFuturesListenKeyExpiredMsg,
41 BinanceFuturesMarginCallMsg, BinanceFuturesMarkPriceMsg, BinanceFuturesOrderUpdateMsg,
42 BinanceFuturesTickerMsg, BinanceFuturesTradeLiteMsg, BinanceFuturesTradeMsg,
43 BinanceFuturesWsErrorMsg, BinanceFuturesWsErrorResponse, BinanceFuturesWsStreamsCommand,
44 BinanceFuturesWsStreamsMessage, BinanceFuturesWsSubscribeRequest,
45 BinanceFuturesWsSubscribeResponse,
46 },
47 parse_data::extract_event_type,
48};
49use crate::common::{
50 consts::BINANCE_RATE_LIMIT_KEY_SUBSCRIPTION,
51 enums::{BinanceWsEventType, BinanceWsMethod},
52 websocket::{
53 PendingSubscriptionRequest, PendingSubscriptionRequests, reset_requests_after_reconnect,
54 },
55};
56
57pub struct BinanceFuturesDataWsFeedHandler {
63 #[allow(dead_code)]
64 signal: Arc<AtomicBool>,
65 cmd_rx: tokio::sync::mpsc::UnboundedReceiver<BinanceFuturesWsStreamsCommand>,
66 raw_rx: tokio::sync::mpsc::UnboundedReceiver<Vec<u8>>,
67 #[allow(dead_code)]
68 out_tx: tokio::sync::mpsc::UnboundedSender<BinanceFuturesWsStreamsMessage>,
69 inner: Option<WebSocketClient>,
70 subscriptions_state: SubscriptionState,
71 request_id_counter: Arc<AtomicU64>,
72 pending_requests: PendingSubscriptionRequests,
73 unsubscribe_correlations: AHashMap<u64, (u64, Vec<String>)>,
74}
75
76impl Debug for BinanceFuturesDataWsFeedHandler {
77 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
78 f.debug_struct(stringify!(BinanceFuturesDataWsFeedHandler))
79 .field("pending_requests", &self.pending_requests.len())
80 .finish_non_exhaustive()
81 }
82}
83
84impl BinanceFuturesDataWsFeedHandler {
85 pub fn new(
87 signal: Arc<AtomicBool>,
88 cmd_rx: tokio::sync::mpsc::UnboundedReceiver<BinanceFuturesWsStreamsCommand>,
89 raw_rx: tokio::sync::mpsc::UnboundedReceiver<Vec<u8>>,
90 out_tx: tokio::sync::mpsc::UnboundedSender<BinanceFuturesWsStreamsMessage>,
91 subscriptions_state: SubscriptionState,
92 request_id_counter: Arc<AtomicU64>,
93 ) -> Self {
94 Self {
95 signal,
96 cmd_rx,
97 raw_rx,
98 out_tx,
99 inner: None,
100 subscriptions_state,
101 request_id_counter,
102 pending_requests: PendingSubscriptionRequests::default(),
103 unsubscribe_correlations: AHashMap::new(),
104 }
105 }
106
107 pub async fn next(&mut self) -> Option<BinanceFuturesWsStreamsMessage> {
111 loop {
112 if self.signal.load(Ordering::Relaxed) {
113 return None;
114 }
115
116 tokio::select! {
117 Some(cmd) = self.cmd_rx.recv() => {
118 self.handle_command(cmd).await;
119 }
120 Some(raw) = self.raw_rx.recv() => {
121 if let Some(msg) = self.handle_raw_message(raw).await {
122 return Some(msg);
123 }
124 }
125 else => {
126 return None;
127 }
128 }
129 }
130 }
131
132 async fn handle_command(&mut self, cmd: BinanceFuturesWsStreamsCommand) {
133 match cmd {
134 BinanceFuturesWsStreamsCommand::SetClient(client) => {
135 self.inner = Some(client);
136 }
137 BinanceFuturesWsStreamsCommand::Disconnect => {
138 if let Some(client) = &self.inner {
139 let () = client.disconnect().await;
140 }
141 self.inner = None;
142 }
143 BinanceFuturesWsStreamsCommand::Subscribe { streams } => {
144 self.send_subscribe(streams).await;
145 }
146 BinanceFuturesWsStreamsCommand::Unsubscribe {
147 streams,
148 correlation,
149 } => {
150 self.send_unsubscribe(streams, correlation).await;
151 }
152 }
153 }
154
155 async fn send_subscribe(&mut self, streams: Vec<String>) {
156 for stream in &streams {
157 self.subscriptions_state.mark_subscribe(stream);
158 }
159
160 let Some(client) = &self.inner else {
161 log::warn!("Cannot subscribe: no client connected");
162 return;
163 };
164
165 let request_id = self.request_id_counter.fetch_add(1, Ordering::Relaxed);
166
167 let request = BinanceFuturesWsSubscribeRequest {
168 method: BinanceWsMethod::Subscribe,
169 params: streams.clone(),
170 id: request_id,
171 };
172
173 let json = match serde_json::to_string(&request) {
174 Ok(j) => j,
175 Err(e) => {
176 log::error!("Failed to serialize subscribe request: {e}");
177 return;
178 }
179 };
180
181 self.pending_requests
182 .insert(request_id, PendingSubscriptionRequest::subscribe(streams));
183
184 if let Err(e) = client
185 .send_text(json, Some(BINANCE_RATE_LIMIT_KEY_SUBSCRIPTION.as_slice()))
186 .await
187 {
188 if let Some(request) = self.pending_requests.take(request_id) {
189 request.mark_failure(&self.subscriptions_state);
190 }
191 log::error!("Failed to send subscribe request: {e}");
192 }
193 }
194
195 async fn send_unsubscribe(&mut self, streams: Vec<String>, correlation: Option<u64>) {
196 for stream in &streams {
197 self.subscriptions_state.mark_unsubscribe(stream);
198 }
199
200 let request_id = self.request_id_counter.fetch_add(1, Ordering::Relaxed);
201
202 if let Some(correlation) = correlation {
205 self.unsubscribe_correlations
206 .insert(request_id, (correlation, streams.clone()));
207 }
208
209 let Some(client) = &self.inner else {
210 log::warn!("Cannot unsubscribe: no client connected");
211 return;
212 };
213
214 let request = BinanceFuturesWsSubscribeRequest {
215 method: BinanceWsMethod::Unsubscribe,
216 params: streams.clone(),
217 id: request_id,
218 };
219
220 let json = match serde_json::to_string(&request) {
221 Ok(j) => j,
222 Err(e) => {
223 log::error!("Failed to serialize unsubscribe request: {e}");
224 return;
225 }
226 };
227
228 self.pending_requests
229 .insert(request_id, PendingSubscriptionRequest::unsubscribe(streams));
230
231 if let Err(e) = client
232 .send_text(json, Some(BINANCE_RATE_LIMIT_KEY_SUBSCRIPTION.as_slice()))
233 .await
234 {
235 self.pending_requests.take(request_id);
236 log::error!("Failed to send unsubscribe request: {e}");
237 }
238 }
239
240 async fn handle_raw_message(&mut self, raw: Vec<u8>) -> Option<BinanceFuturesWsStreamsMessage> {
241 if let Ok(text) = std::str::from_utf8(&raw)
242 && text == RECONNECTED
243 {
244 let abandoned = self
245 .unsubscribe_correlations
246 .values()
247 .map(|(correlation, _)| *correlation)
248 .collect();
249 self.unsubscribe_correlations.clear();
250 reset_requests_after_reconnect(&mut self.pending_requests, &self.subscriptions_state);
251 log::debug!("WebSocket reconnected signal received");
252 return Some(BinanceFuturesWsStreamsMessage::Reconnected(abandoned));
253 }
254
255 let json: serde_json::Value = match serde_json::from_slice(&raw) {
256 Ok(j) => j,
257 Err(e) => {
258 log::warn!("Failed to parse JSON message: {e}");
259 return None;
260 }
261 };
262
263 if json.get("result").is_some() || json.get("id").is_some() {
264 return self.handle_subscription_response(&json);
265 }
266
267 if let Some(code) = json.get("code")
268 && let Some(code) = code.as_i64()
269 {
270 let msg = json
271 .get("msg")
272 .and_then(|m| m.as_str())
273 .unwrap_or("Unknown error")
274 .to_string();
275 return Some(BinanceFuturesWsStreamsMessage::Error(
276 BinanceFuturesWsErrorMsg { code, msg },
277 ));
278 }
279
280 self.handle_stream_data(&json)
281 }
282
283 fn handle_subscription_response(
284 &mut self,
285 json: &serde_json::Value,
286 ) -> Option<BinanceFuturesWsStreamsMessage> {
287 if let Ok(error) = serde_json::from_value::<BinanceFuturesWsErrorResponse>(json.clone()) {
288 if let Some(id) = error.id
289 && let Some(request) = self.pending_requests.take(id)
290 {
291 request.mark_failure(&self.subscriptions_state);
294 }
295 log::warn!(
296 "WebSocket error response: code={}, msg={}",
297 error.code,
298 error.msg
299 );
300 } else if let Ok(response) =
301 serde_json::from_value::<BinanceFuturesWsSubscribeResponse>(json.clone())
302 {
303 if response.result.is_none() {
304 if let Some(request) = self.pending_requests.take(response.id) {
305 request.confirm(&self.subscriptions_state);
306 log::debug!("Subscription request confirmed: request={request:?}");
307 }
308
309 if let Some((correlation, streams)) =
313 self.unsubscribe_correlations.remove(&response.id)
314 {
315 return Some(BinanceFuturesWsStreamsMessage::Unsubscribed {
316 streams,
317 correlation: Some(correlation),
318 });
319 }
320 } else if let Some(request) = self.pending_requests.take(response.id) {
321 request.mark_failure(&self.subscriptions_state);
324 log::warn!(
325 "Subscription request failed: request={request:?}, result={:?}",
326 response.result
327 );
328 }
329 }
330
331 None
332 }
333
334 fn handle_stream_data(
335 &self,
336 json: &serde_json::Value,
337 ) -> Option<BinanceFuturesWsStreamsMessage> {
338 let event_type = extract_event_type(json)?;
339
340 match event_type {
341 BinanceWsEventType::AggTrade => {
342 serde_json::from_value::<BinanceFuturesAggTradeMsg>(json.clone())
343 .map(BinanceFuturesWsStreamsMessage::AggTrade)
344 .map_err(|e| log::warn!("Failed to parse aggregate trade: {e}"))
345 .ok()
346 }
347 BinanceWsEventType::Trade => {
348 serde_json::from_value::<BinanceFuturesTradeMsg>(json.clone())
349 .map(BinanceFuturesWsStreamsMessage::Trade)
350 .map_err(|e| log::warn!("Failed to parse trade: {e}"))
351 .ok()
352 }
353 BinanceWsEventType::BookTicker => {
354 serde_json::from_value::<BinanceFuturesBookTickerMsg>(json.clone())
355 .map(BinanceFuturesWsStreamsMessage::BookTicker)
356 .map_err(|e| log::warn!("Failed to parse book ticker: {e}"))
357 .ok()
358 }
359 BinanceWsEventType::DepthUpdate => {
360 serde_json::from_value::<BinanceFuturesDepthUpdateMsg>(json.clone())
361 .map(BinanceFuturesWsStreamsMessage::DepthUpdate)
362 .map_err(|e| log::warn!("Failed to parse depth update: {e}"))
363 .ok()
364 }
365 BinanceWsEventType::MarkPriceUpdate => {
366 serde_json::from_value::<BinanceFuturesMarkPriceMsg>(json.clone())
367 .map(BinanceFuturesWsStreamsMessage::MarkPrice)
368 .map_err(|e| log::warn!("Failed to parse mark price: {e}"))
369 .ok()
370 }
371 BinanceWsEventType::Kline => {
372 serde_json::from_value::<BinanceFuturesKlineMsg>(json.clone())
373 .map(BinanceFuturesWsStreamsMessage::Kline)
374 .map_err(|e| log::warn!("Failed to parse kline: {e}"))
375 .ok()
376 }
377 BinanceWsEventType::ForceOrder => {
378 serde_json::from_value::<BinanceFuturesLiquidationMsg>(json.clone())
379 .map(BinanceFuturesWsStreamsMessage::ForceOrder)
380 .map_err(|e| log::warn!("Failed to parse force order: {e}"))
381 .ok()
382 }
383 BinanceWsEventType::Ticker24Hr => {
384 serde_json::from_value::<BinanceFuturesTickerMsg>(json.clone())
385 .map(BinanceFuturesWsStreamsMessage::Ticker)
386 .map_err(|e| log::warn!("Failed to parse ticker: {e}"))
387 .ok()
388 }
389 BinanceWsEventType::MiniTicker24Hr => {
390 log::debug!("Mini ticker not yet supported, skipping");
391 None
392 }
393 BinanceWsEventType::AccountUpdate => {
394 serde_json::from_value::<BinanceFuturesAccountUpdateMsg>(json.clone())
395 .map(|msg| {
396 log::debug!(
397 "Account update: reason={:?}, balances={}, positions={}",
398 msg.account.reason,
399 msg.account.balances.len(),
400 msg.account.positions.len()
401 );
402 BinanceFuturesWsStreamsMessage::AccountUpdate(msg)
403 })
404 .map_err(|e| log::warn!("Failed to parse account update: {e}"))
405 .ok()
406 }
407 BinanceWsEventType::OrderTradeUpdate => {
408 serde_json::from_value::<BinanceFuturesOrderUpdateMsg>(json.clone())
409 .map(|msg| {
410 log::debug!(
411 "Order update: symbol={}, order_id={}, exec={:?}, status={:?}",
412 msg.order.symbol,
413 msg.order.order_id,
414 msg.order.execution_type,
415 msg.order.order_status
416 );
417 BinanceFuturesWsStreamsMessage::OrderUpdate(Box::new(msg))
418 })
419 .map_err(|e| log::warn!("Failed to parse order update: {e}"))
420 .ok()
421 }
422 BinanceWsEventType::TradeLite => {
423 serde_json::from_value::<BinanceFuturesTradeLiteMsg>(json.clone())
424 .map(|msg| {
425 log::debug!(
426 "Trade lite: symbol={}, order_id={}, trade_id={}",
427 msg.symbol,
428 msg.order_id,
429 msg.trade_id
430 );
431 BinanceFuturesWsStreamsMessage::TradeLite(Box::new(msg))
432 })
433 .map_err(|e| log::warn!("Failed to parse trade lite: {e}"))
434 .ok()
435 }
436 BinanceWsEventType::AlgoUpdate => {
437 serde_json::from_value::<BinanceFuturesAlgoUpdateMsg>(json.clone())
438 .map(|msg| {
439 log::debug!(
440 "Algo order update: symbol={}, algo_id={}, status={:?}",
441 msg.algo_order.symbol,
442 msg.algo_order.algo_id,
443 msg.algo_order.algo_status
444 );
445 BinanceFuturesWsStreamsMessage::AlgoUpdate(Box::new(msg))
446 })
447 .map_err(|e| log::warn!("Failed to parse algo order update: {e}"))
448 .ok()
449 }
450 BinanceWsEventType::MarginCall => {
451 serde_json::from_value::<BinanceFuturesMarginCallMsg>(json.clone())
452 .map(|msg| {
453 log::warn!(
454 "Margin call: cross_wallet_balance={}, positions_at_risk={}",
455 msg.cross_wallet_balance,
456 msg.positions.len()
457 );
458 BinanceFuturesWsStreamsMessage::MarginCall(msg)
459 })
460 .map_err(|e| log::warn!("Failed to parse margin call: {e}"))
461 .ok()
462 }
463 BinanceWsEventType::AccountConfigUpdate => {
464 serde_json::from_value::<BinanceFuturesAccountConfigMsg>(json.clone())
465 .map(|msg| {
466 if let Some(ref lc) = msg.leverage_config {
467 log::debug!(
468 "Account config update: symbol={}, leverage={}",
469 lc.symbol,
470 lc.leverage
471 );
472 }
473 BinanceFuturesWsStreamsMessage::AccountConfigUpdate(msg)
474 })
475 .map_err(|e| log::warn!("Failed to parse account config update: {e}"))
476 .ok()
477 }
478 BinanceWsEventType::ListenKeyExpired => {
479 if let Ok(msg) =
480 serde_json::from_value::<BinanceFuturesListenKeyExpiredMsg>(json.clone())
481 {
482 log::warn!("Listen key expired at {}", msg.event_time);
483 }
484 Some(BinanceFuturesWsStreamsMessage::ListenKeyExpired)
485 }
486 BinanceWsEventType::Unknown => {
487 log::warn!("Unknown event type in message: {json}");
488 None
489 }
490 }
491 }
492}
493
494#[cfg(test)]
495mod tests {
496 use rstest::rstest;
497
498 use super::*;
499
500 #[rstest]
501 #[tokio::test]
502 async fn test_subscription_intent_is_preserved_without_active_client() {
503 let signal = Arc::new(AtomicBool::new(false));
504 let request_id_counter = Arc::new(AtomicU64::new(1));
505 let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
506 let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
507 let (out_tx, _out_rx) = tokio::sync::mpsc::unbounded_channel();
508 let subscriptions = SubscriptionState::new('@');
509 let subscribe_topic = "btcusdt@aggTrade";
510 let unsubscribe_topic = "ethusdt@aggTrade";
511 subscriptions.mark_subscribe(unsubscribe_topic);
512 subscriptions.confirm_subscribe(unsubscribe_topic);
513
514 let mut handler = BinanceFuturesDataWsFeedHandler::new(
515 signal,
516 cmd_rx,
517 raw_rx,
518 out_tx,
519 subscriptions.clone(),
520 request_id_counter,
521 );
522
523 handler
524 .send_subscribe(vec![subscribe_topic.to_string()])
525 .await;
526 handler
527 .send_unsubscribe(vec![unsubscribe_topic.to_string()], None)
528 .await;
529
530 assert_eq!(
531 subscriptions.pending_subscribe_topics(),
532 [subscribe_topic.to_string()]
533 );
534 assert_eq!(
535 subscriptions.pending_unsubscribe_topics(),
536 [unsubscribe_topic.to_string()]
537 );
538 assert_eq!(subscriptions.len(), 0);
539 assert_eq!(handler.pending_requests.len(), 0);
540 }
541
542 #[rstest]
543 #[tokio::test]
544 async fn test_correlated_unsubscribe_without_client_is_abandoned_on_reconnect() {
545 let signal = Arc::new(AtomicBool::new(false));
546 let request_id_counter = Arc::new(AtomicU64::new(1));
547 let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
548 let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
549 let (out_tx, mut out_rx) = tokio::sync::mpsc::unbounded_channel();
550 let subscriptions = SubscriptionState::new('@');
551
552 let mut handler = BinanceFuturesDataWsFeedHandler::new(
553 signal,
554 cmd_rx,
555 raw_rx,
556 out_tx,
557 subscriptions,
558 request_id_counter,
559 );
560
561 handler
564 .send_unsubscribe(vec!["btcusdt@depth5@100ms".to_string()], Some(7))
565 .await;
566
567 let msg = handler
568 .handle_raw_message(RECONNECTED.as_bytes().to_vec())
569 .await;
570 let Some(BinanceFuturesWsStreamsMessage::Reconnected(abandoned)) = msg else {
571 panic!("expected Reconnected");
572 };
573 assert_eq!(abandoned, vec![7]);
574 assert!(out_rx.try_recv().is_err());
575 }
576
577 #[rstest]
578 fn test_error_responses_preserve_subscription_intent() {
579 let signal = Arc::new(AtomicBool::new(false));
580 let request_id_counter = Arc::new(AtomicU64::new(3));
581 let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
582 let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
583 let (out_tx, _out_rx) = tokio::sync::mpsc::unbounded_channel();
584 let subscriptions = SubscriptionState::new('@');
585 let subscribe_topic = "btcusdt@aggTrade";
586 let unsubscribe_topic = "ethusdt@aggTrade";
587 subscriptions.mark_subscribe(subscribe_topic);
588 subscriptions.mark_subscribe(unsubscribe_topic);
589 subscriptions.confirm_subscribe(unsubscribe_topic);
590 subscriptions.mark_unsubscribe(unsubscribe_topic);
591
592 let mut handler = BinanceFuturesDataWsFeedHandler::new(
593 signal,
594 cmd_rx,
595 raw_rx,
596 out_tx,
597 subscriptions.clone(),
598 request_id_counter,
599 );
600 handler.pending_requests.insert(
601 1,
602 PendingSubscriptionRequest::subscribe(vec![subscribe_topic.to_string()]),
603 );
604 handler.pending_requests.insert(
605 2,
606 PendingSubscriptionRequest::unsubscribe(vec![unsubscribe_topic.to_string()]),
607 );
608 let error = serde_json::json!({"code": 2, "msg": "Invalid request", "id": 1});
609
610 handler.handle_subscription_response(&error);
611
612 assert_eq!(
613 subscriptions.pending_subscribe_topics(),
614 [subscribe_topic.to_string()]
615 );
616 assert_eq!(
617 subscriptions.pending_unsubscribe_topics(),
618 [unsubscribe_topic.to_string()]
619 );
620 assert_eq!(subscriptions.len(), 0);
621 assert_eq!(handler.pending_requests.len(), 1);
622
623 let error = serde_json::json!({"code": 2, "msg": "Invalid request", "id": 2});
624 handler.handle_subscription_response(&error);
625
626 assert_eq!(
627 subscriptions.pending_subscribe_topics(),
628 [subscribe_topic.to_string()]
629 );
630 assert_eq!(
631 subscriptions.pending_unsubscribe_topics(),
632 [unsubscribe_topic.to_string()]
633 );
634 assert_eq!(subscriptions.len(), 0);
635 assert_eq!(handler.pending_requests.len(), 0);
636 }
637}