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

nautilus_binance/futures/websocket/streams/
client.rs
