nautilus_binance/futures/websocket/streams/
client.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! Binance Futures WebSocket client for JSON market data streams.
17//!
18//! ## Connection Details
19//!
20//! - USD-M Endpoint: `wss://fstream.binance.com/market/ws`
21//! - COIN-M Endpoint: `wss://dstream.binance.com/ws`
22//! - Max streams: 200 per connection
23//! - Max connections: 20 per pool (up to 4,000 total streams)
24//! - Connection validity: 24 hours
25//! - Ping/pong: Every 3 minutes
26
27use std::{
28    fmt::Debug,
29    sync::{
30        Arc,
31        atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering},
32    },
33};
34
35use futures_util::Stream;
36use nautilus_core::{
37    AtomicMap,
38    string::secret::{REDACTED, SecretString},
39};
40use nautilus_live::{
41    SocketControl, SocketControlFactory,
42    task::{TaskJoinOutcome, TaskSlot, finish_task},
43};
44use nautilus_model::instruments::{Instrument, InstrumentAny};
45use nautilus_network::{
46    http::create_standard_nautilus_headers,
47    mode::ConnectionMode,
48    websocket::{
49        PingHandler, SubscriptionState, TransportBackend, WebSocketClient, WebSocketConfig,
50        channel_message_handler,
51    },
52};
53use parking_lot::Mutex;
54use tokio_tungstenite::tungstenite::Message;
55use tokio_util::sync::CancellationToken;
56use ustr::Ustr;
57
58use super::{
59    error::{BinanceWsError, BinanceWsResult},
60    handler::BinanceFuturesDataWsFeedHandler,
61    messages::{BinanceFuturesWsStreamsCommand, BinanceFuturesWsStreamsMessage},
62};
63use crate::common::{
64    consts::{
65        BINANCE_API_KEY_HEADER, BINANCE_RATE_LIMIT_KEY_SUBSCRIPTION, BINANCE_WS_CONNECTION_QUOTA,
66        BINANCE_WS_SUBSCRIPTION_QUOTA,
67    },
68    credential::SigningCredential,
69    enums::{BinanceEnvironment, BinanceProductType},
70    urls::get_ws_base_url,
71};
72
73/// Maximum streams per WebSocket connection for Futures.
74pub const MAX_STREAMS_PER_CONNECTION: usize = 200;
75
76/// Maximum connections per pool.
77const MAX_CONNECTIONS: usize = 20;
78
79// State for a single WebSocket connection within the pool
80struct ConnectionSlot {
81    cmd_tx: tokio::sync::mpsc::UnboundedSender<BinanceFuturesWsStreamsCommand>,
82    streams: Vec<String>,
83    subscriptions_state: SubscriptionState,
84    handler_task: TaskSlot<()>,
85    bytes_task: TaskSlot<()>,
86    cancellation_token: CancellationToken,
87    connection_mode: Arc<AtomicU8>,
88    socket_control: Option<SocketControl>,
89    shutdown_errors: Vec<String>,
90}
91
92/// Binance Futures WebSocket client for JSON market data streams.
93///
94/// Manages a pool of up to 20 connections, each supporting up to 200 streams.
95/// New connections are created automatically when subscribing exceeds the current
96/// connection's stream limit. All connections feed into a single output stream,
97/// transparent to the data client.
98#[derive(Clone)]
99pub struct BinanceFuturesWebSocketClient {
100    url: SecretString,
101    product_type: BinanceProductType,
102    credential: Option<Arc<SigningCredential>>,
103    heartbeat: Option<u64>,
104    signal: Arc<AtomicBool>,
105    slots: Arc<ConnectionSlots>,
106    connect_lock: Arc<tokio::sync::Mutex<()>>,
107    out_tx: Arc<Mutex<Option<tokio::sync::mpsc::UnboundedSender<BinanceFuturesWsStreamsMessage>>>>,
108    out_rx:
109        Arc<Mutex<Option<tokio::sync::mpsc::UnboundedReceiver<BinanceFuturesWsStreamsMessage>>>>,
110    request_id_counter: Arc<AtomicU64>,
111    instruments_cache: Arc<AtomicMap<Ustr, InstrumentAny>>,
112    transport_backend: TransportBackend,
113    proxy_url: Option<SecretString>,
114    socket_factory: Option<SocketControlFactory>,
115    socket_endpoint: Option<String>,
116}
117
118impl Debug for BinanceFuturesWebSocketClient {
119    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
120        f.debug_struct(stringify!(BinanceFuturesWebSocketClient))
121            .field("url", &REDACTED)
122            .field("product_type", &self.product_type)
123            .field("credential", &self.credential.as_ref().map(|_| REDACTED))
124            .field("heartbeat", &self.heartbeat)
125            .finish_non_exhaustive()
126    }
127}
128
129impl BinanceFuturesWebSocketClient {
130    /// Creates a new [`BinanceFuturesWebSocketClient`] instance.
131    ///
132    /// # Errors
133    ///
134    /// Returns an error if:
135    /// - `product_type` is not a futures type (UsdM or CoinM).
136    /// - Credential creation fails.
137    pub fn new(
138        product_type: BinanceProductType,
139        environment: BinanceEnvironment,
140        api_key: Option<String>,
141        api_secret: Option<String>,
142        url_override: Option<String>,
143        heartbeat: Option<u64>,
144        transport_backend: TransportBackend,
145    ) -> anyhow::Result<Self> {
146        match product_type {
147            BinanceProductType::UsdM | BinanceProductType::CoinM => {}
148            _ => {
149                anyhow::bail!(
150                    "BinanceFuturesWebSocketClient requires UsdM or CoinM product type, was {product_type:?}"
151                );
152            }
153        }
154
155        let url = SecretString::from(
156            url_override.unwrap_or_else(|| get_ws_base_url(product_type, environment).to_string()),
157        );
158
159        let credential = match (api_key, api_secret) {
160            (Some(key), Some(secret)) => Some(Arc::new(SigningCredential::new(key, secret))),
161            _ => None,
162        };
163
164        Ok(Self {
165            url,
166            product_type,
167            credential,
168            heartbeat,
169            signal: Arc::new(AtomicBool::new(false)),
170            slots: Arc::new(ConnectionSlots(Mutex::new(Vec::new()))),
171            connect_lock: Arc::new(tokio::sync::Mutex::new(())),
172            out_tx: Arc::new(Mutex::new(None)),
173            out_rx: Arc::new(Mutex::new(None)),
174            request_id_counter: Arc::new(AtomicU64::new(1)),
175            instruments_cache: Arc::new(AtomicMap::new()),
176            transport_backend,
177            proxy_url: None,
178            socket_factory: None,
179            socket_endpoint: None,
180        })
181    }
182
183    /// Configures the proxy used by every connection in the stream pool.
184    #[must_use]
185    pub fn with_proxy(mut self, proxy_url: Option<String>) -> Self {
186        self.proxy_url = proxy_url.map(SecretString::from);
187        self
188    }
189
190    /// Configures socket state reporting and reconnect control for the stream pool.
191    #[must_use]
192    pub fn with_socket_control(
193        mut self,
194        factory: SocketControlFactory,
195        endpoint: impl Into<String>,
196    ) -> Self {
197        self.socket_factory = Some(factory);
198        self.socket_endpoint = Some(endpoint.into());
199        self
200    }
201
202    /// Returns the product type (UsdM or CoinM).
203    #[must_use]
204    pub const fn product_type(&self) -> BinanceProductType {
205        self.product_type
206    }
207
208    /// Returns whether any connection in the pool is active.
209    #[must_use]
210    pub fn is_active(&self) -> bool {
211        let slots = self.slots.lock();
212        slots
213            .iter()
214            .any(|s| s.connection_mode.load(Ordering::Relaxed) == ConnectionMode::Active as u8)
215    }
216
217    /// Returns whether all connections in the pool are closed.
218    #[must_use]
219    pub fn is_closed(&self) -> bool {
220        let slots = self.slots.lock();
221        slots.is_empty()
222            || slots
223                .iter()
224                .all(|s| s.connection_mode.load(Ordering::Relaxed) == ConnectionMode::Closed as u8)
225    }
226
227    /// Returns the total number of confirmed subscriptions across all connections.
228    #[must_use]
229    pub fn subscription_count(&self) -> usize {
230        let slots = self.slots.lock();
231        slots.iter().map(|s| s.subscriptions_state.len()).sum()
232    }
233
234    /// Connects the first WebSocket connection in the pool.
235    ///
236    /// # Errors
237    ///
238    /// Returns an error if connection fails.
239    pub async fn connect(&mut self) -> BinanceWsResult<()> {
240        let connect_lock = Arc::clone(&self.connect_lock);
241        let _connect_guard = connect_lock.lock().await;
242
243        if !self.slots.lock().is_empty() {
244            self.close_connections().await?;
245        }
246
247        {
248            let _slots = self.slots.lock();
249            self.signal.store(false, Ordering::Release);
250        }
251
252        let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel();
253        *self.out_tx.lock() = Some(out_tx);
254        *self.out_rx.lock() = Some(out_rx);
255
256        let slot = self.create_connection(0).await?;
257        let shutdown = {
258            let mut slots = self.slots.lock();
259            let shutdown = self.signal.load(Ordering::Acquire);
260            slots.push(slot);
261            shutdown
262        };
263
264        if shutdown {
265            let rollback = self.close_connections().await;
266            return Err(BinanceWsError::ClientError(match rollback {
267                Ok(()) => "Binance Futures stream pool shutdown began during connect".to_string(),
268                Err(e) => format!(
269                    "Binance Futures stream pool shutdown began during connect; rollback failed: {e}"
270                ),
271            }));
272        }
273
274        log::debug!(
275            "Connected to Binance Futures stream pool: product_type={:?}",
276            self.product_type
277        );
278        Ok(())
279    }
280
281    /// Closes all WebSocket connections in the pool.
282    ///
283    /// # Errors
284    ///
285    /// Returns an error if disconnect fails.
286    pub async fn close(&mut self) -> BinanceWsResult<()> {
287        self.begin_shutdown();
288        let connect_lock = Arc::clone(&self.connect_lock);
289        let _connect_guard = connect_lock.lock().await;
290        self.close_connections().await
291    }
292
293    pub(crate) fn begin_shutdown(&self) {
294        let slots = self.slots.lock();
295        self.signal.store(true, Ordering::Release);
296
297        for slot in slots.iter() {
298            if let Some(control) = &slot.socket_control {
299                control.deregister();
300            }
301            slot.cancellation_token.cancel();
302            let _ = slot.cmd_tx.send(BinanceFuturesWsStreamsCommand::Disconnect);
303        }
304    }
305
306    async fn close_connections(&self) -> BinanceWsResult<()> {
307        self.begin_shutdown();
308
309        let mut batch = ConnectionSlotBatch::take(&self.slots);
310        let mut index = batch.slots.len();
311        while index > 0 {
312            index -= 1;
313            let slot = &mut batch.slots[index];
314            if let Some(control) = &slot.socket_control {
315                control.deregister();
316            }
317            slot.cancellation_token.cancel();
318            let _ = slot.cmd_tx.send(BinanceFuturesWsStreamsCommand::Disconnect);
319            if let Some(error) =
320                finish_slot_task(&mut slot.handler_task, "Futures stream handler").await
321            {
322                slot.shutdown_errors.push(error);
323            }
324
325            if let Some(error) = finish_slot_task(&mut slot.bytes_task, "Futures byte stream").await
326            {
327                slot.shutdown_errors.push(error);
328            }
329
330            if slot.handler_task.is_none()
331                && slot.bytes_task.is_none()
332                && slot.shutdown_errors.is_empty()
333            {
334                batch.slots.remove(index);
335            }
336        }
337
338        *self.out_tx.lock() = None;
339        *self.out_rx.lock() = None;
340
341        let errors = batch
342            .slots
343            .iter_mut()
344            .flat_map(|slot| std::mem::take(&mut slot.shutdown_errors))
345            .collect::<Vec<_>>();
346        batch
347            .slots
348            .retain(|slot| slot.handler_task.is_some() || slot.bytes_task.is_some());
349
350        if !errors.is_empty() {
351            return Err(BinanceWsError::ClientError(errors.join("; ")));
352        }
353        log::debug!("Disconnected from Binance Futures stream pool");
354        Ok(())
355    }
356
357    /// Subscribes to the specified streams.
358    ///
359    /// Streams are distributed across pool connections. New connections are created
360    /// automatically when existing ones reach the 200-stream limit, up to a maximum
361    /// of 20 connections.
362    ///
363    /// # Errors
364    ///
365    /// Returns an error if the pool is exhausted or command delivery fails.
366    pub async fn subscribe(&self, streams: Vec<String>) -> BinanceWsResult<()> {
367        // Serialize all phases so concurrent subscribers see a consistent
368        // pool state and can't trigger spurious `Pool exhausted`.
369        let _connect_guard = self.connect_lock.lock().await;
370
371        // Phase 1: filter already-subscribed streams (brief lock).
372        let new_streams: Vec<String> = {
373            let slots = self.slots.lock();
374
375            if self.signal.load(Ordering::Acquire) {
376                return Err(BinanceWsError::ClientError(
377                    "Binance Futures stream pool is shutting down".to_string(),
378                ));
379            }
380            streams
381                .into_iter()
382                .filter(|s| !slots.iter().any(|slot| slot.streams.contains(s)))
383                .collect()
384        };
385
386        if new_streams.is_empty() {
387            return Ok(());
388        }
389
390        // Phase 2: create connections if needed.
391
392        loop {
393            let (remaining_capacity, slot_count) = {
394                let slots = self.slots.lock();
395                let cap: usize = slots
396                    .iter()
397                    .map(|s| MAX_STREAMS_PER_CONNECTION - s.streams.len())
398                    .sum();
399                (cap, slots.len())
400            };
401
402            if remaining_capacity >= new_streams.len() || slot_count >= MAX_CONNECTIONS {
403                break;
404            }
405
406            let new_slot = self.create_connection(slot_count).await?;
407            let (slot_count, shutdown) = {
408                let mut slots = self.slots.lock();
409                let shutdown = self.signal.load(Ordering::Acquire);
410                slots.push(new_slot);
411                (slots.len(), shutdown)
412            };
413
414            if shutdown {
415                let client = self.clone();
416                let rollback = client.close_connections().await;
417                return Err(BinanceWsError::ClientError(match rollback {
418                    Ok(()) => {
419                        "Binance Futures stream pool shutdown began during subscribe".to_string()
420                    }
421                    Err(e) => format!(
422                        "Binance Futures stream pool shutdown began during subscribe; rollback failed: {e}"
423                    ),
424                }));
425            }
426            log::debug!(
427                "Pool slot {} connected: product_type={:?}",
428                slot_count - 1,
429                self.product_type
430            );
431        }
432
433        // Phase 3: assign streams to slots and send commands (brief lock).
434        // Stage assignments first so a capacity error leaves slots unchanged.
435        let mut slots = self.slots.lock();
436
437        if self.signal.load(Ordering::Acquire) {
438            return Err(BinanceWsError::ClientError(
439                "Binance Futures stream pool is shutting down".to_string(),
440            ));
441        }
442        let mut slot_batches: Vec<(usize, Vec<String>)> = Vec::new();
443        let mut slot_counts: Vec<usize> = slots.iter().map(|s| s.streams.len()).collect();
444
445        for stream in &new_streams {
446            let slot_idx = slot_counts
447                .iter()
448                .position(|&count| count < MAX_STREAMS_PER_CONNECTION)
449                .ok_or_else(|| {
450                    let max_total = MAX_CONNECTIONS * MAX_STREAMS_PER_CONNECTION;
451                    BinanceWsError::ClientError(format!(
452                        "Pool exhausted: {max_total} total subscriptions \
453                         ({MAX_CONNECTIONS} connections x {MAX_STREAMS_PER_CONNECTION} streams)"
454                    ))
455                })?;
456
457            slot_counts[slot_idx] += 1;
458
459            if let Some(batch) = slot_batches.iter_mut().find(|(i, _)| *i == slot_idx) {
460                batch.1.push(stream.clone());
461            } else {
462                slot_batches.push((slot_idx, vec![stream.clone()]));
463            }
464        }
465
466        // Send commands first; only update slot state on success
467        for (slot_idx, batch) in &slot_batches {
468            slots[*slot_idx]
469                .cmd_tx
470                .send(BinanceFuturesWsStreamsCommand::Subscribe {
471                    streams: batch.clone(),
472                })
473                .map_err(|e| {
474                    BinanceWsError::ClientError(format!(
475                        "Handler not available for pool slot {slot_idx}: {e}"
476                    ))
477                })?;
478            slots[*slot_idx].streams.extend(batch.iter().cloned());
479        }
480
481        Ok(())
482    }
483
484    /// Unsubscribes from the specified streams.
485    ///
486    /// Returns the streams for which an unsubscribe command was delivered. Streams not
487    /// assigned to a pool connection are skipped and absent from the result.
488    ///
489    /// # Errors
490    ///
491    /// Returns an error if command delivery fails.
492    pub async fn unsubscribe(&self, streams: Vec<String>) -> BinanceWsResult<Vec<String>> {
493        self.unsubscribe_inner(streams, None).await
494    }
495
496    /// Unsubscribes with a correlation ID that the venue confirmation carries back.
497    pub(crate) async fn unsubscribe_correlated(
498        &self,
499        streams: Vec<String>,
500        correlation: u64,
501    ) -> BinanceWsResult<Vec<String>> {
502        self.unsubscribe_inner(streams, Some(correlation)).await
503    }
504
505    async fn unsubscribe_inner(
506        &self,
507        streams: Vec<String>,
508        correlation: Option<u64>,
509    ) -> BinanceWsResult<Vec<String>> {
510        let _connect_guard = self.connect_lock.lock().await;
511        let mut slots = self.slots.lock();
512
513        if self.signal.load(Ordering::Acquire) {
514            return Err(BinanceWsError::ClientError(
515                "Binance Futures stream pool is shutting down".to_string(),
516            ));
517        }
518        let mut slot_batches: Vec<(usize, Vec<String>)> = Vec::new();
519
520        for stream in &streams {
521            if let Some(slot_idx) = slots.iter().position(|s| s.streams.contains(stream)) {
522                if let Some(batch) = slot_batches.iter_mut().find(|(i, _)| *i == slot_idx) {
523                    batch.1.push(stream.clone());
524                } else {
525                    slot_batches.push((slot_idx, vec![stream.clone()]));
526                }
527            }
528        }
529
530        // Send commands first; only update slot state on success
531        for (slot_idx, batch) in &slot_batches {
532            slots[*slot_idx]
533                .cmd_tx
534                .send(BinanceFuturesWsStreamsCommand::Unsubscribe {
535                    streams: batch.clone(),
536                    correlation,
537                })
538                .map_err(|e| {
539                    BinanceWsError::ClientError(format!(
540                        "Handler not available for pool slot {slot_idx}: {e}"
541                    ))
542                })?;
543
544            for stream in batch {
545                slots[*slot_idx].streams.retain(|s| s != stream);
546            }
547        }
548
549        Ok(slot_batches
550            .into_iter()
551            .flat_map(|(_, batch)| batch)
552            .collect())
553    }
554
555    /// Returns a stream of messages from all WebSocket connections.
556    ///
557    /// This method can only be called once per connection lifecycle. Subsequent calls
558    /// return an empty stream.
559    pub fn stream(&self) -> impl Stream<Item = BinanceFuturesWsStreamsMessage> + 'static {
560        let out_rx = self.out_rx.lock().take();
561        async_stream::stream! {
562            if let Some(mut rx) = out_rx {
563                while let Some(msg) = rx.recv().await {
564                    yield msg;
565                }
566            }
567        }
568    }
569
570    /// Bulk initialize the instrument cache.
571    pub fn cache_instruments(&self, instruments: &[InstrumentAny]) {
572        self.instruments_cache.rcu(|m| {
573            for inst in instruments {
574                m.insert(inst.raw_symbol().inner(), inst.clone());
575            }
576        });
577    }
578
579    /// Replaces the complete instrument cache.
580    pub fn replace_instruments(&self, instruments: &[InstrumentAny]) {
581        let cache = instruments
582            .iter()
583            .map(|instrument| (instrument.raw_symbol().inner(), instrument.clone()))
584            .collect();
585        self.instruments_cache.store(cache);
586    }
587
588    /// Update a single instrument in the cache.
589    pub fn cache_instrument(&self, instrument: InstrumentAny) {
590        self.instruments_cache
591            .insert(instrument.raw_symbol().inner(), instrument);
592    }
593
594    /// Returns a shared reference to the instruments cache.
595    #[must_use]
596    pub fn instruments_cache(&self) -> Arc<AtomicMap<Ustr, InstrumentAny>> {
597        self.instruments_cache.clone()
598    }
599
600    /// Returns an instrument from the cache by raw symbol.
601    #[must_use]
602    pub fn get_instrument(&self, symbol: &str) -> Option<InstrumentAny> {
603        self.instruments_cache.get_cloned(&Ustr::from(symbol))
604    }
605
606    async fn create_connection(&self, slot_index: usize) -> BinanceWsResult<ConnectionSlot> {
607        let out_tx = self.out_tx.lock().clone().ok_or_else(|| {
608            BinanceWsError::ClientError("Output channel not initialized".to_string())
609        })?;
610
611        let (raw_handler, raw_rx) = channel_message_handler();
612        let ping_handler: PingHandler = Arc::new(move |_| {});
613
614        let mut headers = create_standard_nautilus_headers();
615
616        if let Some(ref cred) = self.credential {
617            headers.push((
618                BINANCE_API_KEY_HEADER.to_string(),
619                cred.api_key().to_string(),
620            ));
621        }
622
623        let config = WebSocketConfig {
624            url: self.url.expose_secret().to_owned(),
625            headers,
626            heartbeat_interval_secs: self.heartbeat,
627            heartbeat_payload: None,
628            connect_timeout_ms: Some(5_000),
629            reconnect_delay_initial_ms: Some(500),
630            reconnect_delay_max_ms: Some(5_000),
631            reconnect_backoff_factor: Some(2.0),
632            reconnect_jitter_ms: Some(250),
633            reconnect_max_attempts: None,
634            heartbeat_timeout_secs: None,
635            idle_timeout_ms: None,
636            writer_capacity: None,
637            backend: self.transport_backend,
638            proxy_url: self
639                .proxy_url
640                .as_ref()
641                .map(|value| value.expose_secret().to_owned()),
642            max_message_size_bytes: None,
643            max_frame_size_bytes: None,
644        };
645
646        let keyed_quotas = vec![(
647            BINANCE_RATE_LIMIT_KEY_SUBSCRIPTION[0].to_string(),
648            *BINANCE_WS_SUBSCRIPTION_QUOTA,
649        )];
650
651        let socket_control = self
652            .socket_factory
653            .as_ref()
654            .zip(self.socket_endpoint.as_ref())
655            .map(|(factory, endpoint)| {
656                let endpoint = if slot_index == 0 {
657                    endpoint.clone()
658                } else {
659                    format!("{endpoint}-{slot_index}")
660                };
661                factory.control(endpoint)
662            });
663        let client = WebSocketClient::builder()
664            .config(config)
665            .message_handler(raw_handler)
666            .ping_handler(ping_handler)
667            .keyed_quotas(keyed_quotas)
668            .default_quota(*BINANCE_WS_CONNECTION_QUOTA)
669            .maybe_state_sink(socket_control.as_ref().map(SocketControl::sink))
670            .connect()
671            .await
672            .map_err(|e| BinanceWsError::NetworkError(e.to_string()))?;
673
674        let connection_mode = client.connection_mode_atomic();
675        let reconnect_handle = client.reconnect_handle();
676        let subscriptions_state = SubscriptionState::new('@');
677        let cancellation_token = CancellationToken::new();
678
679        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
680
681        // Convert raw Message frames to Vec<u8> for the JSON handler
682        let (bytes_tx, bytes_rx) = tokio::sync::mpsc::unbounded_channel::<Vec<u8>>();
683
684        let mut bytes_task = TaskSlot::new();
685        if let Err(e) = bytes_task.spawn(async move {
686            let mut raw_rx = raw_rx;
687            while let Some(msg) = raw_rx.recv().await {
688                let data = match msg {
689                    Message::Binary(data) => data.to_vec(),
690                    Message::Text(text) => text.as_bytes().to_vec(),
691                    Message::Close(_) => break,
692                    Message::Ping(_) | Message::Pong(_) | Message::Frame(_) => continue,
693                };
694
695                if bytes_tx.send(data).is_err() {
696                    break;
697                }
698            }
699        }) {
700            let shutdown_error =
701                finish_slot_task(&mut bytes_task, "Binance futures WS bytes").await;
702            return Err(BinanceWsError::ClientError(match shutdown_error {
703                Some(shutdown_error) => format!(
704                    "Failed to start futures WS bytes task: {e}; startup rollback failed: \
705                     {shutdown_error}"
706                ),
707                None => format!("Failed to start futures WS bytes task: {e}"),
708            }));
709        }
710
711        let mut handler = BinanceFuturesDataWsFeedHandler::new(
712            self.signal.clone(),
713            cmd_rx,
714            bytes_rx,
715            out_tx.clone(),
716            subscriptions_state.clone(),
717            self.request_id_counter.clone(),
718        );
719
720        cmd_tx
721            .send(BinanceFuturesWsStreamsCommand::SetClient(client))
722            .map_err(|e| BinanceWsError::ClientError(format!("Failed to set client: {e}")))?;
723
724        let signal = self.signal.clone();
725        let token = cancellation_token.clone();
726        let subs = subscriptions_state.clone();
727        let resubscribe_tx = cmd_tx.clone();
728
729        let mut handler_task = TaskSlot::new();
730        if let Err(e) = handler_task.spawn(async move {
731            loop {
732                tokio::select! {
733                    () = token.cancelled() => {
734                        log::debug!("Handler task cancelled");
735                        break;
736                    }
737                    result = handler.next() => {
738                        match result {
739                            Some(BinanceFuturesWsStreamsMessage::Reconnected(abandoned)) => {
740                                log::info!("WebSocket reconnected, restoring subscriptions");
741                                let all_topics = subs.all_topics();
742                                for topic in &all_topics {
743                                    subs.mark_failure(topic);
744                                }
745
746                                let streams = subs.all_topics();
747                                if !streams.is_empty()
748                                    && let Err(e) = resubscribe_tx.send(BinanceFuturesWsStreamsCommand::Subscribe { streams }) {
749                                        log::error!("Failed to resubscribe after reconnect: {e}");
750                                    }
751
752                                if out_tx
753                                    .send(BinanceFuturesWsStreamsMessage::Reconnected(abandoned))
754                                    .is_err()
755                                {
756                                    log::debug!("Output channel closed");
757                                    break;
758                                }
759                            }
760                            Some(msg) => {
761                                if out_tx.send(msg).is_err() {
762                                    log::debug!("Output channel closed");
763                                    break;
764                                }
765                            }
766                            None => {
767                                if signal.load(Ordering::Relaxed) {
768                                    log::debug!("Handler received shutdown signal");
769                                } else {
770                                    log::warn!("Handler loop ended unexpectedly");
771                                }
772                                break;
773                            }
774                        }
775                    }
776                }
777            }
778        }) {
779            cancellation_token.cancel();
780            bytes_task.abort();
781            let mut shutdown_errors = Vec::new();
782
783            if let Some(error) =
784                finish_slot_task(&mut handler_task, "Binance futures WS handler").await
785            {
786                shutdown_errors.push(error);
787            }
788
789            if let Some(error) =
790                finish_slot_task(&mut bytes_task, "Binance futures WS bytes").await
791            {
792                shutdown_errors.push(error);
793            }
794            return Err(BinanceWsError::ClientError(if shutdown_errors.is_empty() {
795                format!("Failed to start futures WS handler task: {e}")
796            } else {
797                format!(
798                    "Failed to start futures WS handler task: {e}; startup rollback failed: {}",
799                    shutdown_errors.join("; ")
800                )
801            }));
802        }
803
804        if let Some(control) = &socket_control {
805            control.register(move || reconnect_handle.request_reconnect());
806        }
807
808        Ok(ConnectionSlot {
809            cmd_tx,
810            streams: Vec::new(),
811            subscriptions_state,
812            handler_task,
813            bytes_task,
814            cancellation_token,
815            connection_mode,
816            socket_control,
817            shutdown_errors: Vec::new(),
818        })
819    }
820}
821
822struct ConnectionSlots(Mutex<Vec<ConnectionSlot>>);
823
824impl std::ops::Deref for ConnectionSlots {
825    type Target = Mutex<Vec<ConnectionSlot>>;
826
827    fn deref(&self) -> &Self::Target {
828        &self.0
829    }
830}
831
832impl Drop for ConnectionSlots {
833    fn drop(&mut self) {
834        for slot in self.0.get_mut().iter() {
835            slot.cancellation_token.cancel();
836            if let Some(handle) = slot.handler_task.as_ref() {
837                handle.abort();
838            }
839
840            if let Some(handle) = slot.bytes_task.as_ref() {
841                handle.abort();
842            }
843
844            if let Some(control) = &slot.socket_control {
845                control.deregister();
846            }
847        }
848    }
849}
850
851struct ConnectionSlotBatch<'a> {
852    owner: &'a Mutex<Vec<ConnectionSlot>>,
853    slots: Vec<ConnectionSlot>,
854}
855
856impl<'a> ConnectionSlotBatch<'a> {
857    fn take(owner: &'a Mutex<Vec<ConnectionSlot>>) -> Self {
858        let slots = std::mem::take(&mut *owner.lock());
859        Self { owner, slots }
860    }
861}
862
863impl Drop for ConnectionSlotBatch<'_> {
864    fn drop(&mut self) {
865        self.owner.lock().extend(self.slots.drain(..));
866    }
867}
868
869async fn finish_slot_task(slot: &mut TaskSlot<()>, owner: &str) -> Option<String> {
870    let outcome = finish_task(
871        slot,
872        std::time::Duration::from_secs(2),
873        std::time::Duration::from_secs(2),
874    )
875    .await?;
876
877    match outcome {
878        TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted => None,
879        TaskJoinOutcome::Failed(e) => Some(format!("{owner} task failed: {e}")),
880        TaskJoinOutcome::Incomplete => Some(format!("{owner} task did not stop after abort")),
881    }
882}
883
884#[cfg(test)]
885mod tests {
886    use rstest::rstest;
887
888    use super::*;
889
890    #[tokio::test]
891    async fn test_cancelled_close_retains_connection_slot() {
892        let mut client = BinanceFuturesWebSocketClient::new(
893            BinanceProductType::UsdM,
894            BinanceEnvironment::Testnet,
895            None,
896            None,
897            None,
898            None,
899            TransportBackend::default(),
900        )
901        .unwrap();
902        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
903        client.slots.lock().push(ConnectionSlot {
904            cmd_tx,
905            streams: Vec::new(),
906            subscriptions_state: SubscriptionState::new('@'),
907            handler_task: TaskSlot::from_handle(tokio::spawn(std::future::pending())),
908            bytes_task: TaskSlot::new(),
909            cancellation_token: CancellationToken::new(),
910            connection_mode: Arc::new(AtomicU8::new(ConnectionMode::Active as u8)),
911            socket_control: None,
912            shutdown_errors: Vec::new(),
913        });
914
915        {
916            let close = client.close();
917            tokio::pin!(close);
918            tokio::select! {
919                result = &mut close => panic!("close completed unexpectedly: {result:?}"),
920                command = cmd_rx.recv() => assert!(command.is_some()),
921            }
922        }
923
924        let slots = client.slots.lock();
925        assert_eq!(slots.len(), 1);
926        assert!(slots[0].handler_task.is_some());
927    }
928
929    #[rstest]
930    fn test_with_proxy_preserves_proxy_url() {
931        let client = BinanceFuturesWebSocketClient::new(
932            BinanceProductType::UsdM,
933            BinanceEnvironment::Testnet,
934            None,
935            None,
936            None,
937            None,
938            TransportBackend::default(),
939        )
940        .unwrap()
941        .with_proxy(Some("socks5://proxy.example:1080".to_string()));
942
943        assert_eq!(
944            client.proxy_url.as_ref().map(SecretString::expose_secret),
945            Some("socks5://proxy.example:1080")
946        );
947    }
948}
nautilus_binance/futures/websocket/streams/client.rs

nautilus_binance/futures/websocket/streams/
client.rs
