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

nautilus_dydx/websocket/
client.rs
