nautilus_dydx/websocket/
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//! WebSocket client for dYdX v4 API.
17//!
18//! This client provides streaming connectivity to dYdX's WebSocket API for both
19//! public market data and private account updates.
20//!
21//! # Authentication
22//!
23//! dYdX v4 uses Cosmos SDK wallet-based authentication. Unlike traditional exchanges:
24//! - **Public channels** require no authentication.
25//! - **Private channels** (subaccounts) only require the wallet address in the subscription message.
26//! - No signature or API key is needed for WebSocket connections themselves.
27//!
28//! # Connection pool
29//!
30//! The Indexer caps each WebSocket connection at 32 subscriptions per channel
31//! (`v4_trades`, `v4_candles`, `v4_orderbook`, `v4_markets`). To scale past that
32//! limit the client maintains a small pool of connection slots and routes each
33//! new subscription to the first slot with capacity, lazily spawning additional
34//! connections up to `max_ws_connections`. The shape mirrors
35//! `BinanceFuturesWebSocketClient` (including its `connect_lock` race fix),
36//! adapted so capacity is tracked per channel kind rather than as a single flat
37//! stream count.
38//!
39//! # References
40//!
41//! <https://docs.dydx.trade/developers/indexer/websockets>
42
43/// Pre-interned rate limit key for subscription operations (subscribe/unsubscribe).
44///
45/// dYdX allows up to 2 subscription messages per second per connection.
46/// See: <https://docs.dydx.trade/developers/indexer/websockets#rate-limits>
47pub static DYDX_RATE_LIMIT_KEY_SUBSCRIPTION: LazyLock<[Ustr; 1]> =
48    LazyLock::new(|| [Ustr::from("subscription")]);
49
50/// WebSocket topic delimiter for dYdX (channel:symbol format).
51pub const DYDX_WS_TOPIC_DELIMITER: char = ':';
52
53/// Default WebSocket quota for dYdX subscriptions (2 messages per second).
54pub static DYDX_WS_SUBSCRIPTION_QUOTA: LazyLock<Quota> = LazyLock::new(|| {
55    Quota::per_second(NonZeroU32::new(2).expect("non-zero")).expect("valid constant")
56});
57
58/// Default maximum number of WebSocket connections in the Indexer pool.
59pub const DEFAULT_MAX_WS_CONNECTIONS: usize = 8;
60
61/// Default per-connection subscription limit for sharded channels.
62pub const DEFAULT_PER_CHANNEL_SUBSCRIPTION_LIMIT: usize = 32;
63
64use std::{
65    fmt::Debug,
66    num::NonZeroU32,
67    sync::{
68        Arc, LazyLock,
69        atomic::{AtomicBool, AtomicU8, Ordering},
70    },
71    time::Duration,
72};
73
74use ahash::{AHashMap, AHashSet};
75use arc_swap::ArcSwap;
76use dashmap::DashMap;
77use nautilus_core::string::secret::SecretString;
78use nautilus_live::{
79    SocketControl, SocketControlFactory,
80    task::{TaskJoinOutcome, TaskSlot, finish_task},
81};
82use nautilus_model::{
83    data::BarType,
84    identifiers::{AccountId, InstrumentId},
85    instruments::InstrumentAny,
86};
87use nautilus_network::{
88    http::create_standard_nautilus_headers,
89    mode::ConnectionMode,
90    ratelimiter::quota::Quota,
91    websocket::{
92        AuthTracker, SubscriptionState, TransportBackend, WebSocketClient, WebSocketConfig,
93        channel_message_handler,
94    },
95};
96use parking_lot::Mutex;
97use ustr::Ustr;
98
99use super::{
100    dispatch::DydxWsDispatchState,
101    enums::{DydxWsChannel, DydxWsOperation, DydxWsOutputMessage},
102    error::{DydxWsError, DydxWsResult},
103    handler::{FeedHandler, HandlerCommand},
104    messages::DydxSubscription,
105};
106use crate::{
107    common::{credential::DydxCredential, instrument_cache::InstrumentCache},
108    execution::encoder::ClientOrderIdEncoder,
109};
110
111/// Identifies a dYdX channel for per-channel capacity accounting in the pool.
112#[derive(Copy, Clone, Debug)]
113#[repr(u8)]
114enum ChannelKind {
115    Trades = 0,
116    Candles = 1,
117    Orderbook = 2,
118    Markets = 3,
119}
120
121const CHANNEL_KIND_COUNT: usize = 4;
122
123/// Per-connection state inside the pool.
124#[derive(Debug)]
125struct ConnectionSlot {
126    cmd_tx: tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
127    topics: AHashMap<String, u32>,
128    channel_counts: [u16; CHANNEL_KIND_COUNT],
129    subscriptions_state: SubscriptionState,
130    handler_task: TaskSlot<()>,
131    connection_mode: Arc<AtomicU8>,
132    socket_control: Option<SocketControl>,
133}
134
135/// WebSocket client for dYdX v4 market data and account streams.
136///
137/// # Authentication
138///
139/// dYdX v4 does not require traditional API key signatures for WebSocket connections.
140/// Public channels work without any credentials. Private channels (subaccounts) only
141/// need the wallet address included in the subscription message.
142///
143/// The [`DydxCredential`] stored in this client is used for:
144/// - Providing the wallet address for private channel subscriptions
145/// - Transaction signing (when placing orders via the validator node)
146///
147/// It is **NOT** used for WebSocket message signing or authentication.
148///
149/// # Architecture
150///
151/// The client owns a small pool of connection slots. Each slot has its own
152/// `WebSocketClient`, [`FeedHandler`] task, command channel, and
153/// [`SubscriptionState`]. All slots write parsed events into a single shared
154/// output channel so callers see one merged stream.
155#[derive(Debug)]
156pub struct DydxWebSocketClient {
157    url: String,
158    credential: Option<Arc<DydxCredential>>,
159    requires_auth: bool,
160    auth_tracker: AuthTracker,
161    slots: Arc<ConnectionSlots>,
162    admission: Arc<Mutex<PoolAdmission>>,
163    connect_lock: Arc<tokio::sync::Mutex<()>>,
164    connection_mode: Arc<ArcSwap<AtomicU8>>,
165    signal: Arc<AtomicBool>,
166    instrument_cache: Arc<InstrumentCache>,
167    account_id: Option<AccountId>,
168    heartbeat: Option<u64>,
169    out_tx: Arc<Mutex<Option<tokio::sync::mpsc::UnboundedSender<DydxWsOutputMessage>>>>,
170    out_rx: Arc<Mutex<Option<tokio::sync::mpsc::UnboundedReceiver<DydxWsOutputMessage>>>>,
171    encoder: Arc<ClientOrderIdEncoder>,
172    bar_types: Arc<DashMap<String, BarType>>,
173    bars_timestamp_on_close: Arc<AtomicBool>,
174    ws_dispatch_state: Arc<DydxWsDispatchState>,
175    transport_backend: TransportBackend,
176    proxy_url: Option<SecretString>,
177    max_ws_connections: usize,
178    per_channel_limit: usize,
179    socket_factory: Option<SocketControlFactory>,
180}
181
182impl Clone for DydxWebSocketClient {
183    fn clone(&self) -> Self {
184        Self {
185            url: self.url.clone(),
186            credential: self.credential.clone(),
187            requires_auth: self.requires_auth,
188            auth_tracker: self.auth_tracker.clone(),
189            slots: self.slots.clone(),
190            admission: self.admission.clone(),
191            connect_lock: self.connect_lock.clone(),
192            connection_mode: self.connection_mode.clone(),
193            signal: self.signal.clone(),
194            instrument_cache: self.instrument_cache.clone(),
195            account_id: self.account_id,
196            heartbeat: self.heartbeat,
197            out_tx: self.out_tx.clone(),
198            out_rx: self.out_rx.clone(),
199            encoder: self.encoder.clone(),
200            bar_types: self.bar_types.clone(),
201            bars_timestamp_on_close: self.bars_timestamp_on_close.clone(),
202            ws_dispatch_state: self.ws_dispatch_state.clone(),
203            transport_backend: self.transport_backend,
204            proxy_url: self.proxy_url.clone(),
205            max_ws_connections: self.max_ws_connections,
206            per_channel_limit: self.per_channel_limit,
207            socket_factory: self.socket_factory.clone(),
208        }
209    }
210}
211
212impl DydxWebSocketClient {
213    /// Creates a new public WebSocket client for market data.
214    ///
215    /// This creates a new independent instrument cache. To share a cache with
216    /// the HTTP client, use [`Self::new_public_with_cache`] instead.
217    #[must_use]
218    pub fn new_public(url: String, heartbeat: Option<u64>, proxy_url: Option<String>) -> Self {
219        Self::new_public_with_cache(
220            url,
221            Arc::new(InstrumentCache::new()),
222            heartbeat,
223            TransportBackend::default(),
224            proxy_url,
225        )
226    }
227
228    /// Creates a new public WebSocket client with a shared instrument cache.
229    ///
230    /// Use this when you want to share instrument data with the HTTP client.
231    #[must_use]
232    pub fn new_public_with_cache(
233        url: String,
234        instrument_cache: Arc<InstrumentCache>,
235        heartbeat: Option<u64>,
236        transport_backend: TransportBackend,
237        proxy_url: Option<String>,
238    ) -> Self {
239        Self::new_public_with_cache_and_pool(
240            url,
241            instrument_cache,
242            heartbeat,
243            transport_backend,
244            proxy_url,
245            DEFAULT_MAX_WS_CONNECTIONS,
246            DEFAULT_PER_CHANNEL_SUBSCRIPTION_LIMIT,
247        )
248    }
249
250    /// Creates a new public WebSocket client with full pool configuration.
251    #[must_use]
252    pub fn new_public_with_cache_and_pool(
253        url: String,
254        instrument_cache: Arc<InstrumentCache>,
255        heartbeat: Option<u64>,
256        transport_backend: TransportBackend,
257        proxy_url: Option<String>,
258        max_ws_connections: usize,
259        per_channel_limit: usize,
260    ) -> Self {
261        Self::new_inner(
262            url,
263            None,
264            false,
265            instrument_cache,
266            None,
267            heartbeat,
268            transport_backend,
269            proxy_url,
270            max_ws_connections,
271            per_channel_limit,
272        )
273    }
274
275    /// Creates a new private WebSocket client for account updates.
276    ///
277    /// This creates a new independent instrument cache. To share a cache with
278    /// the HTTP client, use [`Self::new_private_with_cache`] instead.
279    #[must_use]
280    pub fn new_private(
281        url: String,
282        credential: DydxCredential,
283        account_id: AccountId,
284        heartbeat: Option<u64>,
285        proxy_url: Option<String>,
286    ) -> Self {
287        Self::new_private_with_cache(
288            url,
289            credential,
290            account_id,
291            Arc::new(InstrumentCache::new()),
292            heartbeat,
293            TransportBackend::default(),
294            proxy_url,
295        )
296    }
297
298    /// Creates a new private WebSocket client with a shared instrument cache.
299    ///
300    /// Use this when you want to share instrument data with the HTTP client.
301    #[must_use]
302    pub fn new_private_with_cache(
303        url: String,
304        credential: DydxCredential,
305        account_id: AccountId,
306        instrument_cache: Arc<InstrumentCache>,
307        heartbeat: Option<u64>,
308        transport_backend: TransportBackend,
309        proxy_url: Option<String>,
310    ) -> Self {
311        Self::new_inner(
312            url,
313            Some(Arc::new(credential)),
314            true,
315            instrument_cache,
316            Some(account_id),
317            heartbeat,
318            transport_backend,
319            proxy_url,
320            DEFAULT_MAX_WS_CONNECTIONS,
321            DEFAULT_PER_CHANNEL_SUBSCRIPTION_LIMIT,
322        )
323    }
324
325    #[allow(clippy::too_many_arguments)]
326    fn new_inner(
327        url: String,
328        credential: Option<Arc<DydxCredential>>,
329        requires_auth: bool,
330        instrument_cache: Arc<InstrumentCache>,
331        account_id: Option<AccountId>,
332        heartbeat: Option<u64>,
333        transport_backend: TransportBackend,
334        proxy_url: Option<String>,
335        max_ws_connections: usize,
336        per_channel_limit: usize,
337    ) -> Self {
338        Self {
339            url,
340            credential,
341            requires_auth,
342            auth_tracker: AuthTracker::new(),
343            slots: Arc::new(ConnectionSlots::new()),
344            admission: Arc::new(Mutex::new(PoolAdmission {
345                generation: 0,
346                closed: false,
347            })),
348            connect_lock: Arc::new(tokio::sync::Mutex::new(())),
349            connection_mode: Arc::new(ArcSwap::from_pointee(AtomicU8::new(
350                ConnectionMode::Closed as u8,
351            ))),
352            signal: Arc::new(AtomicBool::new(false)),
353            instrument_cache,
354            account_id,
355            heartbeat,
356            out_tx: Arc::new(Mutex::new(None)),
357            out_rx: Arc::new(Mutex::new(None)),
358            encoder: Arc::new(ClientOrderIdEncoder::new()),
359            bar_types: Arc::new(DashMap::new()),
360            bars_timestamp_on_close: Arc::new(AtomicBool::new(true)),
361            ws_dispatch_state: Arc::new(DydxWsDispatchState::default()),
362            transport_backend,
363            proxy_url: proxy_url.map(SecretString::from),
364            max_ws_connections: max_ws_connections.max(1),
365            per_channel_limit: per_channel_limit.max(1),
366            socket_factory: None,
367        }
368    }
369
370    pub(crate) fn begin_shutdown(&self) {
371        let mut admission = self.admission.lock();
372        admission.generation = admission.generation.wrapping_add(1);
373        admission.closed = true;
374        self.signal.store(true, Ordering::Release);
375
376        for slot in self.slots.lock().iter() {
377            let _ = slot.cmd_tx.send(HandlerCommand::Disconnect);
378        }
379    }
380
381    fn open_generation(&self) -> u64 {
382        let mut admission = self.admission.lock();
383        admission.generation = admission.generation.wrapping_add(1);
384        admission.closed = false;
385        admission.generation
386    }
387
388    fn admission_generation(&self) -> DydxWsResult<u64> {
389        let admission = self.admission.lock();
390        if admission.closed {
391            Err(DydxWsError::Transport(
392                "WebSocket connection pool is closed".to_string(),
393            ))
394        } else {
395            Ok(admission.generation)
396        }
397    }
398
399    /// Configures socket state reporting and reconnect control for each pool slot.
400    #[must_use]
401    pub fn with_socket_factory(mut self, factory: SocketControlFactory) -> Self {
402        self.socket_factory = Some(factory);
403        self
404    }
405
406    /// Returns the credential associated with this client, if any.
407    #[must_use]
408    pub fn credential(&self) -> Option<&Arc<DydxCredential>> {
409        self.credential.as_ref()
410    }
411
412    /// Returns `true` when any connection in the pool is connected.
413    #[must_use]
414    pub fn is_connected(&self) -> bool {
415        let slots = self.slots.lock();
416        slots.iter().any(|s| {
417            let mode = s.connection_mode.load(Ordering::Relaxed);
418            mode == ConnectionMode::Active as u8 || mode == ConnectionMode::Reconnect as u8
419        })
420    }
421
422    /// Returns the URL of this WebSocket client.
423    #[must_use]
424    pub fn url(&self) -> &str {
425        &self.url
426    }
427
428    /// Returns a clone of the connection mode atomic reference.
429    ///
430    /// With sharding, the returned atomic tracks the **primary** slot (slot 0)
431    /// only; use [`Self::is_connected`] for a pool-wide check.
432    #[must_use]
433    pub fn connection_mode_atomic(&self) -> Arc<ArcSwap<AtomicU8>> {
434        self.connection_mode.clone()
435    }
436
437    /// Returns the current number of active slots in the pool.
438    #[must_use]
439    pub fn pool_size(&self) -> usize {
440        self.slots.lock().len()
441    }
442
443    /// Returns the configured maximum number of pool connections.
444    #[must_use]
445    pub const fn max_ws_connections(&self) -> usize {
446        self.max_ws_connections
447    }
448
449    /// Returns the configured per-channel subscription limit.
450    #[must_use]
451    pub const fn per_channel_limit(&self) -> usize {
452        self.per_channel_limit
453    }
454
455    /// Sets the account ID for account message parsing.
456    pub fn set_account_id(&mut self, account_id: AccountId) {
457        self.account_id = Some(account_id);
458    }
459
460    /// Returns the account ID if set.
461    #[must_use]
462    pub fn account_id(&self) -> Option<AccountId> {
463        self.account_id
464    }
465
466    /// Replaces the instrument cache with an externally shared one.
467    ///
468    /// Use this to share the HTTP client's cache (which includes CLOB pair ID
469    /// and market ticker indices) with the WebSocket client. Must be called
470    /// before `connect()`.
471    pub fn set_instrument_cache(&mut self, cache: Arc<InstrumentCache>) {
472        self.instrument_cache = cache;
473    }
474
475    /// Caches a single instrument.
476    ///
477    /// Any existing instrument with the same ID will be replaced.
478    pub fn cache_instrument(&self, instrument: InstrumentAny) {
479        self.instrument_cache.insert_instrument_only(instrument);
480    }
481
482    /// Caches multiple instruments.
483    ///
484    /// Any existing instruments with the same IDs will be replaced.
485    pub fn cache_instruments(&self, instruments: Vec<InstrumentAny>) {
486        log::debug!(
487            "Caching {} instruments in WebSocket client",
488            instruments.len()
489        );
490        self.instrument_cache.insert_instruments_only(instruments);
491    }
492
493    /// Returns a reference to the shared instrument cache.
494    #[must_use]
495    pub fn instrument_cache(&self) -> &Arc<InstrumentCache> {
496        &self.instrument_cache
497    }
498
499    /// Returns a reference to the shared client order ID encoder.
500    #[must_use]
501    pub fn encoder(&self) -> &Arc<ClientOrderIdEncoder> {
502        &self.encoder
503    }
504
505    /// Returns a reference to the bar type registrations map.
506    #[must_use]
507    pub fn bar_types(&self) -> &Arc<DashMap<String, BarType>> {
508        &self.bar_types
509    }
510
511    /// Returns a reference to the shared WebSocket dispatch state.
512    pub fn ws_dispatch_state(&self) -> &Arc<DydxWsDispatchState> {
513        &self.ws_dispatch_state
514    }
515
516    /// Sets whether bar timestamps use the close time.
517    pub fn set_bars_timestamp_on_close(&self, value: bool) {
518        self.bars_timestamp_on_close.store(value, Ordering::Relaxed);
519    }
520
521    /// Returns whether bar timestamps use the close time.
522    #[must_use]
523    pub fn bars_timestamp_on_close(&self) -> bool {
524        self.bars_timestamp_on_close.load(Ordering::Relaxed)
525    }
526
527    /// Returns all cached instruments.
528    ///
529    /// This is a snapshot of the current cache contents.
530    #[must_use]
531    pub fn all_instruments(&self) -> Vec<InstrumentAny> {
532        self.instrument_cache.all_instruments()
533    }
534
535    /// Returns the number of cached instruments.
536    #[must_use]
537    pub fn cached_instruments_count(&self) -> usize {
538        self.instrument_cache.len()
539    }
540
541    /// Retrieves an instrument from the cache by InstrumentId.
542    ///
543    /// Returns `None` if the instrument is not found.
544    #[must_use]
545    pub fn get_instrument(&self, instrument_id: &InstrumentId) -> Option<InstrumentAny> {
546        self.instrument_cache.get(instrument_id)
547    }
548
549    /// Retrieves an instrument from the cache by market ticker (e.g., "BTC-USD").
550    ///
551    /// Returns `None` if the instrument is not found.
552    #[must_use]
553    pub fn get_instrument_by_market(&self, ticker: &str) -> Option<InstrumentAny> {
554        self.instrument_cache.get_by_market(ticker)
555    }
556
557    /// Takes ownership of the inbound message receiver.
558    /// Returns None if the receiver has already been taken or not connected.
559    pub fn take_receiver(
560        &mut self,
561    ) -> Option<tokio::sync::mpsc::UnboundedReceiver<DydxWsOutputMessage>> {
562        self.out_rx.lock().take()
563    }
564
565    /// Returns a stream of venue-specific WebSocket messages.
566    ///
567    /// Takes ownership of the message receiver and returns it as a `Stream`.
568    ///
569    /// # Panics
570    ///
571    /// Panics if the message receiver has already been taken or the client is not connected.
572    pub fn stream(
573        &mut self,
574    ) -> impl futures_util::Stream<Item = DydxWsOutputMessage> + Send + 'static {
575        let mut rx = self
576            .out_rx
577            .lock()
578            .take()
579            .expect("Message stream receiver already taken or not connected");
580
581        async_stream::stream! {
582            while let Some(msg) = rx.recv().await {
583                yield msg;
584            }
585        }
586    }
587
588    /// Connects the websocket client and opens the primary pool slot.
589    ///
590    /// Additional slots are spawned lazily by `subscribe_*` methods once the
591    /// per-channel limit is reached on every existing slot.
592    ///
593    /// # Errors
594    ///
595    /// Returns an error if the connection cannot be established.
596    pub async fn connect(&mut self) -> DydxWsResult<()> {
597        let connect_lock = Arc::clone(&self.connect_lock);
598        let _connect_guard = connect_lock.lock().await;
599
600        let already_connected = {
601            let admission = self.admission.lock();
602            !admission.closed && self.is_connected()
603        };
604
605        if already_connected {
606            return Ok(());
607        }
608
609        if !self.slots.lock().is_empty() {
610            self.disconnect_connections().await?;
611        }
612
613        let generation = self.open_generation();
614        self.signal.store(false, Ordering::Release);
615
616        let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<DydxWsOutputMessage>();
617        {
618            let mut guard = self.out_tx.lock();
619            *guard = Some(out_tx);
620        }
621        {
622            let mut guard = self.out_rx.lock();
623            *guard = Some(out_rx);
624        }
625
626        let slot = match self.create_connection(0).await {
627            Ok(slot) => slot,
628            Err(e) => {
629                self.begin_shutdown();
630                *self.out_tx.lock() = None;
631                *self.out_rx.lock() = None;
632                return Err(e);
633            }
634        };
635        let admission = self.admission.lock();
636        let mut slots = self.slots.lock();
637
638        if admission.closed || admission.generation != generation {
639            let _ = slot.cmd_tx.send(HandlerCommand::Disconnect);
640            slots.push(slot);
641            return Err(DydxWsError::Transport(
642                "WebSocket connection was canceled by shutdown".to_string(),
643            ));
644        }
645        self.connection_mode.store(slot.connection_mode.clone());
646        slots.push(slot);
647        drop(slots);
648        drop(admission);
649
650        log::debug!("Connected dYdX WebSocket pool: {}", self.url);
651        Ok(())
652    }
653
654    /// Disconnects all websocket connections in the pool.
655    ///
656    /// # Errors
657    ///
658    /// Returns an error if the underlying clients cannot be accessed.
659    pub async fn disconnect(&mut self) -> DydxWsResult<()> {
660        self.begin_shutdown();
661        let connect_lock = Arc::clone(&self.connect_lock);
662        let _connect_guard = connect_lock.lock().await;
663        self.disconnect_connections().await
664    }
665
666    async fn disconnect_connections(&self) -> DydxWsResult<()> {
667        self.begin_shutdown();
668
669        let mut slots = ConnectionSlotBatch::take(&self.slots);
670
671        for slot in &mut slots.slots {
672            if let Some(control) = &slot.socket_control {
673                control.deregister();
674            }
675            let _ = slot.cmd_tx.send(HandlerCommand::Disconnect);
676
677            if let Some(outcome) = finish_task(
678                &mut slot.handler_task,
679                Duration::from_secs(2),
680                Duration::from_secs(2),
681            )
682            .await
683            {
684                match outcome {
685                    TaskJoinOutcome::Completed(()) => log::debug!("Handler task completed"),
686                    TaskJoinOutcome::Aborted => {}
687                    TaskJoinOutcome::Failed(error) => {
688                        self.slots
689                            .push_shutdown_error(format!("handler task failed: {error}"));
690                    }
691                    TaskJoinOutcome::Incomplete => {
692                        self.slots.push_shutdown_error(
693                            "handler task did not stop after abort".to_string(),
694                        );
695                    }
696                }
697            }
698        }
699
700        slots.slots.retain(|slot| slot.handler_task.is_some());
701        let has_incomplete_tasks = !slots.slots.is_empty();
702
703        if !has_incomplete_tasks {
704            self.connection_mode
705                .store(Arc::new(AtomicU8::new(ConnectionMode::Closed as u8)));
706            *self.out_tx.lock() = None;
707            *self.out_rx.lock() = None;
708        }
709
710        let join_errors = self.slots.take_shutdown_errors();
711        if !join_errors.is_empty() {
712            return Err(DydxWsError::Transport(join_errors.join("; ")));
713        }
714
715        self.connection_mode
716            .store(Arc::new(AtomicU8::new(ConnectionMode::Closed as u8)));
717
718        log::debug!("Disconnected dYdX WebSocket pool");
719        Ok(())
720    }
721
722    /// Sends a command directly to the primary slot (slot 0).
723    ///
724    /// # Errors
725    ///
726    /// Returns an error if no slot exists or the handler task has terminated.
727    pub fn send_command(&self, cmd: HandlerCommand) -> DydxWsResult<()> {
728        let admission = self.admission.lock();
729        if admission.closed {
730            return Err(DydxWsError::Transport(
731                "WebSocket connection pool is closed".to_string(),
732            ));
733        }
734        let slots = self.slots.lock();
735        let slot = slots
736            .first()
737            .ok_or_else(|| DydxWsError::Transport("No pool slots available".to_string()))?;
738        slot.cmd_tx.send(cmd).map_err(|e| {
739            DydxWsError::Transport(format!("Failed to send command to slot 0: {e}"))
740        })?;
741        Ok(())
742    }
743
744    async fn create_connection(&self, slot_index: usize) -> DydxWsResult<ConnectionSlot> {
745        let (message_handler, raw_rx) = channel_message_handler();
746        let headers = create_standard_nautilus_headers();
747
748        let cfg = WebSocketConfig {
749            url: self.url.clone(),
750            headers,
751            heartbeat_interval_secs: self.heartbeat,
752            heartbeat_payload: None,
753            connect_timeout_ms: Some(15_000),
754            reconnect_delay_initial_ms: Some(250),
755            reconnect_delay_max_ms: Some(5_000),
756            reconnect_backoff_factor: Some(2.0),
757            reconnect_jitter_ms: Some(200),
758            reconnect_max_attempts: None,
759            heartbeat_timeout_secs: None,
760            idle_timeout_ms: None,
761            backend: self.transport_backend,
762            proxy_url: self
763                .proxy_url
764                .as_ref()
765                .map(|value| value.expose_secret().to_owned()),
766        };
767
768        let socket_control = self.socket_factory.as_ref().map(|factory| {
769            let kind = if self.requires_auth { "user" } else { "data" };
770            let endpoint = format!("dydx-{kind}-streams");
771            if slot_index == 0 {
772                factory.control(endpoint)
773            } else {
774                factory.control(format!("{endpoint}-{slot_index}"))
775            }
776        });
777        let client = WebSocketClient::builder()
778            .config(cfg)
779            .message_handler(message_handler)
780            .default_quota(*DYDX_WS_SUBSCRIPTION_QUOTA)
781            .maybe_state_sink(
782                socket_control
783                    .as_ref()
784                    .map(nautilus_live::SocketControl::sink),
785            )
786            .connect()
787            .await
788            .map_err(|e| DydxWsError::Transport(e.to_string()))?;
789
790        let connection_mode = client.connection_mode_atomic();
791        let reconnect_handle = client.reconnect_handle();
792        let subscriptions_state = SubscriptionState::new(DYDX_WS_TOPIC_DELIMITER);
793
794        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
795
796        let out_tx =
797            self.out_tx.lock().clone().ok_or_else(|| {
798                DydxWsError::Transport("Output channel not initialized".to_string())
799            })?;
800
801        let signal = self.signal.clone();
802        let subscriptions = subscriptions_state.clone();
803
804        let mut handler_task = TaskSlot::new();
805        if let Err(e) = handler_task.spawn(async move {
806            let mut handler =
807                FeedHandler::new(cmd_rx, out_tx, raw_rx, client, signal, subscriptions);
808            handler.run().await;
809        }) {
810            let shutdown_error = match finish_task(
811                &mut handler_task,
812                std::time::Duration::ZERO,
813                std::time::Duration::from_secs(2),
814            )
815            .await
816            {
817                Some(TaskJoinOutcome::Failed(error)) => {
818                    Some(format!("handler task failed: {error}"))
819                }
820                Some(TaskJoinOutcome::Incomplete) => {
821                    Some("handler task did not stop after abort".to_string())
822                }
823                None | Some(TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted) => None,
824            };
825            return Err(DydxWsError::Transport(match shutdown_error {
826                Some(shutdown_error) => format!(
827                    "Failed to start handler task: {e}; startup rollback failed: {shutdown_error}"
828                ),
829                None => format!("Failed to start handler task: {e}"),
830            }));
831        }
832
833        if let Some(control) = &socket_control {
834            control.register(move || reconnect_handle.request_reconnect());
835        }
836
837        Ok(ConnectionSlot {
838            cmd_tx,
839            topics: AHashMap::new(),
840            channel_counts: [0; CHANNEL_KIND_COUNT],
841            subscriptions_state,
842            handler_task,
843            connection_mode,
844            socket_control,
845        })
846    }
847
848    fn ticker_from_instrument_id(instrument_id: &InstrumentId) -> String {
849        let mut s = instrument_id.symbol.as_str().to_string();
850        if let Some(stripped) = s.strip_suffix("-PERP") {
851            s = stripped.to_string();
852        }
853        s
854    }
855
856    fn topic(channel: DydxWsChannel, id: Option<&str>) -> String {
857        match id {
858            Some(id) => format!("{}{}{}", channel.as_ref(), DYDX_WS_TOPIC_DELIMITER, id),
859            None => channel.as_ref().to_string(),
860        }
861    }
862
863    async fn subscribe_topic(
864        &self,
865        channel: ChannelKind,
866        topic: String,
867        sub_msg: DydxSubscription,
868    ) -> DydxWsResult<()> {
869        let _connect_guard = self.connect_lock.lock().await;
870        let generation = self.admission_generation()?;
871
872        {
873            let admission = self.admission.lock();
874            if admission.closed || admission.generation != generation {
875                return Err(DydxWsError::Transport(
876                    "WebSocket connection pool is closed".to_string(),
877                ));
878            }
879            let mut slots = self.slots.lock();
880            if let Some(slot) = slots.iter_mut().find(|s| s.topics.contains_key(&topic)) {
881                *slot.topics.get_mut(&topic).expect("topic refcount present") += 1;
882                return Ok(());
883            }
884        }
885
886        let target_idx = loop {
887            {
888                let admission = self.admission.lock();
889                if admission.closed || admission.generation != generation {
890                    return Err(DydxWsError::Transport(
891                        "WebSocket connection pool is closed".to_string(),
892                    ));
893                }
894                let slots = self.slots.lock();
895                if let Some(idx) = slots.iter().position(|s| {
896                    (s.channel_counts[channel as usize] as usize) < self.per_channel_limit
897                }) {
898                    break idx;
899                }
900
901                if slots.len() >= self.max_ws_connections {
902                    return Err(DydxWsError::Subscription(format!(
903                        "Pool exhausted: {} connections x {} {:?} subscriptions",
904                        self.max_ws_connections, self.per_channel_limit, channel,
905                    )));
906                }
907            }
908
909            let slot_index = self.slots.lock().len();
910            let new_slot = self.create_connection(slot_index).await?;
911            let new_idx = {
912                let admission = self.admission.lock();
913                let mut slots = self.slots.lock();
914
915                if admission.closed || admission.generation != generation {
916                    let _ = new_slot.cmd_tx.send(HandlerCommand::Disconnect);
917                    slots.push(new_slot);
918                    return Err(DydxWsError::Transport(
919                        "WebSocket connection was canceled by shutdown".to_string(),
920                    ));
921                }
922                slots.push(new_slot);
923                slots.len() - 1
924            };
925            log::debug!(
926                "dYdX pool slot {new_idx} connected: url={}, channel={:?}",
927                self.url,
928                channel,
929            );
930        };
931
932        let admission = self.admission.lock();
933        if admission.closed || admission.generation != generation {
934            return Err(DydxWsError::Transport(
935                "WebSocket connection pool is closed".to_string(),
936            ));
937        }
938        let mut slots = self.slots.lock();
939        let slot = &mut slots[target_idx];
940
941        slot.subscriptions_state.mark_subscribe(&topic);
942        slot.cmd_tx
943            .send(HandlerCommand::RegisterSubscription {
944                topic: topic.clone(),
945                subscription: sub_msg.clone(),
946            })
947            .map_err(|e| {
948                slot.subscriptions_state.mark_failure(&topic);
949                DydxWsError::Transport(format!("Slot {target_idx} unavailable: {e}"))
950            })?;
951
952        let payload = serde_json::to_string(&sub_msg)?;
953        if let Err(e) = slot.cmd_tx.send(HandlerCommand::SendText(payload)) {
954            slot.subscriptions_state.mark_failure(&topic);
955            let _ = slot.cmd_tx.send(HandlerCommand::UnregisterSubscription {
956                topic: topic.clone(),
957            });
958            return Err(DydxWsError::Transport(format!(
959                "Slot {target_idx} send failed: {e}"
960            )));
961        }
962
963        slot.topics.insert(topic, 1);
964        slot.channel_counts[channel as usize] =
965            slot.channel_counts[channel as usize].saturating_add(1);
966
967        Ok(())
968    }
969
970    async fn unsubscribe_topic(
971        &self,
972        channel: ChannelKind,
973        topic: String,
974        unsub_msg: DydxSubscription,
975    ) -> DydxWsResult<()> {
976        let _connect_guard = self.connect_lock.lock().await;
977        let _generation = self.admission_generation()?;
978        let admission = self.admission.lock();
979        if admission.closed {
980            return Err(DydxWsError::Transport(
981                "WebSocket connection pool is closed".to_string(),
982            ));
983        }
984        let mut slots = self.slots.lock();
985        let Some(slot_idx) = slots.iter().position(|s| s.topics.contains_key(&topic)) else {
986            return Ok(());
987        };
988
989        let slot = &mut slots[slot_idx];
990        let refcount = slot.topics.get_mut(&topic).expect("topic present");
991        if *refcount > 1 {
992            *refcount -= 1;
993            return Ok(());
994        }
995
996        slot.subscriptions_state.mark_unsubscribe(&topic);
997        let payload = serde_json::to_string(&unsub_msg)?;
998        if let Err(e) = slot.cmd_tx.send(HandlerCommand::SendText(payload)) {
999            slot.subscriptions_state.mark_subscribe(&topic);
1000            return Err(DydxWsError::Transport(format!(
1001                "Slot {slot_idx} send failed: {e}"
1002            )));
1003        }
1004        let _ = slot.cmd_tx.send(HandlerCommand::UnregisterSubscription {
1005            topic: topic.clone(),
1006        });
1007
1008        slot.topics.remove(&topic);
1009        slot.channel_counts[channel as usize] =
1010            slot.channel_counts[channel as usize].saturating_sub(1);
1011
1012        Ok(())
1013    }
1014
1015    /// Subscribes to public trade updates for a specific instrument.
1016    ///
1017    /// # Errors
1018    ///
1019    /// Returns an error if the subscription request fails.
1020    ///
1021    /// # References
1022    ///
1023    /// <https://docs.dydx.trade/developers/indexer/websockets#trades-channel>
1024    pub async fn subscribe_trades(&self, instrument_id: InstrumentId) -> DydxWsResult<()> {
1025        let ticker = Self::ticker_from_instrument_id(&instrument_id);
1026        let topic = Self::topic(DydxWsChannel::Trades, Some(&ticker));
1027        let sub = DydxSubscription {
1028            op: DydxWsOperation::Subscribe,
1029            channel: DydxWsChannel::Trades,
1030            id: Some(ticker),
1031        };
1032        self.subscribe_topic(ChannelKind::Trades, topic, sub).await
1033    }
1034
1035    /// Unsubscribes from public trade updates for a specific instrument.
1036    ///
1037    /// # Errors
1038    ///
1039    /// Returns an error if the unsubscription request fails.
1040    pub async fn unsubscribe_trades(&self, instrument_id: InstrumentId) -> DydxWsResult<()> {
1041        let ticker = Self::ticker_from_instrument_id(&instrument_id);
1042        let topic = Self::topic(DydxWsChannel::Trades, Some(&ticker));
1043        let sub = DydxSubscription {
1044            op: DydxWsOperation::Unsubscribe,
1045            channel: DydxWsChannel::Trades,
1046            id: Some(ticker),
1047        };
1048        self.unsubscribe_topic(ChannelKind::Trades, topic, sub)
1049            .await
1050    }
1051
1052    /// Subscribes to orderbook updates for a specific instrument.
1053    ///
1054    /// # Errors
1055    ///
1056    /// Returns an error if the subscription request fails.
1057    ///
1058    /// # References
1059    ///
1060    /// <https://docs.dydx.trade/developers/indexer/websockets#orderbook-channel>
1061    pub async fn subscribe_orderbook(&self, instrument_id: InstrumentId) -> DydxWsResult<()> {
1062        let ticker = Self::ticker_from_instrument_id(&instrument_id);
1063        let topic = Self::topic(DydxWsChannel::Orderbook, Some(&ticker));
1064        let sub = DydxSubscription {
1065            op: DydxWsOperation::Subscribe,
1066            channel: DydxWsChannel::Orderbook,
1067            id: Some(ticker),
1068        };
1069        self.subscribe_topic(ChannelKind::Orderbook, topic, sub)
1070            .await
1071    }
1072
1073    /// Unsubscribes from orderbook updates for a specific instrument.
1074    ///
1075    /// # Errors
1076    ///
1077    /// Returns an error if the unsubscription request fails.
1078    pub async fn unsubscribe_orderbook(&self, instrument_id: InstrumentId) -> DydxWsResult<()> {
1079        let ticker = Self::ticker_from_instrument_id(&instrument_id);
1080        let topic = Self::topic(DydxWsChannel::Orderbook, Some(&ticker));
1081        let sub = DydxSubscription {
1082            op: DydxWsOperation::Unsubscribe,
1083            channel: DydxWsChannel::Orderbook,
1084            id: Some(ticker),
1085        };
1086        self.unsubscribe_topic(ChannelKind::Orderbook, topic, sub)
1087            .await
1088    }
1089
1090    /// Subscribes to candle/kline updates for a specific instrument.
1091    ///
1092    /// # Errors
1093    ///
1094    /// Returns an error if the subscription request fails.
1095    ///
1096    /// # References
1097    ///
1098    /// <https://docs.dydx.trade/developers/indexer/websockets#candles-channel>
1099    pub async fn subscribe_candles(
1100        &self,
1101        instrument_id: InstrumentId,
1102        resolution: &str,
1103    ) -> DydxWsResult<()> {
1104        let ticker = Self::ticker_from_instrument_id(&instrument_id);
1105        let id = format!("{ticker}/{resolution}");
1106        let topic = Self::topic(DydxWsChannel::Candles, Some(&id));
1107        let sub = DydxSubscription {
1108            op: DydxWsOperation::Subscribe,
1109            channel: DydxWsChannel::Candles,
1110            id: Some(id),
1111        };
1112        self.subscribe_topic(ChannelKind::Candles, topic, sub).await
1113    }
1114
1115    /// Unsubscribes from candle/kline updates for a specific instrument.
1116    ///
1117    /// # Errors
1118    ///
1119    /// Returns an error if the unsubscription request fails.
1120    pub async fn unsubscribe_candles(
1121        &self,
1122        instrument_id: InstrumentId,
1123        resolution: &str,
1124    ) -> DydxWsResult<()> {
1125        let ticker = Self::ticker_from_instrument_id(&instrument_id);
1126        let id = format!("{ticker}/{resolution}");
1127        let topic = Self::topic(DydxWsChannel::Candles, Some(&id));
1128        let sub = DydxSubscription {
1129            op: DydxWsOperation::Unsubscribe,
1130            channel: DydxWsChannel::Candles,
1131            id: Some(id),
1132        };
1133        self.unsubscribe_topic(ChannelKind::Candles, topic, sub)
1134            .await
1135    }
1136
1137    /// Subscribes to market updates for all instruments.
1138    ///
1139    /// # Errors
1140    ///
1141    /// Returns an error if the subscription request fails.
1142    ///
1143    /// # References
1144    ///
1145    /// <https://docs.dydx.trade/developers/indexer/websockets#markets-channel>
1146    pub async fn subscribe_markets(&self) -> DydxWsResult<()> {
1147        let topic = Self::topic(DydxWsChannel::Markets, None);
1148        let sub = DydxSubscription {
1149            op: DydxWsOperation::Subscribe,
1150            channel: DydxWsChannel::Markets,
1151            id: None,
1152        };
1153        self.subscribe_topic(ChannelKind::Markets, topic, sub).await
1154    }
1155
1156    /// Unsubscribes from market updates.
1157    ///
1158    /// # Errors
1159    ///
1160    /// Returns an error if the unsubscription request fails.
1161    pub async fn unsubscribe_markets(&self) -> DydxWsResult<()> {
1162        let topic = Self::topic(DydxWsChannel::Markets, None);
1163        let sub = DydxSubscription {
1164            op: DydxWsOperation::Unsubscribe,
1165            channel: DydxWsChannel::Markets,
1166            id: None,
1167        };
1168        self.unsubscribe_topic(ChannelKind::Markets, topic, sub)
1169            .await
1170    }
1171
1172    /// Subscribes to subaccount updates (orders, fills, positions, balances).
1173    ///
1174    /// This requires authentication and will only work for private WebSocket clients
1175    /// created with [`Self::new_private`]. Subaccount streams stay pinned to the
1176    /// primary slot: the Indexer caps them at 256 per connection, which is well
1177    /// above realistic per-process usage and keeps related fill/position events
1178    /// on a single in-order stream.
1179    ///
1180    /// # Errors
1181    ///
1182    /// Returns an error if the client was not created with credentials or if the
1183    /// subscription request fails.
1184    ///
1185    /// # References
1186    ///
1187    /// <https://docs.dydx.trade/developers/indexer/websockets#subaccounts-channel>
1188    pub async fn subscribe_subaccount(
1189        &self,
1190        address: &str,
1191        subaccount_number: u32,
1192    ) -> DydxWsResult<()> {
1193        if !self.requires_auth {
1194            return Err(DydxWsError::Authentication(
1195                "Subaccount subscriptions require authentication. Use new_private() to create an authenticated client".to_string(),
1196            ));
1197        }
1198        let id = format!("{address}/{subaccount_number}");
1199        let topic = Self::topic(DydxWsChannel::Subaccounts, Some(&id));
1200        let sub = DydxSubscription {
1201            op: DydxWsOperation::Subscribe,
1202            channel: DydxWsChannel::Subaccounts,
1203            id: Some(id),
1204        };
1205        self.subscribe_pinned(topic, sub).await
1206    }
1207
1208    /// Unsubscribes from subaccount updates.
1209    ///
1210    /// # Errors
1211    ///
1212    /// Returns an error if the unsubscription request fails.
1213    pub async fn unsubscribe_subaccount(
1214        &self,
1215        address: &str,
1216        subaccount_number: u32,
1217    ) -> DydxWsResult<()> {
1218        let id = format!("{address}/{subaccount_number}");
1219        let topic = Self::topic(DydxWsChannel::Subaccounts, Some(&id));
1220        let sub = DydxSubscription {
1221            op: DydxWsOperation::Unsubscribe,
1222            channel: DydxWsChannel::Subaccounts,
1223            id: Some(id),
1224        };
1225        self.unsubscribe_pinned(topic, sub).await
1226    }
1227
1228    /// Subscribes to block height updates.
1229    ///
1230    /// # Errors
1231    ///
1232    /// Returns an error if the subscription request fails.
1233    ///
1234    /// # References
1235    ///
1236    /// <https://docs.dydx.trade/developers/indexer/websockets#block-height-channel>
1237    pub async fn subscribe_block_height(&self) -> DydxWsResult<()> {
1238        let topic = Self::topic(DydxWsChannel::BlockHeight, None);
1239        let sub = DydxSubscription {
1240            op: DydxWsOperation::Subscribe,
1241            channel: DydxWsChannel::BlockHeight,
1242            id: None,
1243        };
1244        self.subscribe_pinned(topic, sub).await
1245    }
1246
1247    /// Unsubscribes from block height updates.
1248    ///
1249    /// # Errors
1250    ///
1251    /// Returns an error if the unsubscription request fails.
1252    pub async fn unsubscribe_block_height(&self) -> DydxWsResult<()> {
1253        let topic = Self::topic(DydxWsChannel::BlockHeight, None);
1254        let sub = DydxSubscription {
1255            op: DydxWsOperation::Unsubscribe,
1256            channel: DydxWsChannel::BlockHeight,
1257            id: None,
1258        };
1259        self.unsubscribe_pinned(topic, sub).await
1260    }
1261
1262    async fn subscribe_pinned(&self, topic: String, sub_msg: DydxSubscription) -> DydxWsResult<()> {
1263        let _connect_guard = self.connect_lock.lock().await;
1264        let generation = self.admission_generation()?;
1265
1266        {
1267            let admission = self.admission.lock();
1268            if admission.closed || admission.generation != generation {
1269                return Err(DydxWsError::Transport(
1270                    "WebSocket connection pool is closed".to_string(),
1271                ));
1272            }
1273            let mut slots = self.slots.lock();
1274            if let Some(slot) = slots.iter_mut().find(|s| s.topics.contains_key(&topic)) {
1275                *slot.topics.get_mut(&topic).expect("topic refcount present") += 1;
1276                return Ok(());
1277            }
1278        }
1279
1280        if self.slots.lock().is_empty() {
1281            let new_slot = self.create_connection(0).await?;
1282            let admission = self.admission.lock();
1283            let mut slots = self.slots.lock();
1284
1285            if admission.closed || admission.generation != generation {
1286                let _ = new_slot.cmd_tx.send(HandlerCommand::Disconnect);
1287                slots.push(new_slot);
1288                return Err(DydxWsError::Transport(
1289                    "WebSocket connection was canceled by shutdown".to_string(),
1290                ));
1291            }
1292            self.connection_mode.store(new_slot.connection_mode.clone());
1293            slots.push(new_slot);
1294        }
1295
1296        let admission = self.admission.lock();
1297        if admission.closed || admission.generation != generation {
1298            return Err(DydxWsError::Transport(
1299                "WebSocket connection pool is closed".to_string(),
1300            ));
1301        }
1302        let mut slots = self.slots.lock();
1303        let slot = slots.first_mut().expect("primary slot exists");
1304        slot.subscriptions_state.mark_subscribe(&topic);
1305        slot.cmd_tx
1306            .send(HandlerCommand::RegisterSubscription {
1307                topic: topic.clone(),
1308                subscription: sub_msg.clone(),
1309            })
1310            .map_err(|e| {
1311                slot.subscriptions_state.mark_failure(&topic);
1312                DydxWsError::Transport(format!("Primary slot unavailable: {e}"))
1313            })?;
1314        let payload = serde_json::to_string(&sub_msg)?;
1315        if let Err(e) = slot.cmd_tx.send(HandlerCommand::SendText(payload)) {
1316            slot.subscriptions_state.mark_failure(&topic);
1317            let _ = slot.cmd_tx.send(HandlerCommand::UnregisterSubscription {
1318                topic: topic.clone(),
1319            });
1320            return Err(DydxWsError::Transport(format!(
1321                "Primary slot send failed: {e}"
1322            )));
1323        }
1324        slot.topics.insert(topic, 1);
1325        Ok(())
1326    }
1327
1328    async fn unsubscribe_pinned(
1329        &self,
1330        topic: String,
1331        unsub_msg: DydxSubscription,
1332    ) -> DydxWsResult<()> {
1333        let _connect_guard = self.connect_lock.lock().await;
1334        let _generation = self.admission_generation()?;
1335        let admission = self.admission.lock();
1336        if admission.closed {
1337            return Err(DydxWsError::Transport(
1338                "WebSocket connection pool is closed".to_string(),
1339            ));
1340        }
1341        let mut slots = self.slots.lock();
1342        let Some(slot) = slots.first_mut() else {
1343            return Ok(());
1344        };
1345        let Some(refcount) = slot.topics.get_mut(&topic) else {
1346            return Ok(());
1347        };
1348
1349        if *refcount > 1 {
1350            *refcount -= 1;
1351            return Ok(());
1352        }
1353        slot.subscriptions_state.mark_unsubscribe(&topic);
1354        let payload = serde_json::to_string(&unsub_msg)?;
1355        if let Err(e) = slot.cmd_tx.send(HandlerCommand::SendText(payload)) {
1356            slot.subscriptions_state.mark_subscribe(&topic);
1357            return Err(DydxWsError::Transport(format!(
1358                "Primary slot send failed: {e}"
1359            )));
1360        }
1361        let _ = slot.cmd_tx.send(HandlerCommand::UnregisterSubscription {
1362            topic: topic.clone(),
1363        });
1364        slot.topics.remove(&topic);
1365        Ok(())
1366    }
1367}
1368
1369#[derive(Debug)]
1370struct ConnectionSlots {
1371    slots: Mutex<Vec<ConnectionSlot>>,
1372    shutdown_errors: Mutex<Vec<String>>,
1373}
1374
1375impl ConnectionSlots {
1376    fn new() -> Self {
1377        Self {
1378            slots: Mutex::new(Vec::new()),
1379            shutdown_errors: Mutex::new(Vec::new()),
1380        }
1381    }
1382
1383    fn push_shutdown_error(&self, error: String) {
1384        self.shutdown_errors.lock().push(error);
1385    }
1386
1387    fn take_shutdown_errors(&self) -> Vec<String> {
1388        std::mem::take(&mut *self.shutdown_errors.lock())
1389    }
1390}
1391
1392impl std::ops::Deref for ConnectionSlots {
1393    type Target = Mutex<Vec<ConnectionSlot>>;
1394
1395    fn deref(&self) -> &Self::Target {
1396        &self.slots
1397    }
1398}
1399
1400impl Drop for ConnectionSlots {
1401    fn drop(&mut self) {
1402        for slot in self.slots.get_mut().iter() {
1403            if let Some(handle) = slot.handler_task.as_ref() {
1404                handle.abort();
1405            }
1406
1407            if let Some(control) = &slot.socket_control {
1408                control.deregister();
1409            }
1410        }
1411    }
1412}
1413
1414#[derive(Debug)]
1415struct PoolAdmission {
1416    generation: u64,
1417    closed: bool,
1418}
1419
1420struct ConnectionSlotBatch<'a> {
1421    owner: &'a Mutex<Vec<ConnectionSlot>>,
1422    slots: Vec<ConnectionSlot>,
1423}
1424
1425impl<'a> ConnectionSlotBatch<'a> {
1426    fn take(owner: &'a Mutex<Vec<ConnectionSlot>>) -> Self {
1427        let slots = std::mem::take(&mut *owner.lock());
1428        Self { owner, slots }
1429    }
1430}
1431
1432impl Drop for ConnectionSlotBatch<'_> {
1433    fn drop(&mut self) {
1434        self.owner.lock().extend(self.slots.drain(..));
1435    }
1436}
1437
1438// Scopes per-slot reconnect cleanup of in-progress bars to the candle topics
1439// owned by the reconnecting connection, so one slot's reconnect does not
1440// discard bars still aggregating on other healthy connections.
1441pub(crate) fn candle_ids_from_topics(topics: &[String]) -> AHashSet<String> {
1442    let prefix = format!(
1443        "{}{}",
1444        DydxWsChannel::Candles.as_ref(),
1445        DYDX_WS_TOPIC_DELIMITER
1446    );
1447    topics
1448        .iter()
1449        .filter_map(|topic| topic.strip_prefix(&prefix).map(ToString::to_string))
1450        .collect()
1451}
1452
1453#[cfg(test)]
1454mod tests {
1455    use nautilus_core::string::secret::REDACTED;
1456    use rstest::rstest;
1457
1458    use super::*;
1459
1460    #[rstest]
1461    fn test_debug_redacts_proxy_url() {
1462        let proxy_url = "http://user:password@proxy.example:8080";
1463        let client = DydxWebSocketClient::new_public(
1464            "wss://test".to_string(),
1465            None,
1466            Some(proxy_url.to_string()),
1467        );
1468
1469        let debug = format!("{client:?}");
1470
1471        assert!(debug.contains(REDACTED));
1472        assert!(!debug.contains(proxy_url));
1473    }
1474
1475    #[rstest]
1476    fn test_candle_ids_from_topics_extracts_only_candle_ids() {
1477        let topics = vec![
1478            "v4_candles:BTC-USD/1MIN".to_string(),
1479            "v4_trades:BTC-USD".to_string(),
1480            "v4_orderbook:ETH-USD".to_string(),
1481            "v4_candles:ETH-USD/5MINS".to_string(),
1482        ];
1483
1484        let ids = candle_ids_from_topics(&topics);
1485
1486        assert_eq!(ids.len(), 2);
1487        assert!(ids.contains("BTC-USD/1MIN"));
1488        assert!(ids.contains("ETH-USD/5MINS"));
1489        assert!(!ids.contains("BTC-USD"));
1490    }
1491
1492    #[tokio::test]
1493    async fn test_drop_clone_does_not_stop_connection_pool() {
1494        let client = DydxWebSocketClient::new_public("wss://test".to_string(), None, None);
1495        let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1496        client.slots.lock().push(ConnectionSlot {
1497            cmd_tx,
1498            topics: AHashMap::new(),
1499            channel_counts: [0; CHANNEL_KIND_COUNT],
1500            subscriptions_state: SubscriptionState::new(DYDX_WS_TOPIC_DELIMITER),
1501            handler_task: TaskSlot::from_handle(tokio::spawn(std::future::pending())),
1502            connection_mode: Arc::new(AtomicU8::new(ConnectionMode::Active as u8)),
1503            socket_control: None,
1504        });
1505        let clone = client.clone();
1506
1507        drop(clone);
1508
1509        let slots = client.slots.lock();
1510        assert_eq!(slots.len(), 1);
1511        assert!(
1512            !slots[0]
1513                .handler_task
1514                .as_ref()
1515                .expect("handler task")
1516                .is_finished()
1517        );
1518    }
1519
1520    #[tokio::test]
1521    async fn test_cancelled_disconnect_retains_connection_slot() {
1522        let mut client = DydxWebSocketClient::new_public("wss://test".to_string(), None, None);
1523        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1524        client.slots.lock().push(ConnectionSlot {
1525            cmd_tx,
1526            topics: AHashMap::new(),
1527            channel_counts: [0; CHANNEL_KIND_COUNT],
1528            subscriptions_state: SubscriptionState::new(DYDX_WS_TOPIC_DELIMITER),
1529            handler_task: TaskSlot::from_handle(tokio::spawn(std::future::pending())),
1530            connection_mode: Arc::new(AtomicU8::new(ConnectionMode::Active as u8)),
1531            socket_control: None,
1532        });
1533
1534        {
1535            let disconnect = client.disconnect();
1536            tokio::pin!(disconnect);
1537            tokio::select! {
1538                result = &mut disconnect => panic!("disconnect completed unexpectedly: {result:?}"),
1539                command = cmd_rx.recv() => assert!(command.is_some()),
1540            }
1541        }
1542
1543        let slots = client.slots.lock();
1544        assert_eq!(slots.len(), 1);
1545        assert!(slots[0].handler_task.is_some());
1546    }
1547}
nautilus_dydx/websocket/client.rs

nautilus_dydx/websocket/
client.rs
