nautilus_architect_ax/websocket/data/
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//! Market data WebSocket client for Ax.
17
18use std::{
19    fmt::Debug,
20    sync::{
21        Arc, Mutex,
22        atomic::{AtomicBool, AtomicI64, AtomicU8, Ordering},
23    },
24    time::Duration,
25};
26
27use ahash::AHashSet;
28use arc_swap::ArcSwap;
29use nautilus_common::live::get_runtime;
30use nautilus_core::{AtomicMap, consts::NAUTILUS_USER_AGENT};
31use nautilus_network::{
32    backoff::ExponentialBackoff,
33    http::USER_AGENT,
34    mode::ConnectionMode,
35    websocket::{
36        PingHandler, SubscriptionState, TransportBackend, WebSocketClient, WebSocketConfig,
37        channel_message_handler,
38    },
39};
40use ustr::Ustr;
41
42use super::handler::{AxMdWsFeedHandler, HandlerCommand};
43use crate::{
44    common::enums::{AxCandleWidth, AxMarketDataLevel},
45    websocket::messages::AxDataWsMessage,
46};
47
48/// Subscription topic delimiter for Ax.
49const AX_TOPIC_DELIMITER: char = ':';
50
51/// Result type for Ax WebSocket operations.
52pub type AxWsResult<T> = Result<T, AxWsClientError>;
53
54/// Error type for the Ax WebSocket client.
55#[derive(Debug, Clone)]
56pub enum AxWsClientError {
57    /// Transport/connection error.
58    Transport(String),
59    /// Channel send error.
60    ChannelError(String),
61}
62
63impl core::fmt::Display for AxWsClientError {
64    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
65        match self {
66            Self::Transport(msg) => write!(f, "Transport error: {msg}"),
67            Self::ChannelError(msg) => write!(f, "Channel error: {msg}"),
68        }
69    }
70}
71
72impl std::error::Error for AxWsClientError {}
73
74#[derive(Debug, Default, Clone)]
75pub struct SymbolDataTypes {
76    pub quotes: bool,
77    pub trades: bool,
78    pub mark_prices: bool,
79    pub instrument_status: bool,
80    pub book_level: Option<AxMarketDataLevel>,
81}
82
83impl SymbolDataTypes {
84    pub fn effective_level(&self) -> Option<AxMarketDataLevel> {
85        if let Some(level) = self.book_level {
86            return Some(level);
87        }
88
89        if self.quotes || self.trades || self.mark_prices || self.instrument_status {
90            return Some(AxMarketDataLevel::Level1);
91        }
92        None
93    }
94
95    fn is_empty(&self) -> bool {
96        !self.quotes
97            && !self.trades
98            && !self.mark_prices
99            && !self.instrument_status
100            && self.book_level.is_none()
101    }
102}
103
104/// Market data WebSocket client for Ax.
105///
106/// Provides streaming market data including tickers, trades, order books, and candles.
107/// Requires Bearer token authentication obtained via the HTTP `/api/authenticate` endpoint.
108pub struct AxMdWebSocketClient {
109    url: String,
110    heartbeat: Option<u64>,
111    auth_token: Option<String>,
112    connection_mode: Arc<ArcSwap<AtomicU8>>,
113    cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
114    out_rx: Option<Arc<tokio::sync::mpsc::UnboundedReceiver<AxDataWsMessage>>>,
115    signal: Arc<AtomicBool>,
116    task_handle: Option<tokio::task::JoinHandle<()>>,
117    subscriptions: SubscriptionState,
118    request_id_counter: Arc<AtomicI64>,
119    subscribe_lock: Arc<tokio::sync::Mutex<()>>,
120    symbol_data_types: Arc<AtomicMap<String, SymbolDataTypes>>,
121    status_invalidations: Arc<Mutex<AHashSet<Ustr>>>,
122    transport_backend: TransportBackend,
123    proxy_url: Option<String>,
124}
125
126impl Debug for AxMdWebSocketClient {
127    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
128        f.debug_struct(stringify!(AxMdWebSocketClient))
129            .field("url", &self.url)
130            .field("heartbeat", &self.heartbeat)
131            .field("confirmed_subscriptions", &self.subscriptions.len())
132            .finish()
133    }
134}
135
136impl Clone for AxMdWebSocketClient {
137    fn clone(&self) -> Self {
138        Self {
139            url: self.url.clone(),
140            heartbeat: self.heartbeat,
141            auth_token: self.auth_token.clone(),
142            connection_mode: Arc::clone(&self.connection_mode),
143            cmd_tx: Arc::clone(&self.cmd_tx),
144            out_rx: None,
145            signal: Arc::clone(&self.signal),
146            task_handle: None,
147            subscriptions: self.subscriptions.clone(),
148            subscribe_lock: Arc::clone(&self.subscribe_lock),
149            request_id_counter: Arc::clone(&self.request_id_counter),
150            symbol_data_types: Arc::clone(&self.symbol_data_types),
151            status_invalidations: Arc::clone(&self.status_invalidations),
152            transport_backend: self.transport_backend,
153            proxy_url: self.proxy_url.clone(),
154        }
155    }
156}
157
158impl AxMdWebSocketClient {
159    /// Creates a new Ax market data WebSocket client.
160    ///
161    /// The `auth_token` is a Bearer token obtained from the HTTP `/api/authenticate` endpoint.
162    #[must_use]
163    pub fn new(
164        url: String,
165        auth_token: String,
166        heartbeat: u64,
167        transport_backend: TransportBackend,
168        proxy_url: Option<String>,
169    ) -> Self {
170        let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
171
172        let initial_mode = AtomicU8::new(ConnectionMode::Closed.as_u8());
173        let connection_mode = Arc::new(ArcSwap::from_pointee(initial_mode));
174
175        Self {
176            url,
177            heartbeat: Some(heartbeat),
178            auth_token: Some(auth_token),
179            connection_mode,
180            cmd_tx: Arc::new(tokio::sync::RwLock::new(cmd_tx)),
181            out_rx: None,
182            signal: Arc::new(AtomicBool::new(false)),
183            task_handle: None,
184            subscriptions: SubscriptionState::new(AX_TOPIC_DELIMITER),
185            request_id_counter: Arc::new(AtomicI64::new(1)),
186            subscribe_lock: Arc::new(tokio::sync::Mutex::new(())),
187            symbol_data_types: Arc::new(AtomicMap::new()),
188            status_invalidations: Arc::new(Mutex::new(AHashSet::new())),
189            transport_backend,
190            proxy_url,
191        }
192    }
193
194    /// Creates a new Ax market data WebSocket client without authentication.
195    ///
196    /// Use [`set_auth_token`](Self::set_auth_token) to set the token before connecting.
197    #[must_use]
198    pub fn without_auth(
199        url: String,
200        heartbeat: u64,
201        transport_backend: TransportBackend,
202        proxy_url: Option<String>,
203    ) -> Self {
204        let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
205
206        let initial_mode = AtomicU8::new(ConnectionMode::Closed.as_u8());
207        let connection_mode = Arc::new(ArcSwap::from_pointee(initial_mode));
208
209        Self {
210            url,
211            heartbeat: Some(heartbeat),
212            auth_token: None,
213            connection_mode,
214            cmd_tx: Arc::new(tokio::sync::RwLock::new(cmd_tx)),
215            out_rx: None,
216            signal: Arc::new(AtomicBool::new(false)),
217            task_handle: None,
218            subscriptions: SubscriptionState::new(AX_TOPIC_DELIMITER),
219            request_id_counter: Arc::new(AtomicI64::new(1)),
220            subscribe_lock: Arc::new(tokio::sync::Mutex::new(())),
221            symbol_data_types: Arc::new(AtomicMap::new()),
222            status_invalidations: Arc::new(Mutex::new(AHashSet::new())),
223            transport_backend,
224            proxy_url,
225        }
226    }
227
228    /// Returns the WebSocket URL.
229    #[must_use]
230    pub fn url(&self) -> &str {
231        &self.url
232    }
233
234    /// Sets the authentication token for subsequent connections.
235    ///
236    /// This should be called before `connect()` if authentication is required.
237    pub fn set_auth_token(&mut self, token: String) {
238        self.auth_token = Some(token);
239    }
240
241    /// Returns whether the client is currently connected and active.
242    #[must_use]
243    pub fn is_active(&self) -> bool {
244        let connection_mode_arc = self.connection_mode.load();
245        ConnectionMode::from_atomic(&connection_mode_arc).is_active()
246            && !self.signal.load(Ordering::Acquire)
247    }
248
249    /// Returns whether the client is closed.
250    #[must_use]
251    pub fn is_closed(&self) -> bool {
252        let connection_mode_arc = self.connection_mode.load();
253        ConnectionMode::from_atomic(&connection_mode_arc).is_closed()
254            || self.signal.load(Ordering::Acquire)
255    }
256
257    /// Returns the number of confirmed subscriptions.
258    #[must_use]
259    pub fn subscription_count(&self) -> usize {
260        self.subscriptions.len()
261    }
262
263    /// Returns the symbol data types map (shared with handler).
264    #[must_use]
265    pub fn symbol_data_types(&self) -> Arc<AtomicMap<String, SymbolDataTypes>> {
266        Arc::clone(&self.symbol_data_types)
267    }
268
269    /// Returns the shared set of symbols whose instrument status cache has been invalidated.
270    pub fn status_invalidations(&self) -> Arc<Mutex<AHashSet<Ustr>>> {
271        Arc::clone(&self.status_invalidations)
272    }
273
274    fn next_request_id(&self) -> i64 {
275        self.request_id_counter.fetch_add(1, Ordering::Relaxed)
276    }
277
278    fn is_subscribed_topic(&self, topic: &str) -> bool {
279        let (channel, symbol) = topic
280            .split_once(AX_TOPIC_DELIMITER)
281            .map_or((topic, None), |(c, s)| (c, Some(s)));
282        let channel_ustr = Ustr::from(channel);
283        let symbol_ustr = symbol.map_or_else(|| Ustr::from(""), Ustr::from);
284        self.subscriptions
285            .is_subscribed(&channel_ustr, &symbol_ustr)
286    }
287
288    /// Establishes the WebSocket connection.
289    ///
290    /// # Errors
291    ///
292    pub async fn connect(&mut self) -> AxWsResult<()> {
293        const MAX_RETRIES: u32 = 5;
294        const CONNECTION_TIMEOUT_SECS: u64 = 10;
295
296        self.signal.store(false, Ordering::Release);
297
298        let (raw_handler, raw_rx) = channel_message_handler();
299
300        // No-op: ping responses are handled internally by the WebSocketClient
301        let ping_handler: PingHandler = Arc::new(move |_payload: Vec<u8>| {});
302
303        let mut headers = vec![(USER_AGENT.to_string(), NAUTILUS_USER_AGENT.to_string())];
304
305        if let Some(ref token) = self.auth_token {
306            headers.push(("Authorization".to_string(), format!("Bearer {token}")));
307        }
308
309        let config = WebSocketConfig {
310            url: self.url.clone(),
311            headers,
312            heartbeat: self.heartbeat,
313            heartbeat_msg: None, // Ax server sends heartbeats
314            reconnect_timeout_ms: Some(5_000),
315            reconnect_delay_initial_ms: Some(500),
316            reconnect_delay_max_ms: Some(5_000),
317            reconnect_backoff_factor: Some(1.5),
318            reconnect_jitter_ms: Some(250),
319            reconnect_max_attempts: None,
320            idle_timeout_ms: None,
321            backend: self.transport_backend,
322            proxy_url: self.proxy_url.clone(),
323        };
324
325        // Retry initial connection with exponential backoff
326        let mut backoff = ExponentialBackoff::new(
327            Duration::from_millis(500),
328            Duration::from_millis(5000),
329            2.0,
330            250,
331            false,
332        )
333        .map_err(|e| AxWsClientError::Transport(e.to_string()))?;
334
335        let mut last_error: String;
336        let mut attempt = 0;
337
338        let client = loop {
339            attempt += 1;
340
341            match tokio::time::timeout(
342                Duration::from_secs(CONNECTION_TIMEOUT_SECS),
343                WebSocketClient::connect(
344                    config.clone(),
345                    Some(raw_handler.clone()),
346                    Some(ping_handler.clone()),
347                    None,
348                    vec![],
349                    None,
350                ),
351            )
352            .await
353            {
354                Ok(Ok(client)) => {
355                    if attempt > 1 {
356                        log::debug!("WebSocket connection established after {attempt} attempts");
357                    }
358                    break client;
359                }
360                Ok(Err(e)) => {
361                    last_error = e.to_string();
362                    log::warn!(
363                        "WebSocket connection attempt failed: attempt={attempt}/{MAX_RETRIES}, url={}, error={last_error}",
364                        self.url
365                    );
366                }
367                Err(_) => {
368                    last_error = format!("Connection timeout after {CONNECTION_TIMEOUT_SECS}s");
369                    log::warn!(
370                        "WebSocket connection attempt timed out: attempt={attempt}/{MAX_RETRIES}, url={}",
371                        self.url
372                    );
373                }
374            }
375
376            if attempt >= MAX_RETRIES {
377                return Err(AxWsClientError::Transport(format!(
378                    "Failed to connect to {} after {MAX_RETRIES} attempts: {}",
379                    self.url,
380                    if last_error.is_empty() {
381                        "unknown error"
382                    } else {
383                        &last_error
384                    }
385                )));
386            }
387
388            let delay = backoff.next_duration();
389            log::debug!(
390                "Retrying in {delay:?} (attempt {}/{MAX_RETRIES})",
391                attempt + 1
392            );
393            tokio::time::sleep(delay).await;
394        };
395
396        self.connection_mode.store(client.connection_mode_atomic());
397
398        let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<AxDataWsMessage>();
399        self.out_rx = Some(Arc::new(out_rx));
400
401        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
402        *self.cmd_tx.write().await = cmd_tx.clone();
403
404        self.send_cmd(HandlerCommand::SetClient(client)).await?;
405
406        let signal = Arc::clone(&self.signal);
407        let subscriptions = self.subscriptions.clone();
408
409        let stream_handle = get_runtime().spawn(async move {
410            let mut handler =
411                AxMdWsFeedHandler::new(signal.clone(), cmd_rx, raw_rx, subscriptions.clone());
412
413            while let Some(msg) = handler.next().await {
414                if matches!(msg, AxDataWsMessage::Reconnected) {
415                    log::info!("WebSocket reconnected, subscriptions will be replayed");
416                }
417
418                if out_tx.send(msg).is_err() {
419                    log::debug!("Output channel closed");
420                    break;
421                }
422            }
423
424            log::debug!("Handler loop exited");
425        });
426
427        self.task_handle = Some(stream_handle);
428
429        Ok(())
430    }
431
432    /// Subscribes to order book deltas for a symbol.
433    ///
434    /// Uses reference counting so the underlying AX subscription is only
435    /// removed when all data types have been unsubscribed.
436    ///
437    /// # Errors
438    ///
439    /// Returns an error if the subscription command cannot be sent.
440    pub async fn subscribe_book_deltas(
441        &self,
442        symbol: &str,
443        level: AxMarketDataLevel,
444    ) -> AxWsResult<()> {
445        let _guard = self.subscribe_lock.lock().await;
446
447        let current = self
448            .symbol_data_types
449            .load()
450            .get(symbol)
451            .cloned()
452            .unwrap_or_default();
453
454        // AX allows only one subscription per symbol, skip if book already subscribed
455        if current.book_level.is_some() {
456            log::debug!("Book deltas already subscribed for {symbol}, skipping");
457            return Ok(());
458        }
459
460        let old_level = current.effective_level();
461        let mut next = current.clone();
462        next.book_level = Some(level);
463        let new_level = next.effective_level();
464
465        self.update_data_subscription(symbol, old_level, new_level)
466            .await?;
467
468        self.symbol_data_types.rcu(|m| {
469            let entry = m.entry(symbol.to_string()).or_default();
470            entry.book_level = Some(level);
471        });
472
473        Ok(())
474    }
475
476    /// Subscribes to quote data for a symbol.
477    ///
478    /// Uses reference counting so the underlying AX subscription is only
479    /// removed when all data types have been unsubscribed.
480    ///
481    /// # Errors
482    ///
483    /// Returns an error if the subscription command cannot be sent.
484    pub async fn subscribe_quotes(&self, symbol: &str) -> AxWsResult<()> {
485        let _guard = self.subscribe_lock.lock().await;
486
487        let current = self
488            .symbol_data_types
489            .load()
490            .get(symbol)
491            .cloned()
492            .unwrap_or_default();
493        let old_level = current.effective_level();
494        let mut next = current.clone();
495        next.quotes = true;
496        let new_level = next.effective_level();
497
498        self.update_data_subscription(symbol, old_level, new_level)
499            .await?;
500
501        self.symbol_data_types.rcu(|m| {
502            m.entry(symbol.to_string()).or_default().quotes = true;
503        });
504
505        Ok(())
506    }
507
508    /// Subscribes to trade data for a symbol.
509    ///
510    /// Uses reference counting so the underlying AX subscription is only
511    /// removed when all data types have been unsubscribed.
512    ///
513    /// # Errors
514    ///
515    /// Returns an error if the subscription command cannot be sent.
516    pub async fn subscribe_trades(&self, symbol: &str) -> AxWsResult<()> {
517        let _guard = self.subscribe_lock.lock().await;
518
519        let current = self
520            .symbol_data_types
521            .load()
522            .get(symbol)
523            .cloned()
524            .unwrap_or_default();
525        let old_level = current.effective_level();
526        let mut next = current.clone();
527        next.trades = true;
528        let new_level = next.effective_level();
529
530        self.update_data_subscription(symbol, old_level, new_level)
531            .await?;
532
533        self.symbol_data_types.rcu(|m| {
534            m.entry(symbol.to_string()).or_default().trades = true;
535        });
536
537        Ok(())
538    }
539
540    /// Unsubscribes from order book deltas for a symbol.
541    ///
542    /// The underlying AX subscription is only removed when all data types
543    /// (quotes, trades, book) have been unsubscribed.
544    ///
545    /// # Errors
546    ///
547    /// Returns an error if the unsubscribe command cannot be sent.
548    pub async fn unsubscribe_book_deltas(&self, symbol: &str) -> AxWsResult<()> {
549        let _guard = self.subscribe_lock.lock().await;
550
551        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
552            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe book deltas");
553            return Ok(());
554        };
555        let old_level = current.effective_level();
556        let mut next = current.clone();
557        next.book_level = None;
558        let new_level = next.effective_level();
559
560        self.update_data_subscription(symbol, old_level, new_level)
561            .await?;
562
563        self.symbol_data_types.rcu(|m| {
564            if let Some(entry) = m.get_mut(symbol) {
565                entry.book_level = None;
566                if entry.is_empty() {
567                    m.remove(symbol);
568                }
569            }
570        });
571
572        Ok(())
573    }
574
575    /// Unsubscribes from quote data for a symbol.
576    ///
577    /// The underlying AX subscription is only removed when all data types
578    /// (quotes, trades, book) have been unsubscribed.
579    ///
580    /// # Errors
581    ///
582    /// Returns an error if the unsubscribe command cannot be sent.
583    pub async fn unsubscribe_quotes(&self, symbol: &str) -> AxWsResult<()> {
584        let _guard = self.subscribe_lock.lock().await;
585
586        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
587            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe quotes");
588            return Ok(());
589        };
590        let old_level = current.effective_level();
591        let mut next = current.clone();
592        next.quotes = false;
593        let new_level = next.effective_level();
594
595        self.update_data_subscription(symbol, old_level, new_level)
596            .await?;
597
598        self.symbol_data_types.rcu(|m| {
599            if let Some(entry) = m.get_mut(symbol) {
600                entry.quotes = false;
601                if entry.is_empty() {
602                    m.remove(symbol);
603                }
604            }
605        });
606
607        Ok(())
608    }
609
610    /// Unsubscribes from trade data for a symbol.
611    ///
612    /// The underlying AX subscription is only removed when all data types
613    /// (quotes, trades, book) have been unsubscribed.
614    ///
615    /// # Errors
616    ///
617    /// Returns an error if the unsubscribe command cannot be sent.
618    pub async fn unsubscribe_trades(&self, symbol: &str) -> AxWsResult<()> {
619        let _guard = self.subscribe_lock.lock().await;
620
621        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
622            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe trades");
623            return Ok(());
624        };
625        let old_level = current.effective_level();
626        let mut next = current.clone();
627        next.trades = false;
628        let new_level = next.effective_level();
629
630        self.update_data_subscription(symbol, old_level, new_level)
631            .await?;
632
633        self.symbol_data_types.rcu(|m| {
634            if let Some(entry) = m.get_mut(symbol) {
635                entry.trades = false;
636                if entry.is_empty() {
637                    m.remove(symbol);
638                }
639            }
640        });
641
642        Ok(())
643    }
644
645    /// Subscribes to mark prices for a symbol.
646    ///
647    /// Ensures at least an L1 subscription so that ticker messages
648    /// (which carry the mark price field) are received.
649    ///
650    /// # Errors
651    ///
652    /// Returns an error if the subscription command cannot be sent.
653    pub async fn subscribe_mark_prices(&self, symbol: &str) -> AxWsResult<()> {
654        let _guard = self.subscribe_lock.lock().await;
655
656        let current = self
657            .symbol_data_types
658            .load()
659            .get(symbol)
660            .cloned()
661            .unwrap_or_default();
662        let old_level = current.effective_level();
663        let mut next = current.clone();
664        next.mark_prices = true;
665        let new_level = next.effective_level();
666
667        self.update_data_subscription(symbol, old_level, new_level)
668            .await?;
669
670        self.symbol_data_types.rcu(|m| {
671            m.entry(symbol.to_string()).or_default().mark_prices = true;
672        });
673
674        Ok(())
675    }
676
677    /// Unsubscribes from mark prices for a symbol.
678    ///
679    /// The underlying AX subscription is only removed when all data types
680    /// have been unsubscribed.
681    ///
682    /// # Errors
683    ///
684    /// Returns an error if the unsubscribe command cannot be sent.
685    pub async fn unsubscribe_mark_prices(&self, symbol: &str) -> AxWsResult<()> {
686        let _guard = self.subscribe_lock.lock().await;
687
688        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
689            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe mark prices");
690            return Ok(());
691        };
692        let old_level = current.effective_level();
693        let mut next = current.clone();
694        next.mark_prices = false;
695        let new_level = next.effective_level();
696
697        self.update_data_subscription(symbol, old_level, new_level)
698            .await?;
699
700        self.symbol_data_types.rcu(|m| {
701            if let Some(entry) = m.get_mut(symbol) {
702                entry.mark_prices = false;
703                if entry.is_empty() {
704                    m.remove(symbol);
705                }
706            }
707        });
708
709        Ok(())
710    }
711
712    /// Subscribes to instrument status for a symbol.
713    ///
714    /// Ensures at least an L1 subscription so that ticker messages
715    /// (which carry the instrument state field) are received.
716    ///
717    /// # Errors
718    ///
719    /// Returns an error if the subscription command cannot be sent.
720    pub async fn subscribe_instrument_status(&self, symbol: &str) -> AxWsResult<()> {
721        let _guard = self.subscribe_lock.lock().await;
722
723        let current = self
724            .symbol_data_types
725            .load()
726            .get(symbol)
727            .cloned()
728            .unwrap_or_default();
729        let old_level = current.effective_level();
730        let mut next = current.clone();
731        next.instrument_status = true;
732        let new_level = next.effective_level();
733
734        self.update_data_subscription(symbol, old_level, new_level)
735            .await?;
736
737        self.symbol_data_types.rcu(|m| {
738            m.entry(symbol.to_string()).or_default().instrument_status = true;
739        });
740
741        Ok(())
742    }
743
744    /// Unsubscribes from instrument status for a symbol.
745    ///
746    /// The underlying AX subscription is only removed when all data types
747    /// have been unsubscribed.
748    ///
749    /// # Errors
750    ///
751    /// Returns an error if the unsubscribe command cannot be sent.
752    pub async fn unsubscribe_instrument_status(&self, symbol: &str) -> AxWsResult<()> {
753        let _guard = self.subscribe_lock.lock().await;
754
755        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
756            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe instrument status");
757            return Ok(());
758        };
759        let old_level = current.effective_level();
760        let mut next = current.clone();
761        next.instrument_status = false;
762        let new_level = next.effective_level();
763
764        self.update_data_subscription(symbol, old_level, new_level)
765            .await?;
766
767        self.symbol_data_types.rcu(|m| {
768            if let Some(entry) = m.get_mut(symbol) {
769                entry.instrument_status = false;
770                if entry.is_empty() {
771                    m.remove(symbol);
772                }
773            }
774        });
775
776        if let Ok(mut invalidations) = self.status_invalidations.lock() {
777            invalidations.insert(Ustr::from(symbol));
778        }
779
780        Ok(())
781    }
782
783    async fn update_data_subscription(
784        &self,
785        symbol: &str,
786        old_level: Option<AxMarketDataLevel>,
787        new_level: Option<AxMarketDataLevel>,
788    ) -> AxWsResult<()> {
789        if old_level == new_level {
790            return Ok(());
791        }
792
793        match (old_level, new_level) {
794            (None, Some(level)) => {
795                log::debug!("Subscribing {symbol} at {level:?}");
796                self.send_subscribe(symbol, level).await
797            }
798            (Some(_), None) => {
799                log::debug!("Unsubscribing {symbol} (no remaining data types)");
800                self.send_unsubscribe(symbol).await
801            }
802            (Some(old), Some(new)) => {
803                log::debug!("Resubscribing {symbol}: {old:?} -> {new:?}");
804                self.send_unsubscribe(symbol).await?;
805                if let Err(e) = self.send_subscribe(symbol, new).await {
806                    log::warn!("Resubscribe failed for {symbol} at {new:?}: {e}");
807                    if let Err(restore_err) = self.send_subscribe(symbol, old).await {
808                        // Channel dead, mark old topic for reconnection replay
809                        log::error!(
810                            "Failed to restore {symbol} at {old:?}: {restore_err}, \
811                             reconnection required"
812                        );
813                        let old_topic = format!("{symbol}:{old:?}");
814                        self.subscriptions.mark_subscribe(&old_topic);
815                    }
816                    return Err(e);
817                }
818                Ok(())
819            }
820            (None, None) => Ok(()),
821        }
822    }
823
824    async fn send_subscribe(&self, symbol: &str, level: AxMarketDataLevel) -> AxWsResult<()> {
825        let topic = format!("{symbol}:{level:?}");
826        let request_id = self.next_request_id();
827
828        self.subscriptions.mark_subscribe(&topic);
829
830        if let Err(e) = self
831            .send_cmd(HandlerCommand::Subscribe {
832                request_id,
833                symbol: Ustr::from(symbol),
834                level,
835            })
836            .await
837        {
838            self.subscriptions.mark_unsubscribe(&topic);
839            return Err(e);
840        }
841
842        Ok(())
843    }
844
845    async fn send_unsubscribe(&self, symbol: &str) -> AxWsResult<()> {
846        let request_id = self.next_request_id();
847
848        self.send_cmd(HandlerCommand::Unsubscribe {
849            request_id,
850            symbol: Ustr::from(symbol),
851        })
852        .await?;
853
854        for level in [
855            AxMarketDataLevel::Level1,
856            AxMarketDataLevel::Level2,
857            AxMarketDataLevel::Level3,
858        ] {
859            let topic = format!("{symbol}:{level:?}");
860            self.subscriptions.mark_unsubscribe(&topic);
861        }
862
863        Ok(())
864    }
865
866    /// Subscribes to candle data for a symbol.
867    ///
868    /// Skips sending if already subscribed or subscription is pending.
869    ///
870    /// # Errors
871    ///
872    /// Returns an error if the subscription command cannot be sent.
873    pub async fn subscribe_candles(&self, symbol: &str, width: AxCandleWidth) -> AxWsResult<()> {
874        let _guard = self.subscribe_lock.lock().await;
875        let topic = format!("candles:{symbol}:{width:?}");
876
877        // Skip if already subscribed or pending
878        if self.is_subscribed_topic(&topic) {
879            log::debug!("Already subscribed to {topic}, skipping");
880            return Ok(());
881        }
882
883        let request_id = self.next_request_id();
884
885        // Mark pending BEFORE sending to prevent race conditions with concurrent subscribes
886        self.subscriptions.mark_subscribe(&topic);
887
888        if let Err(e) = self
889            .send_cmd(HandlerCommand::SubscribeCandles {
890                request_id,
891                symbol: Ustr::from(symbol),
892                width,
893            })
894            .await
895        {
896            // Rollback pending state on send failure
897            self.subscriptions.mark_unsubscribe(&topic);
898            return Err(e);
899        }
900
901        Ok(())
902    }
903
904    /// Unsubscribes from candle data for a symbol.
905    ///
906    /// # Errors
907    ///
908    /// Returns an error if the unsubscribe command cannot be sent.
909    pub async fn unsubscribe_candles(&self, symbol: &str, width: AxCandleWidth) -> AxWsResult<()> {
910        let _guard = self.subscribe_lock.lock().await;
911        let request_id = self.next_request_id();
912        let topic = format!("candles:{symbol}:{width:?}");
913
914        self.subscriptions.mark_unsubscribe(&topic);
915
916        self.send_cmd(HandlerCommand::UnsubscribeCandles {
917            request_id,
918            symbol: Ustr::from(symbol),
919            width,
920        })
921        .await
922    }
923
924    /// Returns a stream of WebSocket messages.
925    ///
926    /// # Panics
927    ///
928    /// Panics if called before `connect()` or if the stream has already been taken.
929    pub fn stream(&mut self) -> impl futures_util::Stream<Item = AxDataWsMessage> + 'static {
930        let rx = self
931            .out_rx
932            .take()
933            .expect("Stream receiver already taken or client not connected - stream() can only be called once");
934        let mut rx = Arc::try_unwrap(rx).expect(
935            "Cannot take ownership of stream - client was cloned and other references exist",
936        );
937        async_stream::stream! {
938            while let Some(msg) = rx.recv().await {
939                yield msg;
940            }
941        }
942    }
943
944    /// Disconnects the WebSocket connection gracefully.
945    pub async fn disconnect(&self) {
946        log::debug!("Disconnecting WebSocket");
947        let _ = self.send_cmd(HandlerCommand::Disconnect).await;
948    }
949
950    /// Closes the WebSocket connection and cleans up resources.
951    pub async fn close(&mut self) {
952        log::debug!("Closing WebSocket client");
953
954        // Send disconnect first to allow graceful cleanup before signal
955        let _ = self.send_cmd(HandlerCommand::Disconnect).await;
956        tokio::time::sleep(Duration::from_millis(50)).await;
957        self.signal.store(true, Ordering::Release);
958
959        if let Some(handle) = self.task_handle.take() {
960            const CLOSE_TIMEOUT: Duration = Duration::from_secs(2);
961            let abort_handle = handle.abort_handle();
962
963            match tokio::time::timeout(CLOSE_TIMEOUT, handle).await {
964                Ok(Ok(())) => log::debug!("Handler task completed gracefully"),
965                Ok(Err(e)) => log::warn!("Handler task panicked: {e}"),
966                Err(_) => {
967                    log::warn!("Handler task did not complete within timeout, aborting");
968                    abort_handle.abort();
969                }
970            }
971        }
972    }
973
974    async fn send_cmd(&self, cmd: HandlerCommand) -> AxWsResult<()> {
975        let guard = self.cmd_tx.read().await;
976        guard
977            .send(cmd)
978            .map_err(|e| AxWsClientError::ChannelError(e.to_string()))
979    }
980}
981
982#[cfg(test)]
983mod tests {
984    use rstest::rstest;
985
986    use super::*;
987
988    #[rstest]
989    fn test_effective_level_empty_returns_none() {
990        let sdt = SymbolDataTypes::default();
991        assert_eq!(sdt.effective_level(), None);
992        assert!(sdt.is_empty());
993    }
994
995    #[rstest]
996    fn test_effective_level_book_level_takes_precedence() {
997        let sdt = SymbolDataTypes {
998            book_level: Some(AxMarketDataLevel::Level2),
999            quotes: true,
1000            ..Default::default()
1001        };
1002        assert_eq!(sdt.effective_level(), Some(AxMarketDataLevel::Level2));
1003        assert!(!sdt.is_empty());
1004    }
1005
1006    #[rstest]
1007    #[case(true, false, false, false)]
1008    #[case(false, true, false, false)]
1009    #[case(false, false, true, false)]
1010    #[case(false, false, false, true)]
1011    fn test_effective_level_any_flag_returns_level1(
1012        #[case] quotes: bool,
1013        #[case] trades: bool,
1014        #[case] mark_prices: bool,
1015        #[case] instrument_status: bool,
1016    ) {
1017        let sdt = SymbolDataTypes {
1018            quotes,
1019            trades,
1020            mark_prices,
1021            instrument_status,
1022            book_level: None,
1023        };
1024        assert_eq!(sdt.effective_level(), Some(AxMarketDataLevel::Level1));
1025        assert!(!sdt.is_empty());
1026    }
1027}
nautilus_architect_ax/websocket/data/client.rs

nautilus_architect_ax/websocket/data/
client.rs
