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, Mutex,
68        atomic::{AtomicBool, AtomicU8, Ordering},
69    },
70    time::Duration,
71};
72
73use ahash::{AHashMap, AHashSet};
74use arc_swap::ArcSwap;
75use dashmap::DashMap;
76use nautilus_common::live::get_runtime;
77use nautilus_model::{
78    data::BarType,
79    identifiers::{AccountId, InstrumentId},
80    instruments::InstrumentAny,
81};
82use nautilus_network::{
83    mode::ConnectionMode,
84    ratelimiter::quota::Quota,
85    websocket::{
86        AuthTracker, SubscriptionState, TransportBackend, WebSocketClient, WebSocketConfig,
87        channel_message_handler,
88    },
89};
90use ustr::Ustr;
91
92use super::{
93    dispatch::DydxWsDispatchState,
94    enums::{DydxWsChannel, DydxWsOperation, DydxWsOutputMessage},
95    error::{DydxWsError, DydxWsResult},
96    handler::{FeedHandler, HandlerCommand},
97    messages::DydxSubscription,
98};
99use crate::{
100    common::{credential::DydxCredential, instrument_cache::InstrumentCache},
101    execution::encoder::ClientOrderIdEncoder,
102};
103
104/// Identifies a dYdX channel for per-channel capacity accounting in the pool.
105#[derive(Copy, Clone, Debug)]
106#[repr(u8)]
107enum ChannelKind {
108    Trades = 0,
109    Candles = 1,
110    Orderbook = 2,
111    Markets = 3,
112}
113
114const CHANNEL_KIND_COUNT: usize = 4;
115
116/// Per-connection state inside the pool.
117#[derive(Debug)]
118struct ConnectionSlot {
119    cmd_tx: tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
120    topics: AHashMap<String, u32>,
121    channel_counts: [u16; CHANNEL_KIND_COUNT],
122    subscriptions_state: SubscriptionState,
123    handler_task: Option<tokio::task::JoinHandle<()>>,
124    connection_mode: Arc<AtomicU8>,
125}
126
127/// WebSocket client for dYdX v4 market data and account streams.
128///
129/// # Authentication
130///
131/// dYdX v4 does not require traditional API key signatures for WebSocket connections.
132/// Public channels work without any credentials. Private channels (subaccounts) only
133/// need the wallet address included in the subscription message.
134///
135/// The [`DydxCredential`] stored in this client is used for:
136/// - Providing the wallet address for private channel subscriptions
137/// - Transaction signing (when placing orders via the validator node)
138///
139/// It is **NOT** used for WebSocket message signing or authentication.
140///
141/// # Architecture
142///
143/// The client owns a small pool of connection slots. Each slot has its own
144/// `WebSocketClient`, [`FeedHandler`] task, command channel, and
145/// [`SubscriptionState`]. All slots write parsed events into a single shared
146/// output channel so callers see one merged stream.
147#[derive(Debug)]
148#[cfg_attr(
149    feature = "python",
150    pyo3::pyclass(module = "nautilus_trader.core.nautilus_pyo3.dydx", from_py_object)
151)]
152#[cfg_attr(
153    feature = "python",
154    pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.adapters.dydx")
155)]
156pub struct DydxWebSocketClient {
157    url: String,
158    credential: Option<Arc<DydxCredential>>,
159    requires_auth: bool,
160    auth_tracker: AuthTracker,
161    slots: Arc<Mutex<Vec<ConnectionSlot>>>,
162    connect_lock: Arc<tokio::sync::Mutex<()>>,
163    connection_mode: Arc<ArcSwap<AtomicU8>>,
164    signal: Arc<AtomicBool>,
165    instrument_cache: Arc<InstrumentCache>,
166    account_id: Option<AccountId>,
167    heartbeat: Option<u64>,
168    out_tx: Arc<Mutex<Option<tokio::sync::mpsc::UnboundedSender<DydxWsOutputMessage>>>>,
169    out_rx: Arc<Mutex<Option<tokio::sync::mpsc::UnboundedReceiver<DydxWsOutputMessage>>>>,
170    encoder: Arc<ClientOrderIdEncoder>,
171    bar_types: Arc<DashMap<String, BarType>>,
172    bars_timestamp_on_close: Arc<AtomicBool>,
173    ws_dispatch_state: Arc<DydxWsDispatchState>,
174    transport_backend: TransportBackend,
175    proxy_url: Option<String>,
176    max_ws_connections: usize,
177    per_channel_limit: usize,
178}
179
180impl Clone for DydxWebSocketClient {
181    fn clone(&self) -> Self {
182        Self {
183            url: self.url.clone(),
184            credential: self.credential.clone(),
185            requires_auth: self.requires_auth,
186            auth_tracker: self.auth_tracker.clone(),
187            slots: self.slots.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        }
205    }
206}
207
208impl DydxWebSocketClient {
209    /// Creates a new public WebSocket client for market data.
210    ///
211    /// This creates a new independent instrument cache. To share a cache with
212    /// the HTTP client, use [`Self::new_public_with_cache`] instead.
213    #[must_use]
214    pub fn new_public(url: String, heartbeat: Option<u64>, proxy_url: Option<String>) -> Self {
215        Self::new_public_with_cache(
216            url,
217            Arc::new(InstrumentCache::new()),
218            heartbeat,
219            TransportBackend::default(),
220            proxy_url,
221        )
222    }
223
224    /// Creates a new public WebSocket client with a shared instrument cache.
225    ///
226    /// Use this when you want to share instrument data with the HTTP client.
227    #[must_use]
228    pub fn new_public_with_cache(
229        url: String,
230        instrument_cache: Arc<InstrumentCache>,
231        heartbeat: Option<u64>,
232        transport_backend: TransportBackend,
233        proxy_url: Option<String>,
234    ) -> Self {
235        Self::new_public_with_cache_and_pool(
236            url,
237            instrument_cache,
238            heartbeat,
239            transport_backend,
240            proxy_url,
241            DEFAULT_MAX_WS_CONNECTIONS,
242            DEFAULT_PER_CHANNEL_SUBSCRIPTION_LIMIT,
243        )
244    }
245
246    /// Creates a new public WebSocket client with full pool configuration.
247    #[must_use]
248    pub fn new_public_with_cache_and_pool(
249        url: String,
250        instrument_cache: Arc<InstrumentCache>,
251        heartbeat: Option<u64>,
252        transport_backend: TransportBackend,
253        proxy_url: Option<String>,
254        max_ws_connections: usize,
255        per_channel_limit: usize,
256    ) -> Self {
257        Self::new_inner(
258            url,
259            None,
260            false,
261            instrument_cache,
262            None,
263            heartbeat,
264            transport_backend,
265            proxy_url,
266            max_ws_connections,
267            per_channel_limit,
268        )
269    }
270
271    /// Creates a new private WebSocket client for account updates.
272    ///
273    /// This creates a new independent instrument cache. To share a cache with
274    /// the HTTP client, use [`Self::new_private_with_cache`] instead.
275    #[must_use]
276    pub fn new_private(
277        url: String,
278        credential: DydxCredential,
279        account_id: AccountId,
280        heartbeat: Option<u64>,
281        proxy_url: Option<String>,
282    ) -> Self {
283        Self::new_private_with_cache(
284            url,
285            credential,
286            account_id,
287            Arc::new(InstrumentCache::new()),
288            heartbeat,
289            TransportBackend::default(),
290            proxy_url,
291        )
292    }
293
294    /// Creates a new private WebSocket client with a shared instrument cache.
295    ///
296    /// Use this when you want to share instrument data with the HTTP client.
297    #[must_use]
298    pub fn new_private_with_cache(
299        url: String,
300        credential: DydxCredential,
301        account_id: AccountId,
302        instrument_cache: Arc<InstrumentCache>,
303        heartbeat: Option<u64>,
304        transport_backend: TransportBackend,
305        proxy_url: Option<String>,
306    ) -> Self {
307        Self::new_inner(
308            url,
309            Some(Arc::new(credential)),
310            true,
311            instrument_cache,
312            Some(account_id),
313            heartbeat,
314            transport_backend,
315            proxy_url,
316            DEFAULT_MAX_WS_CONNECTIONS,
317            DEFAULT_PER_CHANNEL_SUBSCRIPTION_LIMIT,
318        )
319    }
320
321    #[allow(clippy::too_many_arguments)]
322    fn new_inner(
323        url: String,
324        credential: Option<Arc<DydxCredential>>,
325        requires_auth: bool,
326        instrument_cache: Arc<InstrumentCache>,
327        account_id: Option<AccountId>,
328        heartbeat: Option<u64>,
329        transport_backend: TransportBackend,
330        proxy_url: Option<String>,
331        max_ws_connections: usize,
332        per_channel_limit: usize,
333    ) -> Self {
334        Self {
335            url,
336            credential,
337            requires_auth,
338            auth_tracker: AuthTracker::new(),
339            slots: Arc::new(Mutex::new(Vec::new())),
340            connect_lock: Arc::new(tokio::sync::Mutex::new(())),
341            connection_mode: Arc::new(ArcSwap::from_pointee(AtomicU8::new(
342                ConnectionMode::Closed as u8,
343            ))),
344            signal: Arc::new(AtomicBool::new(false)),
345            instrument_cache,
346            account_id,
347            heartbeat,
348            out_tx: Arc::new(Mutex::new(None)),
349            out_rx: Arc::new(Mutex::new(None)),
350            encoder: Arc::new(ClientOrderIdEncoder::new()),
351            bar_types: Arc::new(DashMap::new()),
352            bars_timestamp_on_close: Arc::new(AtomicBool::new(true)),
353            ws_dispatch_state: Arc::new(DydxWsDispatchState::default()),
354            transport_backend,
355            proxy_url,
356            max_ws_connections: max_ws_connections.max(1),
357            per_channel_limit: per_channel_limit.max(1),
358        }
359    }
360
361    /// Returns the credential associated with this client, if any.
362    #[must_use]
363    pub fn credential(&self) -> Option<&Arc<DydxCredential>> {
364        self.credential.as_ref()
365    }
366
367    /// Returns `true` when any connection in the pool is connected.
368    #[must_use]
369    #[expect(clippy::missing_panics_doc, reason = "mutex poisoning is not expected")]
370    pub fn is_connected(&self) -> bool {
371        let slots = self.slots.lock().expect("slots lock poisoned");
372        slots.iter().any(|s| {
373            let mode = s.connection_mode.load(Ordering::Relaxed);
374            mode == ConnectionMode::Active as u8 || mode == ConnectionMode::Reconnect as u8
375        })
376    }
377
378    /// Returns the URL of this WebSocket client.
379    #[must_use]
380    pub fn url(&self) -> &str {
381        &self.url
382    }
383
384    /// Returns a clone of the connection mode atomic reference.
385    ///
386    /// With sharding, the returned atomic tracks the **primary** slot (slot 0)
387    /// only; use [`Self::is_connected`] for a pool-wide check.
388    #[must_use]
389    pub fn connection_mode_atomic(&self) -> Arc<ArcSwap<AtomicU8>> {
390        self.connection_mode.clone()
391    }
392
393    /// Returns the current number of active slots in the pool.
394    #[must_use]
395    #[expect(clippy::missing_panics_doc, reason = "mutex poisoning is not expected")]
396    pub fn pool_size(&self) -> usize {
397        self.slots.lock().expect("slots lock poisoned").len()
398    }
399
400    /// Returns the configured maximum number of pool connections.
401    #[must_use]
402    pub const fn max_ws_connections(&self) -> usize {
403        self.max_ws_connections
404    }
405
406    /// Returns the configured per-channel subscription limit.
407    #[must_use]
408    pub const fn per_channel_limit(&self) -> usize {
409        self.per_channel_limit
410    }
411
412    /// Sets the account ID for account message parsing.
413    pub fn set_account_id(&mut self, account_id: AccountId) {
414        self.account_id = Some(account_id);
415    }
416
417    /// Returns the account ID if set.
418    #[must_use]
419    pub fn account_id(&self) -> Option<AccountId> {
420        self.account_id
421    }
422
423    /// Replaces the instrument cache with an externally shared one.
424    ///
425    /// Use this to share the HTTP client's cache (which includes CLOB pair ID
426    /// and market ticker indices) with the WebSocket client. Must be called
427    /// before `connect()`.
428    pub fn set_instrument_cache(&mut self, cache: Arc<InstrumentCache>) {
429        self.instrument_cache = cache;
430    }
431
432    /// Caches a single instrument.
433    ///
434    /// Any existing instrument with the same ID will be replaced.
435    pub fn cache_instrument(&self, instrument: InstrumentAny) {
436        self.instrument_cache.insert_instrument_only(instrument);
437    }
438
439    /// Caches multiple instruments.
440    ///
441    /// Any existing instruments with the same IDs will be replaced.
442    pub fn cache_instruments(&self, instruments: Vec<InstrumentAny>) {
443        log::debug!(
444            "Caching {} instruments in WebSocket client",
445            instruments.len()
446        );
447        self.instrument_cache.insert_instruments_only(instruments);
448    }
449
450    /// Returns a reference to the shared instrument cache.
451    #[must_use]
452    pub fn instrument_cache(&self) -> &Arc<InstrumentCache> {
453        &self.instrument_cache
454    }
455
456    /// Returns a reference to the shared client order ID encoder.
457    #[must_use]
458    pub fn encoder(&self) -> &Arc<ClientOrderIdEncoder> {
459        &self.encoder
460    }
461
462    /// Returns a reference to the bar type registrations map.
463    #[must_use]
464    pub fn bar_types(&self) -> &Arc<DashMap<String, BarType>> {
465        &self.bar_types
466    }
467
468    /// Returns a reference to the shared WebSocket dispatch state.
469    pub fn ws_dispatch_state(&self) -> &Arc<DydxWsDispatchState> {
470        &self.ws_dispatch_state
471    }
472
473    /// Sets whether bar timestamps use the close time.
474    pub fn set_bars_timestamp_on_close(&self, value: bool) {
475        self.bars_timestamp_on_close.store(value, Ordering::Relaxed);
476    }
477
478    /// Returns whether bar timestamps use the close time.
479    #[must_use]
480    pub fn bars_timestamp_on_close(&self) -> bool {
481        self.bars_timestamp_on_close.load(Ordering::Relaxed)
482    }
483
484    /// Returns all cached instruments.
485    ///
486    /// This is a snapshot of the current cache contents.
487    #[must_use]
488    pub fn all_instruments(&self) -> Vec<InstrumentAny> {
489        self.instrument_cache.all_instruments()
490    }
491
492    /// Returns the number of cached instruments.
493    #[must_use]
494    pub fn cached_instruments_count(&self) -> usize {
495        self.instrument_cache.len()
496    }
497
498    /// Retrieves an instrument from the cache by InstrumentId.
499    ///
500    /// Returns `None` if the instrument is not found.
501    #[must_use]
502    pub fn get_instrument(&self, instrument_id: &InstrumentId) -> Option<InstrumentAny> {
503        self.instrument_cache.get(instrument_id)
504    }
505
506    /// Retrieves an instrument from the cache by market ticker (e.g., "BTC-USD").
507    ///
508    /// Returns `None` if the instrument is not found.
509    #[must_use]
510    pub fn get_instrument_by_market(&self, ticker: &str) -> Option<InstrumentAny> {
511        self.instrument_cache.get_by_market(ticker)
512    }
513
514    /// Takes ownership of the inbound message receiver.
515    /// Returns None if the receiver has already been taken or not connected.
516    #[expect(clippy::missing_panics_doc, reason = "mutex poisoning is not expected")]
517    pub fn take_receiver(
518        &mut self,
519    ) -> Option<tokio::sync::mpsc::UnboundedReceiver<DydxWsOutputMessage>> {
520        self.out_rx.lock().expect("out_rx lock poisoned").take()
521    }
522
523    /// Returns a stream of venue-specific WebSocket messages.
524    ///
525    /// Takes ownership of the message receiver and returns it as a `Stream`.
526    ///
527    /// # Panics
528    ///
529    /// Panics if the receiver has already been taken or the receiver mutex is poisoned.
530    pub fn stream(
531        &mut self,
532    ) -> impl futures_util::Stream<Item = DydxWsOutputMessage> + Send + 'static {
533        let mut rx = self
534            .out_rx
535            .lock()
536            .expect("out_rx lock poisoned")
537            .take()
538            .expect("Message stream receiver already taken or not connected");
539
540        async_stream::stream! {
541            while let Some(msg) = rx.recv().await {
542                yield msg;
543            }
544        }
545    }
546
547    /// Connects the websocket client and opens the primary pool slot.
548    ///
549    /// Additional slots are spawned lazily by `subscribe_*` methods once the
550    /// per-channel limit is reached on every existing slot.
551    ///
552    /// # Errors
553    ///
554    /// Returns an error if the connection cannot be established.
555    #[expect(clippy::missing_panics_doc, reason = "mutex poisoning is not expected")]
556    pub async fn connect(&mut self) -> DydxWsResult<()> {
557        if self.is_connected() {
558            return Ok(());
559        }
560
561        self.signal.store(false, Ordering::Release);
562
563        let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<DydxWsOutputMessage>();
564        {
565            let mut guard = self.out_tx.lock().expect("out_tx lock poisoned");
566            *guard = Some(out_tx);
567        }
568        {
569            let mut guard = self.out_rx.lock().expect("out_rx lock poisoned");
570            *guard = Some(out_rx);
571        }
572
573        let slot = self.create_connection().await?;
574        self.connection_mode.store(slot.connection_mode.clone());
575        self.slots.lock().expect("slots lock poisoned").push(slot);
576
577        log::debug!("Connected dYdX WebSocket pool: {}", self.url);
578        Ok(())
579    }
580
581    /// Disconnects all websocket connections in the pool.
582    ///
583    /// # Errors
584    ///
585    /// Returns an error if the underlying clients cannot be accessed.
586    #[expect(clippy::missing_panics_doc, reason = "mutex poisoning is not expected")]
587    pub async fn disconnect(&mut self) -> DydxWsResult<()> {
588        self.signal.store(true, Ordering::Release);
589
590        let slots: Vec<ConnectionSlot> = {
591            let mut guard = self.slots.lock().expect("slots lock poisoned");
592            guard.drain(..).collect()
593        };
594
595        for mut slot in slots {
596            let _ = slot.cmd_tx.send(HandlerCommand::Disconnect);
597            if let Some(task) = slot.handler_task.take() {
598                let abort_handle = task.abort_handle();
599                match tokio::time::timeout(Duration::from_secs(2), task).await {
600                    Ok(Ok(())) => log::debug!("Handler task completed"),
601                    Ok(Err(e)) => log::error!("Handler task error: {e:?}"),
602                    Err(_) => {
603                        log::warn!("Timeout waiting for handler task, aborting");
604                        abort_handle.abort();
605                    }
606                }
607            }
608        }
609
610        self.connection_mode
611            .store(Arc::new(AtomicU8::new(ConnectionMode::Closed as u8)));
612
613        *self.out_tx.lock().expect("out_tx lock poisoned") = None;
614        *self.out_rx.lock().expect("out_rx lock poisoned") = None;
615
616        log::debug!("Disconnected dYdX WebSocket pool");
617        Ok(())
618    }
619
620    /// Sends a command directly to the primary slot (slot 0).
621    ///
622    /// # Errors
623    ///
624    /// Returns an error if no slot exists or the handler task has terminated.
625    #[expect(clippy::missing_panics_doc, reason = "mutex poisoning is not expected")]
626    pub fn send_command(&self, cmd: HandlerCommand) -> DydxWsResult<()> {
627        let slots = self.slots.lock().expect("slots lock poisoned");
628        let slot = slots
629            .first()
630            .ok_or_else(|| DydxWsError::Transport("No pool slots available".to_string()))?;
631        slot.cmd_tx.send(cmd).map_err(|e| {
632            DydxWsError::Transport(format!("Failed to send command to slot 0: {e}"))
633        })?;
634        Ok(())
635    }
636
637    async fn create_connection(&self) -> DydxWsResult<ConnectionSlot> {
638        let (message_handler, raw_rx) = channel_message_handler();
639
640        let cfg = WebSocketConfig {
641            url: self.url.clone(),
642            headers: vec![],
643            heartbeat: self.heartbeat,
644            heartbeat_msg: None,
645            reconnect_timeout_ms: Some(15_000),
646            reconnect_delay_initial_ms: Some(250),
647            reconnect_delay_max_ms: Some(5_000),
648            reconnect_backoff_factor: Some(2.0),
649            reconnect_jitter_ms: Some(200),
650            reconnect_max_attempts: None,
651            idle_timeout_ms: None,
652            backend: self.transport_backend,
653            proxy_url: self.proxy_url.clone(),
654        };
655
656        let client = WebSocketClient::connect(
657            cfg,
658            Some(message_handler),
659            None,
660            None,
661            vec![],
662            Some(*DYDX_WS_SUBSCRIPTION_QUOTA),
663        )
664        .await
665        .map_err(|e| DydxWsError::Transport(e.to_string()))?;
666
667        let connection_mode = client.connection_mode_atomic();
668        let subscriptions_state = SubscriptionState::new(DYDX_WS_TOPIC_DELIMITER);
669
670        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
671
672        let out_tx = self
673            .out_tx
674            .lock()
675            .expect("out_tx lock poisoned")
676            .clone()
677            .ok_or_else(|| DydxWsError::Transport("Output channel not initialized".to_string()))?;
678
679        let signal = self.signal.clone();
680        let subscriptions = subscriptions_state.clone();
681
682        let handler_task = get_runtime().spawn(async move {
683            let mut handler =
684                FeedHandler::new(cmd_rx, out_tx, raw_rx, client, signal, subscriptions);
685            handler.run().await;
686        });
687
688        Ok(ConnectionSlot {
689            cmd_tx,
690            topics: AHashMap::new(),
691            channel_counts: [0; CHANNEL_KIND_COUNT],
692            subscriptions_state,
693            handler_task: Some(handler_task),
694            connection_mode,
695        })
696    }
697
698    fn ticker_from_instrument_id(instrument_id: &InstrumentId) -> String {
699        let mut s = instrument_id.symbol.as_str().to_string();
700        if let Some(stripped) = s.strip_suffix("-PERP") {
701            s = stripped.to_string();
702        }
703        s
704    }
705
706    fn topic(channel: DydxWsChannel, id: Option<&str>) -> String {
707        match id {
708            Some(id) => format!("{}{}{}", channel.as_ref(), DYDX_WS_TOPIC_DELIMITER, id),
709            None => channel.as_ref().to_string(),
710        }
711    }
712
713    async fn subscribe_topic(
714        &self,
715        channel: ChannelKind,
716        topic: String,
717        sub_msg: DydxSubscription,
718    ) -> DydxWsResult<()> {
719        let _connect_guard = self.connect_lock.lock().await;
720
721        {
722            let mut slots = self.slots.lock().expect("slots lock poisoned");
723            if let Some(slot) = slots.iter_mut().find(|s| s.topics.contains_key(&topic)) {
724                *slot.topics.get_mut(&topic).expect("topic refcount present") += 1;
725                return Ok(());
726            }
727        }
728
729        let target_idx = loop {
730            {
731                let slots = self.slots.lock().expect("slots lock poisoned");
732                if let Some(idx) = slots.iter().position(|s| {
733                    (s.channel_counts[channel as usize] as usize) < self.per_channel_limit
734                }) {
735                    break idx;
736                }
737
738                if slots.len() >= self.max_ws_connections {
739                    return Err(DydxWsError::Subscription(format!(
740                        "Pool exhausted: {} connections x {} {:?} subscriptions",
741                        self.max_ws_connections, self.per_channel_limit, channel,
742                    )));
743                }
744            }
745
746            let new_slot = self.create_connection().await?;
747            let new_idx = {
748                let mut slots = self.slots.lock().expect("slots lock poisoned");
749                slots.push(new_slot);
750                slots.len() - 1
751            };
752            log::debug!(
753                "dYdX pool slot {new_idx} connected: url={}, channel={:?}",
754                self.url,
755                channel,
756            );
757        };
758
759        let mut slots = self.slots.lock().expect("slots lock poisoned");
760        let slot = &mut slots[target_idx];
761
762        slot.subscriptions_state.mark_subscribe(&topic);
763        slot.cmd_tx
764            .send(HandlerCommand::RegisterSubscription {
765                topic: topic.clone(),
766                subscription: sub_msg.clone(),
767            })
768            .map_err(|e| {
769                slot.subscriptions_state.mark_failure(&topic);
770                DydxWsError::Transport(format!("Slot {target_idx} unavailable: {e}"))
771            })?;
772
773        let payload = serde_json::to_string(&sub_msg)?;
774        if let Err(e) = slot.cmd_tx.send(HandlerCommand::SendText(payload)) {
775            slot.subscriptions_state.mark_failure(&topic);
776            let _ = slot.cmd_tx.send(HandlerCommand::UnregisterSubscription {
777                topic: topic.clone(),
778            });
779            return Err(DydxWsError::Transport(format!(
780                "Slot {target_idx} send failed: {e}"
781            )));
782        }
783
784        slot.topics.insert(topic, 1);
785        slot.channel_counts[channel as usize] =
786            slot.channel_counts[channel as usize].saturating_add(1);
787
788        Ok(())
789    }
790
791    async fn unsubscribe_topic(
792        &self,
793        channel: ChannelKind,
794        topic: String,
795        unsub_msg: DydxSubscription,
796    ) -> DydxWsResult<()> {
797        let mut slots = self.slots.lock().expect("slots lock poisoned");
798        let Some(slot_idx) = slots.iter().position(|s| s.topics.contains_key(&topic)) else {
799            return Ok(());
800        };
801
802        let slot = &mut slots[slot_idx];
803        let refcount = slot.topics.get_mut(&topic).expect("topic present");
804        if *refcount > 1 {
805            *refcount -= 1;
806            return Ok(());
807        }
808
809        slot.subscriptions_state.mark_unsubscribe(&topic);
810        let payload = serde_json::to_string(&unsub_msg)?;
811        if let Err(e) = slot.cmd_tx.send(HandlerCommand::SendText(payload)) {
812            slot.subscriptions_state.mark_subscribe(&topic);
813            return Err(DydxWsError::Transport(format!(
814                "Slot {slot_idx} send failed: {e}"
815            )));
816        }
817        let _ = slot.cmd_tx.send(HandlerCommand::UnregisterSubscription {
818            topic: topic.clone(),
819        });
820
821        slot.topics.remove(&topic);
822        slot.channel_counts[channel as usize] =
823            slot.channel_counts[channel as usize].saturating_sub(1);
824
825        Ok(())
826    }
827
828    /// Subscribes to public trade updates for a specific instrument.
829    ///
830    /// # Errors
831    ///
832    /// Returns an error if the subscription request fails.
833    ///
834    /// # References
835    ///
836    /// <https://docs.dydx.trade/developers/indexer/websockets#trades-channel>
837    pub async fn subscribe_trades(&self, instrument_id: InstrumentId) -> DydxWsResult<()> {
838        let ticker = Self::ticker_from_instrument_id(&instrument_id);
839        let topic = Self::topic(DydxWsChannel::Trades, Some(&ticker));
840        let sub = DydxSubscription {
841            op: DydxWsOperation::Subscribe,
842            channel: DydxWsChannel::Trades,
843            id: Some(ticker),
844        };
845        self.subscribe_topic(ChannelKind::Trades, topic, sub).await
846    }
847
848    /// Unsubscribes from public trade updates for a specific instrument.
849    ///
850    /// # Errors
851    ///
852    /// Returns an error if the unsubscription request fails.
853    pub async fn unsubscribe_trades(&self, instrument_id: InstrumentId) -> DydxWsResult<()> {
854        let ticker = Self::ticker_from_instrument_id(&instrument_id);
855        let topic = Self::topic(DydxWsChannel::Trades, Some(&ticker));
856        let sub = DydxSubscription {
857            op: DydxWsOperation::Unsubscribe,
858            channel: DydxWsChannel::Trades,
859            id: Some(ticker),
860        };
861        self.unsubscribe_topic(ChannelKind::Trades, topic, sub)
862            .await
863    }
864
865    /// Subscribes to orderbook updates for a specific instrument.
866    ///
867    /// # Errors
868    ///
869    /// Returns an error if the subscription request fails.
870    ///
871    /// # References
872    ///
873    /// <https://docs.dydx.trade/developers/indexer/websockets#orderbook-channel>
874    pub async fn subscribe_orderbook(&self, instrument_id: InstrumentId) -> DydxWsResult<()> {
875        let ticker = Self::ticker_from_instrument_id(&instrument_id);
876        let topic = Self::topic(DydxWsChannel::Orderbook, Some(&ticker));
877        let sub = DydxSubscription {
878            op: DydxWsOperation::Subscribe,
879            channel: DydxWsChannel::Orderbook,
880            id: Some(ticker),
881        };
882        self.subscribe_topic(ChannelKind::Orderbook, topic, sub)
883            .await
884    }
885
886    /// Unsubscribes from orderbook updates for a specific instrument.
887    ///
888    /// # Errors
889    ///
890    /// Returns an error if the unsubscription request fails.
891    pub async fn unsubscribe_orderbook(&self, instrument_id: InstrumentId) -> DydxWsResult<()> {
892        let ticker = Self::ticker_from_instrument_id(&instrument_id);
893        let topic = Self::topic(DydxWsChannel::Orderbook, Some(&ticker));
894        let sub = DydxSubscription {
895            op: DydxWsOperation::Unsubscribe,
896            channel: DydxWsChannel::Orderbook,
897            id: Some(ticker),
898        };
899        self.unsubscribe_topic(ChannelKind::Orderbook, topic, sub)
900            .await
901    }
902
903    /// Subscribes to candle/kline updates for a specific instrument.
904    ///
905    /// # Errors
906    ///
907    /// Returns an error if the subscription request fails.
908    ///
909    /// # References
910    ///
911    /// <https://docs.dydx.trade/developers/indexer/websockets#candles-channel>
912    pub async fn subscribe_candles(
913        &self,
914        instrument_id: InstrumentId,
915        resolution: &str,
916    ) -> DydxWsResult<()> {
917        let ticker = Self::ticker_from_instrument_id(&instrument_id);
918        let id = format!("{ticker}/{resolution}");
919        let topic = Self::topic(DydxWsChannel::Candles, Some(&id));
920        let sub = DydxSubscription {
921            op: DydxWsOperation::Subscribe,
922            channel: DydxWsChannel::Candles,
923            id: Some(id),
924        };
925        self.subscribe_topic(ChannelKind::Candles, topic, sub).await
926    }
927
928    /// Unsubscribes from candle/kline updates for a specific instrument.
929    ///
930    /// # Errors
931    ///
932    /// Returns an error if the unsubscription request fails.
933    pub async fn unsubscribe_candles(
934        &self,
935        instrument_id: InstrumentId,
936        resolution: &str,
937    ) -> DydxWsResult<()> {
938        let ticker = Self::ticker_from_instrument_id(&instrument_id);
939        let id = format!("{ticker}/{resolution}");
940        let topic = Self::topic(DydxWsChannel::Candles, Some(&id));
941        let sub = DydxSubscription {
942            op: DydxWsOperation::Unsubscribe,
943            channel: DydxWsChannel::Candles,
944            id: Some(id),
945        };
946        self.unsubscribe_topic(ChannelKind::Candles, topic, sub)
947            .await
948    }
949
950    /// Subscribes to market updates for all instruments.
951    ///
952    /// # Errors
953    ///
954    /// Returns an error if the subscription request fails.
955    ///
956    /// # References
957    ///
958    /// <https://docs.dydx.trade/developers/indexer/websockets#markets-channel>
959    pub async fn subscribe_markets(&self) -> DydxWsResult<()> {
960        let topic = Self::topic(DydxWsChannel::Markets, None);
961        let sub = DydxSubscription {
962            op: DydxWsOperation::Subscribe,
963            channel: DydxWsChannel::Markets,
964            id: None,
965        };
966        self.subscribe_topic(ChannelKind::Markets, topic, sub).await
967    }
968
969    /// Unsubscribes from market updates.
970    ///
971    /// # Errors
972    ///
973    /// Returns an error if the unsubscription request fails.
974    pub async fn unsubscribe_markets(&self) -> DydxWsResult<()> {
975        let topic = Self::topic(DydxWsChannel::Markets, None);
976        let sub = DydxSubscription {
977            op: DydxWsOperation::Unsubscribe,
978            channel: DydxWsChannel::Markets,
979            id: None,
980        };
981        self.unsubscribe_topic(ChannelKind::Markets, topic, sub)
982            .await
983    }
984
985    /// Subscribes to subaccount updates (orders, fills, positions, balances).
986    ///
987    /// This requires authentication and will only work for private WebSocket clients
988    /// created with [`Self::new_private`]. Subaccount streams stay pinned to the
989    /// primary slot: the Indexer caps them at 256 per connection, which is well
990    /// above realistic per-process usage and keeps related fill/position events
991    /// on a single in-order stream.
992    ///
993    /// # Errors
994    ///
995    /// Returns an error if the client was not created with credentials or if the
996    /// subscription request fails.
997    ///
998    /// # References
999    ///
1000    /// <https://docs.dydx.trade/developers/indexer/websockets#subaccounts-channel>
1001    pub async fn subscribe_subaccount(
1002        &self,
1003        address: &str,
1004        subaccount_number: u32,
1005    ) -> DydxWsResult<()> {
1006        if !self.requires_auth {
1007            return Err(DydxWsError::Authentication(
1008                "Subaccount subscriptions require authentication. Use new_private() to create an authenticated client".to_string(),
1009            ));
1010        }
1011        let id = format!("{address}/{subaccount_number}");
1012        let topic = Self::topic(DydxWsChannel::Subaccounts, Some(&id));
1013        let sub = DydxSubscription {
1014            op: DydxWsOperation::Subscribe,
1015            channel: DydxWsChannel::Subaccounts,
1016            id: Some(id),
1017        };
1018        self.subscribe_pinned(topic, sub).await
1019    }
1020
1021    /// Unsubscribes from subaccount updates.
1022    ///
1023    /// # Errors
1024    ///
1025    /// Returns an error if the unsubscription request fails.
1026    pub async fn unsubscribe_subaccount(
1027        &self,
1028        address: &str,
1029        subaccount_number: u32,
1030    ) -> DydxWsResult<()> {
1031        let id = format!("{address}/{subaccount_number}");
1032        let topic = Self::topic(DydxWsChannel::Subaccounts, Some(&id));
1033        let sub = DydxSubscription {
1034            op: DydxWsOperation::Unsubscribe,
1035            channel: DydxWsChannel::Subaccounts,
1036            id: Some(id),
1037        };
1038        self.unsubscribe_pinned(topic, sub).await
1039    }
1040
1041    /// Subscribes to block height updates.
1042    ///
1043    /// # Errors
1044    ///
1045    /// Returns an error if the subscription request fails.
1046    ///
1047    /// # References
1048    ///
1049    /// <https://docs.dydx.trade/developers/indexer/websockets#block-height-channel>
1050    pub async fn subscribe_block_height(&self) -> DydxWsResult<()> {
1051        let topic = Self::topic(DydxWsChannel::BlockHeight, None);
1052        let sub = DydxSubscription {
1053            op: DydxWsOperation::Subscribe,
1054            channel: DydxWsChannel::BlockHeight,
1055            id: None,
1056        };
1057        self.subscribe_pinned(topic, sub).await
1058    }
1059
1060    /// Unsubscribes from block height updates.
1061    ///
1062    /// # Errors
1063    ///
1064    /// Returns an error if the unsubscription request fails.
1065    pub async fn unsubscribe_block_height(&self) -> DydxWsResult<()> {
1066        let topic = Self::topic(DydxWsChannel::BlockHeight, None);
1067        let sub = DydxSubscription {
1068            op: DydxWsOperation::Unsubscribe,
1069            channel: DydxWsChannel::BlockHeight,
1070            id: None,
1071        };
1072        self.unsubscribe_pinned(topic, sub).await
1073    }
1074
1075    async fn subscribe_pinned(&self, topic: String, sub_msg: DydxSubscription) -> DydxWsResult<()> {
1076        let _connect_guard = self.connect_lock.lock().await;
1077
1078        {
1079            let mut slots = self.slots.lock().expect("slots lock poisoned");
1080            if let Some(slot) = slots.iter_mut().find(|s| s.topics.contains_key(&topic)) {
1081                *slot.topics.get_mut(&topic).expect("topic refcount present") += 1;
1082                return Ok(());
1083            }
1084        }
1085
1086        if self.slots.lock().expect("slots lock poisoned").is_empty() {
1087            let new_slot = self.create_connection().await?;
1088            self.connection_mode.store(new_slot.connection_mode.clone());
1089            self.slots
1090                .lock()
1091                .expect("slots lock poisoned")
1092                .push(new_slot);
1093        }
1094
1095        let mut slots = self.slots.lock().expect("slots lock poisoned");
1096        let slot = slots.first_mut().expect("primary slot exists");
1097        slot.subscriptions_state.mark_subscribe(&topic);
1098        slot.cmd_tx
1099            .send(HandlerCommand::RegisterSubscription {
1100                topic: topic.clone(),
1101                subscription: sub_msg.clone(),
1102            })
1103            .map_err(|e| {
1104                slot.subscriptions_state.mark_failure(&topic);
1105                DydxWsError::Transport(format!("Primary slot unavailable: {e}"))
1106            })?;
1107        let payload = serde_json::to_string(&sub_msg)?;
1108        if let Err(e) = slot.cmd_tx.send(HandlerCommand::SendText(payload)) {
1109            slot.subscriptions_state.mark_failure(&topic);
1110            let _ = slot.cmd_tx.send(HandlerCommand::UnregisterSubscription {
1111                topic: topic.clone(),
1112            });
1113            return Err(DydxWsError::Transport(format!(
1114                "Primary slot send failed: {e}"
1115            )));
1116        }
1117        slot.topics.insert(topic, 1);
1118        Ok(())
1119    }
1120
1121    async fn unsubscribe_pinned(
1122        &self,
1123        topic: String,
1124        unsub_msg: DydxSubscription,
1125    ) -> DydxWsResult<()> {
1126        let mut slots = self.slots.lock().expect("slots lock poisoned");
1127        let Some(slot) = slots.first_mut() else {
1128            return Ok(());
1129        };
1130        let Some(refcount) = slot.topics.get_mut(&topic) else {
1131            return Ok(());
1132        };
1133
1134        if *refcount > 1 {
1135            *refcount -= 1;
1136            return Ok(());
1137        }
1138        slot.subscriptions_state.mark_unsubscribe(&topic);
1139        let payload = serde_json::to_string(&unsub_msg)?;
1140        if let Err(e) = slot.cmd_tx.send(HandlerCommand::SendText(payload)) {
1141            slot.subscriptions_state.mark_subscribe(&topic);
1142            return Err(DydxWsError::Transport(format!(
1143                "Primary slot send failed: {e}"
1144            )));
1145        }
1146        let _ = slot.cmd_tx.send(HandlerCommand::UnregisterSubscription {
1147            topic: topic.clone(),
1148        });
1149        slot.topics.remove(&topic);
1150        Ok(())
1151    }
1152}
1153
1154// Scopes per-slot reconnect cleanup of in-progress bars to the candle topics
1155// owned by the reconnecting connection, so one slot's reconnect does not
1156// discard bars still aggregating on other healthy connections.
1157pub(crate) fn candle_ids_from_topics(topics: &[String]) -> AHashSet<String> {
1158    let prefix = format!(
1159        "{}{}",
1160        DydxWsChannel::Candles.as_ref(),
1161        DYDX_WS_TOPIC_DELIMITER
1162    );
1163    topics
1164        .iter()
1165        .filter_map(|topic| topic.strip_prefix(&prefix).map(ToString::to_string))
1166        .collect()
1167}
1168
1169#[cfg(test)]
1170mod tests {
1171    use rstest::rstest;
1172
1173    use super::*;
1174
1175    #[rstest]
1176    fn test_candle_ids_from_topics_extracts_only_candle_ids() {
1177        let topics = vec![
1178            "v4_candles:BTC-USD/1MIN".to_string(),
1179            "v4_trades:BTC-USD".to_string(),
1180            "v4_orderbook:ETH-USD".to_string(),
1181            "v4_candles:ETH-USD/5MINS".to_string(),
1182        ];
1183
1184        let ids = candle_ids_from_topics(&topics);
1185
1186        assert_eq!(ids.len(), 2);
1187        assert!(ids.contains("BTC-USD/1MIN"));
1188        assert!(ids.contains("ETH-USD/5MINS"));
1189        assert!(!ids.contains("BTC-USD"));
1190    }
1191}
nautilus_dydx/websocket/client.rs

nautilus_dydx/websocket/
client.rs
