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

nautilus_binance/spot/websocket/public_json/
client.rs
