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

nautilus_dydx/websocket/
client.rs
