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

nautilus_binance/spot/websocket/public_json/
client.rs
