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, Display},
20    num::NonZeroU32,
21    sync::{
22        Arc,
23        atomic::{AtomicBool, AtomicI64, AtomicU8, Ordering},
24    },
25    time::Duration,
26};
27
28use ahash::AHashSet;
29use arc_swap::ArcSwap;
30use nautilus_core::{AtomicMap, string::secret::SecretString};
31use nautilus_live::{
32    SocketControl,
33    task::{SharedTaskSlot, TaskJoinOutcome},
34};
35use nautilus_network::{
36    http::create_standard_nautilus_headers,
37    mode::ConnectionMode,
38    websocket::{
39        InitialConnectRetryPolicy, PingHandler, ReconnectHeaders, SubscriptionState,
40        TransportBackend, WebSocketClient, WebSocketConfig, channel_message_handler,
41    },
42};
43use parking_lot::Mutex;
44use tokio_util::sync::CancellationToken;
45use ustr::Ustr;
46
47use super::{
48    AxMdSubscriptionSpec,
49    handler::{AxMdWsFeedHandler, HandlerCommand},
50};
51use crate::{
52    common::enums::{AxCandleWidth, AxMarketDataLevel},
53    websocket::messages::AxDataWsMessage,
54};
55
56/// Subscription topic delimiter for Ax.
57const AX_TOPIC_DELIMITER: char = ':';
58
59/// Result type for Ax WebSocket operations.
60pub type AxWsResult<T> = Result<T, AxWsClientError>;
61
62/// Error type for the Ax WebSocket client.
63#[derive(Debug, Clone)]
64pub enum AxWsClientError {
65    /// Transport/connection error.
66    Transport(String),
67    /// Channel send error.
68    ChannelError(String),
69}
70
71impl Display for AxWsClientError {
72    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
73        match self {
74            Self::Transport(msg) => write!(f, "Transport error: {msg}"),
75            Self::ChannelError(msg) => write!(f, "Channel error: {msg}"),
76        }
77    }
78}
79
80impl std::error::Error for AxWsClientError {}
81
82#[derive(Debug, Default, Clone)]
83pub struct SymbolDataTypes {
84    pub quotes: bool,
85    pub trades: bool,
86    pub mark_prices: bool,
87    pub instrument_status: bool,
88    pub book_level: Option<AxMarketDataLevel>,
89}
90
91impl SymbolDataTypes {
92    fn effective_subscription(&self) -> Option<AxMdSubscriptionSpec> {
93        let ticker = self.mark_prices || self.instrument_status;
94        let book_level = self.book_level.or({
95            if self.quotes || ticker {
96                Some(AxMarketDataLevel::Level1)
97            } else {
98                None
99            }
100        });
101
102        if let Some(level) = book_level {
103            return Some(AxMdSubscriptionSpec::new(
104                level,
105                Some(self.trades),
106                Some(ticker),
107            ));
108        }
109
110        if self.trades {
111            return Some(AxMdSubscriptionSpec::new(
112                AxMarketDataLevel::Trades,
113                None,
114                None,
115            ));
116        }
117
118        None
119    }
120
121    fn is_empty(&self) -> bool {
122        !self.quotes
123            && !self.trades
124            && !self.mark_prices
125            && !self.instrument_status
126            && self.book_level.is_none()
127    }
128}
129
130/// Market data WebSocket client for Ax.
131///
132/// Provides streaming market data including tickers, trades, order books, and candles.
133/// Requires Bearer token authentication obtained via the HTTP `/api/authenticate` endpoint.
134pub struct AxMdWebSocketClient {
135    url: String,
136    heartbeat: Option<u64>,
137    auth_token: Arc<Mutex<Option<SecretString>>>,
138    reconnect_headers: Arc<Mutex<Option<ReconnectHeaders>>>,
139    connection_mode: Arc<ArcSwap<AtomicU8>>,
140    cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
141    out_rx: Option<Arc<tokio::sync::mpsc::UnboundedReceiver<AxDataWsMessage>>>,
142    signal: Arc<AtomicBool>,
143    cancellation_token: Arc<ArcSwap<CancellationToken>>,
144    task_handle: Arc<SharedTaskSlot<()>>,
145    connect_lock: Arc<tokio::sync::Mutex<()>>,
146    subscriptions: SubscriptionState,
147    request_id_counter: Arc<AtomicI64>,
148    subscribe_lock: Arc<tokio::sync::Mutex<()>>,
149    symbol_data_types: Arc<AtomicMap<String, SymbolDataTypes>>,
150    status_invalidations: Arc<Mutex<AHashSet<Ustr>>>,
151    transport_backend: TransportBackend,
152    proxy_url: Option<SecretString>,
153    socket_control: Option<SocketControl>,
154}
155
156impl Debug for AxMdWebSocketClient {
157    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
158        f.debug_struct(stringify!(AxMdWebSocketClient))
159            .field("url", &self.url)
160            .field("heartbeat", &self.heartbeat)
161            .field("confirmed_subscriptions", &self.subscriptions.len())
162            .finish()
163    }
164}
165
166impl Clone for AxMdWebSocketClient {
167    fn clone(&self) -> Self {
168        Self {
169            url: self.url.clone(),
170            heartbeat: self.heartbeat,
171            auth_token: Arc::clone(&self.auth_token),
172            reconnect_headers: Arc::clone(&self.reconnect_headers),
173            connection_mode: Arc::clone(&self.connection_mode),
174            cmd_tx: Arc::clone(&self.cmd_tx),
175            out_rx: None,
176            signal: Arc::clone(&self.signal),
177            cancellation_token: Arc::clone(&self.cancellation_token),
178            task_handle: Arc::clone(&self.task_handle),
179            connect_lock: Arc::clone(&self.connect_lock),
180            subscriptions: self.subscriptions.clone(),
181            subscribe_lock: Arc::clone(&self.subscribe_lock),
182            request_id_counter: Arc::clone(&self.request_id_counter),
183            symbol_data_types: Arc::clone(&self.symbol_data_types),
184            status_invalidations: Arc::clone(&self.status_invalidations),
185            transport_backend: self.transport_backend,
186            proxy_url: self.proxy_url.clone(),
187            socket_control: self.socket_control.clone(),
188        }
189    }
190}
191
192impl AxMdWebSocketClient {
193    fn initial_connect_retry_policy() -> InitialConnectRetryPolicy {
194        InitialConnectRetryPolicy {
195            max_attempts: NonZeroU32::new(5).expect("initial connect attempts must be non-zero"),
196            delay_initial: Duration::from_millis(500),
197            delay_max: Duration::from_secs(5),
198            backoff_factor: 2.0,
199            jitter_ms: 250,
200        }
201    }
202
203    /// Creates a new Ax market data WebSocket client.
204    ///
205    /// The `auth_token` is a Bearer token obtained from the HTTP `/api/authenticate` endpoint.
206    #[must_use]
207    pub fn new(
208        url: String,
209        auth_token: impl Into<SecretString>,
210        heartbeat: u64,
211        transport_backend: TransportBackend,
212        proxy_url: Option<String>,
213    ) -> Self {
214        let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
215
216        let initial_mode = AtomicU8::new(ConnectionMode::Closed.as_u8());
217        let connection_mode = Arc::new(ArcSwap::from_pointee(initial_mode));
218
219        Self {
220            url,
221            heartbeat: Some(heartbeat),
222            auth_token: Arc::new(Mutex::new(Some(auth_token.into()))),
223            reconnect_headers: Arc::new(Mutex::new(None)),
224            connection_mode,
225            cmd_tx: Arc::new(tokio::sync::RwLock::new(cmd_tx)),
226            out_rx: None,
227            signal: Arc::new(AtomicBool::new(false)),
228            cancellation_token: Arc::new(ArcSwap::from_pointee(CancellationToken::new())),
229            task_handle: Arc::new(SharedTaskSlot::new()),
230            connect_lock: Arc::new(tokio::sync::Mutex::new(())),
231            subscriptions: SubscriptionState::new(AX_TOPIC_DELIMITER),
232            request_id_counter: Arc::new(AtomicI64::new(1)),
233            subscribe_lock: Arc::new(tokio::sync::Mutex::new(())),
234            symbol_data_types: Arc::new(AtomicMap::new()),
235            status_invalidations: Arc::new(Mutex::new(AHashSet::new())),
236            transport_backend,
237            proxy_url: proxy_url.map(SecretString::from),
238            socket_control: None,
239        }
240    }
241
242    /// Creates a new Ax market data WebSocket client without authentication.
243    ///
244    /// Use [`set_auth_token`](Self::set_auth_token) to set the token before connecting.
245    #[must_use]
246    pub fn without_auth(
247        url: String,
248        heartbeat: u64,
249        transport_backend: TransportBackend,
250        proxy_url: Option<String>,
251    ) -> Self {
252        let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
253
254        let initial_mode = AtomicU8::new(ConnectionMode::Closed.as_u8());
255        let connection_mode = Arc::new(ArcSwap::from_pointee(initial_mode));
256
257        Self {
258            url,
259            heartbeat: Some(heartbeat),
260            auth_token: Arc::new(Mutex::new(None)),
261            reconnect_headers: Arc::new(Mutex::new(None)),
262            connection_mode,
263            cmd_tx: Arc::new(tokio::sync::RwLock::new(cmd_tx)),
264            out_rx: None,
265            signal: Arc::new(AtomicBool::new(false)),
266            cancellation_token: Arc::new(ArcSwap::from_pointee(CancellationToken::new())),
267            task_handle: Arc::new(SharedTaskSlot::new()),
268            connect_lock: Arc::new(tokio::sync::Mutex::new(())),
269            subscriptions: SubscriptionState::new(AX_TOPIC_DELIMITER),
270            request_id_counter: Arc::new(AtomicI64::new(1)),
271            subscribe_lock: Arc::new(tokio::sync::Mutex::new(())),
272            symbol_data_types: Arc::new(AtomicMap::new()),
273            status_invalidations: Arc::new(Mutex::new(AHashSet::new())),
274            transport_backend,
275            proxy_url: proxy_url.map(SecretString::from),
276            socket_control: None,
277        }
278    }
279
280    /// Configures socket state reporting and reconnect control.
281    #[must_use]
282    pub fn with_socket_control(mut self, control: SocketControl) -> Self {
283        self.socket_control = Some(control);
284        self
285    }
286
287    /// Returns the WebSocket URL.
288    #[must_use]
289    pub fn url(&self) -> &str {
290        &self.url
291    }
292
293    /// Sets the authentication token for subsequent connections.
294    ///
295    /// This should be called before `connect()` if authentication is required.
296    pub fn set_auth_token(&self, token: impl Into<SecretString>) {
297        *self.auth_token.lock() = Some(token.into());
298    }
299
300    /// Updates the token used by future automatic reconnect attempts.
301    ///
302    /// Updating the token does not interrupt the active WebSocket connection.
303    ///
304    /// # Errors
305    ///
306    /// Returns an error if the reconnect header cannot be updated.
307    pub fn update_auth_token(&self, token: SecretString) -> AxWsResult<()> {
308        let value = format!("Bearer {}", token.expose_secret());
309
310        if let Some(headers) = self.reconnect_headers.lock().as_ref() {
311            headers
312                .update("Authorization", &value)
313                .map_err(|e| AxWsClientError::Transport(e.to_string()))?;
314        }
315        self.set_auth_token(token);
316        Ok(())
317    }
318
319    /// Returns whether the client is currently connected and active.
320    #[must_use]
321    pub fn is_active(&self) -> bool {
322        let connection_mode_arc = self.connection_mode.load();
323        ConnectionMode::from_atomic(&connection_mode_arc).is_active()
324            && !self.signal.load(Ordering::Acquire)
325    }
326
327    /// Returns whether the client is closed.
328    #[must_use]
329    pub fn is_closed(&self) -> bool {
330        let connection_mode_arc = self.connection_mode.load();
331        ConnectionMode::from_atomic(&connection_mode_arc).is_closed()
332            || self.signal.load(Ordering::Acquire)
333    }
334
335    /// Returns the number of confirmed subscriptions.
336    #[must_use]
337    pub fn subscription_count(&self) -> usize {
338        self.subscriptions.len()
339    }
340
341    /// Returns the symbol data types map (shared with handler).
342    #[must_use]
343    pub fn symbol_data_types(&self) -> Arc<AtomicMap<String, SymbolDataTypes>> {
344        Arc::clone(&self.symbol_data_types)
345    }
346
347    /// Returns the shared set of symbols whose instrument status cache has been invalidated.
348    pub fn status_invalidations(&self) -> Arc<Mutex<AHashSet<Ustr>>> {
349        Arc::clone(&self.status_invalidations)
350    }
351
352    fn next_request_id(&self) -> i64 {
353        self.request_id_counter.fetch_add(1, Ordering::Relaxed)
354    }
355
356    fn is_subscribed_topic(&self, topic: &str) -> bool {
357        let (channel, symbol) = topic
358            .split_once(AX_TOPIC_DELIMITER)
359            .map_or((topic, None), |(c, s)| (c, Some(s)));
360        let channel_ustr = Ustr::from(channel);
361        let symbol_ustr = symbol.map_or_else(|| Ustr::from(""), Ustr::from);
362        self.subscriptions
363            .is_subscribed(&channel_ustr, &symbol_ustr)
364    }
365
366    /// Establishes the WebSocket connection.
367    ///
368    /// # Errors
369    ///
370    /// Returns an error if the connection cannot be established or the initial handler command
371    /// cannot be sent.
372    pub async fn connect(&mut self) -> AxWsResult<()> {
373        let connect_lock = Arc::clone(&self.connect_lock);
374        let _guard = connect_lock.lock().await;
375
376        if !self.task_handle.is_empty() && !self.task_handle.is_finished() {
377            return Err(AxWsClientError::Transport(
378                "WebSocket handler is already running".to_string(),
379            ));
380        }
381
382        if let Some(outcome) = self
383            .task_handle
384            .finish(Duration::from_secs(2), Duration::from_secs(2))
385            .await
386        {
387            match outcome {
388                TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted => {}
389                TaskJoinOutcome::Failed(error) => {
390                    return Err(AxWsClientError::Transport(format!(
391                        "Previous WebSocket handler failed: {error}"
392                    )));
393                }
394                TaskJoinOutcome::Incomplete => {
395                    return Err(AxWsClientError::Transport(
396                        "Previous WebSocket handler did not stop within shutdown bounds"
397                            .to_string(),
398                    ));
399                }
400            }
401        }
402
403        self.signal.store(false, Ordering::Release);
404        let cancellation_token = CancellationToken::new();
405        self.cancellation_token
406            .store(Arc::new(cancellation_token.clone()));
407
408        let (raw_handler, raw_rx) = channel_message_handler();
409
410        // No-op: ping responses are handled internally by the WebSocketClient
411        let ping_handler: PingHandler = Arc::new(move |_payload: Vec<u8>| {});
412
413        let mut headers = create_standard_nautilus_headers();
414
415        let auth_token = self.auth_token.lock().clone();
416
417        if let Some(token) = auth_token {
418            headers.push((
419                "Authorization".to_string(),
420                format!("Bearer {}", token.expose_secret()),
421            ));
422        }
423
424        let config = WebSocketConfig {
425            url: self.url.clone(),
426            headers,
427            heartbeat_interval_secs: self.heartbeat,
428            heartbeat_payload: None, // Ax server sends heartbeats
429            connect_timeout_ms: Some(5_000),
430            reconnect_delay_initial_ms: Some(500),
431            reconnect_delay_max_ms: Some(5_000),
432            reconnect_backoff_factor: Some(1.5),
433            reconnect_jitter_ms: Some(250),
434            reconnect_max_attempts: None,
435            heartbeat_timeout_secs: None,
436            idle_timeout_ms: None,
437            backend: self.transport_backend,
438            proxy_url: self
439                .proxy_url
440                .as_ref()
441                .map(|url| url.expose_secret().to_owned()),
442        };
443
444        let client = WebSocketClient::builder()
445            .config(config.clone())
446            .message_handler(raw_handler.clone())
447            .ping_handler(ping_handler.clone())
448            .initial_connect_retry_policy(Self::initial_connect_retry_policy())
449            .cancellation_token(cancellation_token)
450            .maybe_state_sink(self.socket_control.as_ref().map(SocketControl::sink))
451            .connect()
452            .await
453            .map_err(|e| {
454                AxWsClientError::Transport(format!("Failed to connect to {}: {e}", self.url))
455            })?;
456
457        self.connection_mode.store(client.connection_mode_atomic());
458        let reconnect_handle = client.reconnect_handle();
459        *self.reconnect_headers.lock() = Some(client.reconnect_headers());
460
461        let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<AxDataWsMessage>();
462        self.out_rx = Some(Arc::new(out_rx));
463
464        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
465        *self.cmd_tx.write().await = cmd_tx.clone();
466
467        self.send_cmd(HandlerCommand::SetClient(client)).await?;
468
469        let signal = Arc::clone(&self.signal);
470        let subscriptions = self.subscriptions.clone();
471
472        if let Err(e) = self.task_handle.spawn(async move {
473            let mut handler =
474                AxMdWsFeedHandler::new(signal.clone(), cmd_rx, raw_rx, subscriptions.clone());
475
476            while let Some(msg) = handler.next().await {
477                if matches!(msg, AxDataWsMessage::Reconnected) {
478                    log::info!("WebSocket reconnected, subscriptions will be replayed");
479                }
480
481                if out_tx.send(msg).is_err() {
482                    log::debug!("Output channel closed");
483                    break;
484                }
485            }
486
487            log::debug!("Handler loop exited");
488        }) {
489            self.out_rx = None;
490            return Err(AxWsClientError::Transport(format!(
491                "Failed to start WebSocket handler task: {e}"
492            )));
493        }
494
495        if let Some(control) = &self.socket_control {
496            control.register(move || reconnect_handle.request_reconnect());
497        }
498
499        Ok(())
500    }
501
502    /// Subscribes to order book deltas for a symbol.
503    ///
504    /// Uses reference counting so the underlying AX subscription is only
505    /// removed when all data types have been unsubscribed.
506    ///
507    /// # Errors
508    ///
509    /// Returns an error if the subscription command cannot be sent.
510    pub async fn subscribe_book_deltas(
511        &self,
512        symbol: &str,
513        level: AxMarketDataLevel,
514    ) -> AxWsResult<()> {
515        let _guard = self.subscribe_lock.lock().await;
516
517        let current = self
518            .symbol_data_types
519            .load()
520            .get(symbol)
521            .cloned()
522            .unwrap_or_default();
523
524        if current.book_level == Some(level) {
525            log::debug!("Book deltas already subscribed for {symbol} at {level:?}, skipping");
526            return Ok(());
527        }
528
529        let old_spec = current.effective_subscription();
530        let mut next = current.clone();
531        next.book_level = Some(level);
532        let new_spec = next.effective_subscription();
533
534        self.update_data_subscription(symbol, old_spec, new_spec)
535            .await?;
536
537        self.symbol_data_types.rcu(|m| {
538            let entry = m.entry(symbol.to_string()).or_default();
539            entry.book_level = Some(level);
540        });
541
542        Ok(())
543    }
544
545    /// Subscribes to quote data for a symbol.
546    ///
547    /// Uses reference counting so the underlying AX subscription is only
548    /// removed when all data types have been unsubscribed.
549    ///
550    /// # Errors
551    ///
552    /// Returns an error if the subscription command cannot be sent.
553    pub async fn subscribe_quotes(&self, symbol: &str) -> AxWsResult<()> {
554        let _guard = self.subscribe_lock.lock().await;
555
556        let current = self
557            .symbol_data_types
558            .load()
559            .get(symbol)
560            .cloned()
561            .unwrap_or_default();
562        let old_spec = current.effective_subscription();
563        let mut next = current.clone();
564        next.quotes = true;
565        let new_spec = next.effective_subscription();
566
567        self.update_data_subscription(symbol, old_spec, new_spec)
568            .await?;
569
570        self.symbol_data_types.rcu(|m| {
571            m.entry(symbol.to_string()).or_default().quotes = true;
572        });
573
574        Ok(())
575    }
576
577    /// Subscribes to trade data for a symbol.
578    ///
579    /// Uses reference counting so the underlying AX subscription is only
580    /// removed when all data types have been unsubscribed.
581    ///
582    /// # Errors
583    ///
584    /// Returns an error if the subscription command cannot be sent.
585    pub async fn subscribe_trades(&self, symbol: &str) -> AxWsResult<()> {
586        let _guard = self.subscribe_lock.lock().await;
587
588        let current = self
589            .symbol_data_types
590            .load()
591            .get(symbol)
592            .cloned()
593            .unwrap_or_default();
594        let old_spec = current.effective_subscription();
595        let mut next = current.clone();
596        next.trades = true;
597        let new_spec = next.effective_subscription();
598
599        self.update_data_subscription(symbol, old_spec, new_spec)
600            .await?;
601
602        self.symbol_data_types.rcu(|m| {
603            m.entry(symbol.to_string()).or_default().trades = true;
604        });
605
606        Ok(())
607    }
608
609    /// Unsubscribes from order book deltas for a symbol.
610    ///
611    /// The underlying AX subscription is only removed when all data types
612    /// (quotes, trades, book) have been unsubscribed.
613    ///
614    /// # Errors
615    ///
616    /// Returns an error if the unsubscribe command cannot be sent.
617    pub async fn unsubscribe_book_deltas(&self, symbol: &str) -> AxWsResult<()> {
618        let _guard = self.subscribe_lock.lock().await;
619
620        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
621            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe book deltas");
622            return Ok(());
623        };
624        let old_spec = current.effective_subscription();
625        let mut next = current.clone();
626        next.book_level = None;
627        let new_spec = next.effective_subscription();
628
629        self.update_data_subscription(symbol, old_spec, new_spec)
630            .await?;
631
632        self.symbol_data_types.rcu(|m| {
633            if let Some(entry) = m.get_mut(symbol) {
634                entry.book_level = None;
635                if entry.is_empty() {
636                    m.remove(symbol);
637                }
638            }
639        });
640
641        Ok(())
642    }
643
644    /// Unsubscribes from quote data for a symbol.
645    ///
646    /// The underlying AX subscription is only removed when all data types
647    /// (quotes, trades, book) have been unsubscribed.
648    ///
649    /// # Errors
650    ///
651    /// Returns an error if the unsubscribe command cannot be sent.
652    pub async fn unsubscribe_quotes(&self, symbol: &str) -> AxWsResult<()> {
653        let _guard = self.subscribe_lock.lock().await;
654
655        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
656            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe quotes");
657            return Ok(());
658        };
659        let old_spec = current.effective_subscription();
660        let mut next = current.clone();
661        next.quotes = false;
662        let new_spec = next.effective_subscription();
663
664        self.update_data_subscription(symbol, old_spec, new_spec)
665            .await?;
666
667        self.symbol_data_types.rcu(|m| {
668            if let Some(entry) = m.get_mut(symbol) {
669                entry.quotes = false;
670                if entry.is_empty() {
671                    m.remove(symbol);
672                }
673            }
674        });
675
676        Ok(())
677    }
678
679    /// Unsubscribes from trade data for a symbol.
680    ///
681    /// The underlying AX subscription is only removed when all data types
682    /// (quotes, trades, book) have been unsubscribed.
683    ///
684    /// # Errors
685    ///
686    /// Returns an error if the unsubscribe command cannot be sent.
687    pub async fn unsubscribe_trades(&self, symbol: &str) -> AxWsResult<()> {
688        let _guard = self.subscribe_lock.lock().await;
689
690        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
691            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe trades");
692            return Ok(());
693        };
694        let old_spec = current.effective_subscription();
695        let mut next = current.clone();
696        next.trades = false;
697        let new_spec = next.effective_subscription();
698
699        self.update_data_subscription(symbol, old_spec, new_spec)
700            .await?;
701
702        self.symbol_data_types.rcu(|m| {
703            if let Some(entry) = m.get_mut(symbol) {
704                entry.trades = false;
705                if entry.is_empty() {
706                    m.remove(symbol);
707                }
708            }
709        });
710
711        Ok(())
712    }
713
714    /// Subscribes to mark prices for a symbol.
715    ///
716    /// Ensures at least an L1 subscription so that ticker messages
717    /// (which carry the mark price field) are received.
718    ///
719    /// # Errors
720    ///
721    /// Returns an error if the subscription command cannot be sent.
722    pub async fn subscribe_mark_prices(&self, symbol: &str) -> AxWsResult<()> {
723        let _guard = self.subscribe_lock.lock().await;
724
725        let current = self
726            .symbol_data_types
727            .load()
728            .get(symbol)
729            .cloned()
730            .unwrap_or_default();
731        let old_spec = current.effective_subscription();
732        let mut next = current.clone();
733        next.mark_prices = true;
734        let new_spec = next.effective_subscription();
735
736        self.update_data_subscription(symbol, old_spec, new_spec)
737            .await?;
738
739        self.symbol_data_types.rcu(|m| {
740            m.entry(symbol.to_string()).or_default().mark_prices = true;
741        });
742
743        Ok(())
744    }
745
746    /// Unsubscribes from mark prices for a symbol.
747    ///
748    /// The underlying AX subscription is only removed when all data types
749    /// have been unsubscribed.
750    ///
751    /// # Errors
752    ///
753    /// Returns an error if the unsubscribe command cannot be sent.
754    pub async fn unsubscribe_mark_prices(&self, symbol: &str) -> AxWsResult<()> {
755        let _guard = self.subscribe_lock.lock().await;
756
757        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
758            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe mark prices");
759            return Ok(());
760        };
761        let old_spec = current.effective_subscription();
762        let mut next = current.clone();
763        next.mark_prices = false;
764        let new_spec = next.effective_subscription();
765
766        self.update_data_subscription(symbol, old_spec, new_spec)
767            .await?;
768
769        self.symbol_data_types.rcu(|m| {
770            if let Some(entry) = m.get_mut(symbol) {
771                entry.mark_prices = false;
772                if entry.is_empty() {
773                    m.remove(symbol);
774                }
775            }
776        });
777
778        Ok(())
779    }
780
781    /// Subscribes to instrument status for a symbol.
782    ///
783    /// Ensures at least an L1 subscription so that ticker messages
784    /// (which carry the instrument state field) are received.
785    ///
786    /// # Errors
787    ///
788    /// Returns an error if the subscription command cannot be sent.
789    pub async fn subscribe_instrument_status(&self, symbol: &str) -> AxWsResult<()> {
790        let _guard = self.subscribe_lock.lock().await;
791
792        let current = self
793            .symbol_data_types
794            .load()
795            .get(symbol)
796            .cloned()
797            .unwrap_or_default();
798        let old_spec = current.effective_subscription();
799        let mut next = current.clone();
800        next.instrument_status = true;
801        let new_spec = next.effective_subscription();
802
803        self.update_data_subscription(symbol, old_spec, new_spec)
804            .await?;
805
806        self.symbol_data_types.rcu(|m| {
807            m.entry(symbol.to_string()).or_default().instrument_status = true;
808        });
809
810        Ok(())
811    }
812
813    /// Unsubscribes from instrument status for a symbol.
814    ///
815    /// The underlying AX subscription is only removed when all data types
816    /// have been unsubscribed.
817    ///
818    /// # Errors
819    ///
820    /// Returns an error if the unsubscribe command cannot be sent.
821    pub async fn unsubscribe_instrument_status(&self, symbol: &str) -> AxWsResult<()> {
822        let _guard = self.subscribe_lock.lock().await;
823
824        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
825            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe instrument status");
826            return Ok(());
827        };
828        let old_spec = current.effective_subscription();
829        let mut next = current.clone();
830        next.instrument_status = false;
831        let new_spec = next.effective_subscription();
832
833        self.update_data_subscription(symbol, old_spec, new_spec)
834            .await?;
835
836        self.symbol_data_types.rcu(|m| {
837            if let Some(entry) = m.get_mut(symbol) {
838                entry.instrument_status = false;
839                if entry.is_empty() {
840                    m.remove(symbol);
841                }
842            }
843        });
844
845        self.status_invalidations.lock().insert(Ustr::from(symbol));
846
847        Ok(())
848    }
849
850    async fn update_data_subscription(
851        &self,
852        symbol: &str,
853        old_spec: Option<AxMdSubscriptionSpec>,
854        new_spec: Option<AxMdSubscriptionSpec>,
855    ) -> AxWsResult<()> {
856        if old_spec == new_spec {
857            return Ok(());
858        }
859
860        match (old_spec, new_spec) {
861            (None, Some(spec)) => {
862                log::debug!("Subscribing {symbol} at {spec:?}");
863                self.send_subscribe(symbol, spec).await
864            }
865            (Some(old), None) => {
866                log::debug!("Unsubscribing {symbol} (no remaining data types)");
867                self.send_unsubscribe(symbol, old).await
868            }
869            (Some(old), Some(new)) => {
870                log::debug!("Resubscribing {symbol}: {old:?} -> {new:?}");
871                self.send_unsubscribe(symbol, old).await?;
872                if let Err(e) = self.send_subscribe(symbol, new).await {
873                    log::warn!("Resubscribe failed for {symbol} at {new:?}: {e}");
874                    if let Err(restore_err) = self.send_subscribe(symbol, old).await {
875                        log::error!(
876                            "Failed to restore {symbol} at {old:?}: {restore_err}, \
877                             reconnection required"
878                        );
879                        self.subscriptions.mark_subscribe(&old.topic(symbol));
880                    }
881                    return Err(e);
882                }
883                Ok(())
884            }
885            (None, None) => Ok(()),
886        }
887    }
888
889    async fn send_subscribe(&self, symbol: &str, spec: AxMdSubscriptionSpec) -> AxWsResult<()> {
890        let topic = spec.topic(symbol);
891        let request_id = self.next_request_id();
892
893        self.subscriptions.mark_subscribe(&topic);
894
895        if let Err(e) = self
896            .send_cmd(HandlerCommand::Subscribe {
897                request_id,
898                symbol: Ustr::from(symbol),
899                spec,
900            })
901            .await
902        {
903            self.subscriptions.mark_unsubscribe(&topic);
904            return Err(e);
905        }
906
907        Ok(())
908    }
909
910    async fn send_unsubscribe(&self, symbol: &str, spec: AxMdSubscriptionSpec) -> AxWsResult<()> {
911        let request_id = self.next_request_id();
912        let topic = spec.topic(symbol);
913        let was_pending = self
914            .subscriptions
915            .pending_subscribe_topics()
916            .contains(&topic);
917
918        self.subscriptions.mark_unsubscribe(&topic);
919
920        if let Err(e) = self
921            .send_cmd(HandlerCommand::Unsubscribe {
922                request_id,
923                symbol: Ustr::from(symbol),
924                topic: topic.clone(),
925            })
926            .await
927        {
928            self.restore_unsubscribe_state(&topic, was_pending);
929            return Err(e);
930        }
931
932        Ok(())
933    }
934
935    /// Subscribes to candle data for a symbol.
936    ///
937    /// Skips sending if already subscribed or subscription is pending.
938    ///
939    /// # Errors
940    ///
941    /// Returns an error if the subscription command cannot be sent.
942    pub async fn subscribe_candles(&self, symbol: &str, width: AxCandleWidth) -> AxWsResult<()> {
943        let _guard = self.subscribe_lock.lock().await;
944        let topic = format!("candles:{symbol}:{width:?}");
945
946        // Skip if already subscribed or pending
947        if self.is_subscribed_topic(&topic) {
948            log::debug!("Already subscribed to {topic}, skipping");
949            return Ok(());
950        }
951
952        let request_id = self.next_request_id();
953
954        // Mark pending BEFORE sending to prevent race conditions with concurrent subscribes
955        self.subscriptions.mark_subscribe(&topic);
956
957        if let Err(e) = self
958            .send_cmd(HandlerCommand::SubscribeCandles {
959                request_id,
960                symbol: Ustr::from(symbol),
961                width,
962            })
963            .await
964        {
965            // Rollback pending state on send failure
966            self.subscriptions.mark_unsubscribe(&topic);
967            return Err(e);
968        }
969
970        Ok(())
971    }
972
973    /// Unsubscribes from candle data for a symbol.
974    ///
975    /// # Errors
976    ///
977    /// Returns an error if the unsubscribe command cannot be sent.
978    pub async fn unsubscribe_candles(&self, symbol: &str, width: AxCandleWidth) -> AxWsResult<()> {
979        let _guard = self.subscribe_lock.lock().await;
980        let request_id = self.next_request_id();
981        let topic = format!("candles:{symbol}:{width:?}");
982        let was_pending = self
983            .subscriptions
984            .pending_subscribe_topics()
985            .contains(&topic);
986
987        if !self.is_subscribed_topic(&topic) {
988            log::debug!("Not subscribed to {topic}, skipping unsubscribe");
989            return Ok(());
990        }
991
992        self.subscriptions.mark_unsubscribe(&topic);
993
994        if let Err(e) = self
995            .send_cmd(HandlerCommand::UnsubscribeCandles {
996                request_id,
997                symbol: Ustr::from(symbol),
998                width,
999                topic: topic.clone(),
1000            })
1001            .await
1002        {
1003            self.restore_unsubscribe_state(&topic, was_pending);
1004            return Err(e);
1005        }
1006
1007        Ok(())
1008    }
1009
1010    fn restore_unsubscribe_state(&self, topic: &str, was_pending: bool) {
1011        self.subscriptions.confirm_unsubscribe(topic);
1012        self.subscriptions.mark_subscribe(topic);
1013        if !was_pending {
1014            self.subscriptions.confirm_subscribe(topic);
1015        }
1016    }
1017
1018    /// Returns a stream of WebSocket messages.
1019    ///
1020    /// # Panics
1021    ///
1022    /// Panics if called before `connect()` or if the stream has already been taken.
1023    pub fn stream(&mut self) -> impl futures_util::Stream<Item = AxDataWsMessage> + 'static {
1024        let rx = self
1025            .out_rx
1026            .take()
1027            .expect("Stream receiver already taken or client not connected - stream() can only be called once");
1028        let mut rx = Arc::try_unwrap(rx).expect(
1029            "Cannot take ownership of stream - client was cloned and other references exist",
1030        );
1031        async_stream::stream! {
1032            while let Some(msg) = rx.recv().await {
1033                yield msg;
1034            }
1035        }
1036    }
1037
1038    pub(crate) fn begin_shutdown(&self) {
1039        self.cancellation_token.load().cancel();
1040        self.signal.store(true, Ordering::Release);
1041    }
1042
1043    /// Disconnects the WebSocket connection gracefully.
1044    pub async fn disconnect(&self) {
1045        log::debug!("Disconnecting WebSocket");
1046        let _ = self.send_cmd(HandlerCommand::Disconnect).await;
1047    }
1048
1049    /// Closes the WebSocket connection and cleans up resources.
1050    ///
1051    /// # Errors
1052    ///
1053    /// Returns an error if the handler task fails or does not stop after abort.
1054    pub async fn close(&mut self) -> anyhow::Result<()> {
1055        let connect_lock = Arc::clone(&self.connect_lock);
1056        let _guard = connect_lock.lock().await;
1057        log::debug!("Closing WebSocket client");
1058
1059        // Send disconnect first to allow graceful cleanup before signal
1060        self.cancellation_token.load().cancel();
1061        let _ = self.send_cmd(HandlerCommand::Disconnect).await;
1062        tokio::time::sleep(Duration::from_millis(50)).await;
1063        self.signal.store(true, Ordering::Release);
1064
1065        let outcome = self
1066            .task_handle
1067            .finish(Duration::from_secs(2), Duration::from_secs(2))
1068            .await;
1069
1070        *self.reconnect_headers.lock() = None;
1071
1072        if let Some(control) = &self.socket_control {
1073            control.deregister();
1074        }
1075
1076        match outcome {
1077            None | Some(TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted) => Ok(()),
1078            Some(TaskJoinOutcome::Failed(error)) => Err(anyhow::anyhow!(
1079                "Architect AX data WebSocket handler failed: {error}"
1080            )),
1081            Some(TaskJoinOutcome::Incomplete) => Err(anyhow::anyhow!(
1082                "Architect AX data WebSocket handler did not stop after abort"
1083            )),
1084        }
1085    }
1086
1087    async fn send_cmd(&self, cmd: HandlerCommand) -> AxWsResult<()> {
1088        let guard = self.cmd_tx.read().await;
1089        guard
1090            .send(cmd)
1091            .map_err(|e| AxWsClientError::ChannelError(e.to_string()))
1092    }
1093}
1094
1095impl Drop for AxMdWebSocketClient {
1096    fn drop(&mut self) {
1097        if Arc::strong_count(&self.task_handle) == 1 && !self.task_handle.is_empty() {
1098            self.cancellation_token.load().cancel();
1099            self.signal.store(true, Ordering::Release);
1100            self.task_handle.abort();
1101
1102            if let Some(control) = &self.socket_control {
1103                control.deregister();
1104            }
1105        }
1106    }
1107}
1108
1109#[cfg(test)]
1110mod tests {
1111    use rstest::rstest;
1112
1113    use super::*;
1114
1115    #[rstest]
1116    fn test_auth_token_uses_secret_string_owner() {
1117        let client = AxMdWebSocketClient::new(
1118            "ws://localhost:9999/md/ws".to_string(),
1119            "initial-token".to_string(),
1120            30,
1121            TransportBackend::default(),
1122            None,
1123        );
1124
1125        client.set_auth_token("replacement-token".to_string());
1126
1127        let token = client.auth_token.lock();
1128        assert_eq!(
1129            token.as_ref().map(SecretString::expose_secret),
1130            Some("replacement-token")
1131        );
1132    }
1133
1134    #[tokio::test]
1135    async fn test_drop_aborts_handler_task() {
1136        let client = AxMdWebSocketClient::new(
1137            "ws://localhost:9999/md/ws".to_string(),
1138            "test_token".to_string(),
1139            30,
1140            TransportBackend::default(),
1141            None,
1142        );
1143        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
1144        let handle = tokio::spawn(async move {
1145            started_tx.send(()).expect("started receiver");
1146            std::future::pending::<()>().await;
1147        });
1148        let abort_handle = handle.abort_handle();
1149        client.task_handle.insert(handle);
1150        started_rx.await.expect("handler task started");
1151
1152        drop(client);
1153
1154        tokio::time::timeout(Duration::from_secs(1), async {
1155            while !abort_handle.is_finished() {
1156                tokio::task::yield_now().await;
1157            }
1158        })
1159        .await
1160        .expect("handler task aborted");
1161    }
1162
1163    #[rstest]
1164    fn test_effective_subscription_empty_returns_none() {
1165        let sdt = SymbolDataTypes::default();
1166        assert_eq!(sdt.effective_subscription(), None);
1167        assert!(sdt.is_empty());
1168    }
1169
1170    #[rstest]
1171    fn test_effective_subscription_book_level_takes_precedence() {
1172        let sdt = SymbolDataTypes {
1173            book_level: Some(AxMarketDataLevel::Level2),
1174            quotes: true,
1175            ..Default::default()
1176        };
1177        assert_eq!(
1178            sdt.effective_subscription(),
1179            Some(AxMdSubscriptionSpec::new(
1180                AxMarketDataLevel::Level2,
1181                Some(false),
1182                Some(false),
1183            ))
1184        );
1185        assert!(!sdt.is_empty());
1186    }
1187
1188    #[rstest]
1189    #[case(
1190        true,
1191        false,
1192        false,
1193        false,
1194        AxMarketDataLevel::Level1,
1195        Some(false),
1196        Some(false)
1197    )]
1198    #[case(false, true, false, false, AxMarketDataLevel::Trades, None, None)]
1199    #[case(
1200        false,
1201        false,
1202        true,
1203        false,
1204        AxMarketDataLevel::Level1,
1205        Some(false),
1206        Some(true)
1207    )]
1208    #[case(
1209        false,
1210        false,
1211        false,
1212        true,
1213        AxMarketDataLevel::Level1,
1214        Some(false),
1215        Some(true)
1216    )]
1217    fn test_effective_subscription_for_single_data_type(
1218        #[case] quotes: bool,
1219        #[case] trades: bool,
1220        #[case] mark_prices: bool,
1221        #[case] instrument_status: bool,
1222        #[case] level: AxMarketDataLevel,
1223        #[case] include_trades: Option<bool>,
1224        #[case] include_ticker: Option<bool>,
1225    ) {
1226        let sdt = SymbolDataTypes {
1227            quotes,
1228            trades,
1229            mark_prices,
1230            instrument_status,
1231            book_level: None,
1232        };
1233        assert_eq!(
1234            sdt.effective_subscription(),
1235            Some(AxMdSubscriptionSpec::new(
1236                level,
1237                include_trades,
1238                include_ticker,
1239            ))
1240        );
1241        assert!(!sdt.is_empty());
1242    }
1243
1244    #[rstest]
1245    #[case(false)]
1246    #[case(true)]
1247    #[tokio::test]
1248    async fn test_unsubscribe_send_failure_restores_subscription(#[case] was_pending: bool) {
1249        let client = AxMdWebSocketClient::new(
1250            "ws://localhost:9999/md/ws".to_string(),
1251            "test_token".to_string(),
1252            30,
1253            TransportBackend::default(),
1254            None,
1255        );
1256        let symbol = "EURUSD-PERP";
1257        let spec = AxMdSubscriptionSpec::new(AxMarketDataLevel::Level2, Some(false), Some(false));
1258        let topic = spec.topic(symbol);
1259        client.subscriptions.mark_subscribe(&topic);
1260        if !was_pending {
1261            client.subscriptions.confirm_subscribe(&topic);
1262        }
1263
1264        let error = client.send_unsubscribe(symbol, spec).await.unwrap_err();
1265
1266        assert_eq!(error.to_string(), "Channel error: channel closed");
1267        assert_eq!(client.subscription_count(), usize::from(!was_pending));
1268        assert_eq!(client.subscriptions.all_topics(), vec![topic]);
1269        assert_eq!(
1270            client.subscriptions.pending_subscribe_topics().len(),
1271            usize::from(was_pending)
1272        );
1273        assert!(client.subscriptions.pending_unsubscribe_topics().is_empty());
1274    }
1275
1276    #[rstest]
1277    #[case(false)]
1278    #[case(true)]
1279    #[tokio::test]
1280    async fn test_unsubscribe_candles_send_failure_restores_subscription(
1281        #[case] was_pending: bool,
1282    ) {
1283        let client = AxMdWebSocketClient::new(
1284            "ws://localhost:9999/md/ws".to_string(),
1285            "test_token".to_string(),
1286            30,
1287            TransportBackend::default(),
1288            None,
1289        );
1290        let symbol = "EURUSD-PERP";
1291        let width = AxCandleWidth::Minutes1;
1292        let topic = format!("candles:{symbol}:{width:?}");
1293        client.subscriptions.mark_subscribe(&topic);
1294        if !was_pending {
1295            client.subscriptions.confirm_subscribe(&topic);
1296        }
1297
1298        let error = client.unsubscribe_candles(symbol, width).await.unwrap_err();
1299
1300        assert_eq!(error.to_string(), "Channel error: channel closed");
1301        assert_eq!(client.subscription_count(), usize::from(!was_pending));
1302        assert_eq!(client.subscriptions.all_topics(), vec![topic]);
1303        assert_eq!(
1304            client.subscriptions.pending_subscribe_topics().len(),
1305            usize::from(was_pending)
1306        );
1307        assert!(client.subscriptions.pending_unsubscribe_topics().is_empty());
1308    }
1309
1310    #[rstest]
1311    #[tokio::test]
1312    async fn test_unsubscribe_candles_skips_untracked_topic() {
1313        let client = AxMdWebSocketClient::new(
1314            "ws://localhost:9999/md/ws".to_string(),
1315            "test_token".to_string(),
1316            30,
1317            TransportBackend::default(),
1318            None,
1319        );
1320
1321        client
1322            .unsubscribe_candles("EURUSD-PERP", AxCandleWidth::Minutes1)
1323            .await
1324            .unwrap();
1325
1326        assert_eq!(client.subscription_count(), 0);
1327        assert!(client.subscriptions.all_topics().is_empty());
1328        assert!(client.subscriptions.pending_subscribe_topics().is_empty());
1329        assert!(client.subscriptions.pending_unsubscribe_topics().is_empty());
1330    }
1331}
nautilus_architect_ax/websocket/data/client.rs

nautilus_architect_ax/websocket/data/
client.rs
