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

nautilus_binance/spot/websocket/public_json/
client.rs
