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    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, consts::NAUTILUS_USER_AGENT};
31use nautilus_live::{
32    SocketControl,
33    task::{SharedTaskSlot, TaskJoinOutcome},
34};
35use nautilus_network::{
36    http::USER_AGENT,
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 core::fmt::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<String>>>,
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<String>,
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: String,
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))),
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,
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,
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: String) {
297        *self.auth_token.lock() = Some(token);
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: String) -> AxWsResult<()> {
308        let value = format!("Bearer {token}");
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 = vec![(USER_AGENT.to_string(), NAUTILUS_USER_AGENT.to_string())];
414
415        let auth_token = self.auth_token.lock().clone();
416
417        if let Some(token) = auth_token {
418            headers.push(("Authorization".to_string(), format!("Bearer {token}")));
419        }
420
421        let config = WebSocketConfig {
422            url: self.url.clone(),
423            headers,
424            heartbeat_interval_secs: self.heartbeat,
425            heartbeat_payload: None, // Ax server sends heartbeats
426            connect_timeout_ms: Some(5_000),
427            reconnect_delay_initial_ms: Some(500),
428            reconnect_delay_max_ms: Some(5_000),
429            reconnect_backoff_factor: Some(1.5),
430            reconnect_jitter_ms: Some(250),
431            reconnect_max_attempts: None,
432            heartbeat_timeout_secs: None,
433            idle_timeout_ms: None,
434            backend: self.transport_backend,
435            proxy_url: self.proxy_url.clone(),
436        };
437
438        let client = WebSocketClient::builder()
439            .config(config.clone())
440            .message_handler(raw_handler.clone())
441            .ping_handler(ping_handler.clone())
442            .initial_connect_retry_policy(Self::initial_connect_retry_policy())
443            .cancellation_token(cancellation_token)
444            .maybe_state_sink(self.socket_control.as_ref().map(SocketControl::sink))
445            .connect()
446            .await
447            .map_err(|e| {
448                AxWsClientError::Transport(format!("Failed to connect to {}: {e}", self.url))
449            })?;
450
451        self.connection_mode.store(client.connection_mode_atomic());
452        let reconnect_handle = client.reconnect_handle();
453        *self.reconnect_headers.lock() = Some(client.reconnect_headers());
454
455        let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<AxDataWsMessage>();
456        self.out_rx = Some(Arc::new(out_rx));
457
458        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
459        *self.cmd_tx.write().await = cmd_tx.clone();
460
461        self.send_cmd(HandlerCommand::SetClient(client)).await?;
462
463        let signal = Arc::clone(&self.signal);
464        let subscriptions = self.subscriptions.clone();
465
466        if let Err(e) = self.task_handle.spawn(async move {
467            let mut handler =
468                AxMdWsFeedHandler::new(signal.clone(), cmd_rx, raw_rx, subscriptions.clone());
469
470            while let Some(msg) = handler.next().await {
471                if matches!(msg, AxDataWsMessage::Reconnected) {
472                    log::info!("WebSocket reconnected, subscriptions will be replayed");
473                }
474
475                if out_tx.send(msg).is_err() {
476                    log::debug!("Output channel closed");
477                    break;
478                }
479            }
480
481            log::debug!("Handler loop exited");
482        }) {
483            self.out_rx = None;
484            return Err(AxWsClientError::Transport(format!(
485                "Failed to start WebSocket handler task: {e}"
486            )));
487        }
488
489        if let Some(control) = &self.socket_control {
490            control.register(move || reconnect_handle.request_reconnect());
491        }
492
493        Ok(())
494    }
495
496    /// Subscribes to order book deltas for a symbol.
497    ///
498    /// Uses reference counting so the underlying AX subscription is only
499    /// removed when all data types have been unsubscribed.
500    ///
501    /// # Errors
502    ///
503    /// Returns an error if the subscription command cannot be sent.
504    pub async fn subscribe_book_deltas(
505        &self,
506        symbol: &str,
507        level: AxMarketDataLevel,
508    ) -> AxWsResult<()> {
509        let _guard = self.subscribe_lock.lock().await;
510
511        let current = self
512            .symbol_data_types
513            .load()
514            .get(symbol)
515            .cloned()
516            .unwrap_or_default();
517
518        if current.book_level == Some(level) {
519            log::debug!("Book deltas already subscribed for {symbol} at {level:?}, skipping");
520            return Ok(());
521        }
522
523        let old_spec = current.effective_subscription();
524        let mut next = current.clone();
525        next.book_level = Some(level);
526        let new_spec = next.effective_subscription();
527
528        self.update_data_subscription(symbol, old_spec, new_spec)
529            .await?;
530
531        self.symbol_data_types.rcu(|m| {
532            let entry = m.entry(symbol.to_string()).or_default();
533            entry.book_level = Some(level);
534        });
535
536        Ok(())
537    }
538
539    /// Subscribes to quote data for a symbol.
540    ///
541    /// Uses reference counting so the underlying AX subscription is only
542    /// removed when all data types have been unsubscribed.
543    ///
544    /// # Errors
545    ///
546    /// Returns an error if the subscription command cannot be sent.
547    pub async fn subscribe_quotes(&self, symbol: &str) -> AxWsResult<()> {
548        let _guard = self.subscribe_lock.lock().await;
549
550        let current = self
551            .symbol_data_types
552            .load()
553            .get(symbol)
554            .cloned()
555            .unwrap_or_default();
556        let old_spec = current.effective_subscription();
557        let mut next = current.clone();
558        next.quotes = true;
559        let new_spec = next.effective_subscription();
560
561        self.update_data_subscription(symbol, old_spec, new_spec)
562            .await?;
563
564        self.symbol_data_types.rcu(|m| {
565            m.entry(symbol.to_string()).or_default().quotes = true;
566        });
567
568        Ok(())
569    }
570
571    /// Subscribes to trade data for a symbol.
572    ///
573    /// Uses reference counting so the underlying AX subscription is only
574    /// removed when all data types have been unsubscribed.
575    ///
576    /// # Errors
577    ///
578    /// Returns an error if the subscription command cannot be sent.
579    pub async fn subscribe_trades(&self, symbol: &str) -> AxWsResult<()> {
580        let _guard = self.subscribe_lock.lock().await;
581
582        let current = self
583            .symbol_data_types
584            .load()
585            .get(symbol)
586            .cloned()
587            .unwrap_or_default();
588        let old_spec = current.effective_subscription();
589        let mut next = current.clone();
590        next.trades = true;
591        let new_spec = next.effective_subscription();
592
593        self.update_data_subscription(symbol, old_spec, new_spec)
594            .await?;
595
596        self.symbol_data_types.rcu(|m| {
597            m.entry(symbol.to_string()).or_default().trades = true;
598        });
599
600        Ok(())
601    }
602
603    /// Unsubscribes from order book deltas for a symbol.
604    ///
605    /// The underlying AX subscription is only removed when all data types
606    /// (quotes, trades, book) have been unsubscribed.
607    ///
608    /// # Errors
609    ///
610    /// Returns an error if the unsubscribe command cannot be sent.
611    pub async fn unsubscribe_book_deltas(&self, symbol: &str) -> AxWsResult<()> {
612        let _guard = self.subscribe_lock.lock().await;
613
614        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
615            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe book deltas");
616            return Ok(());
617        };
618        let old_spec = current.effective_subscription();
619        let mut next = current.clone();
620        next.book_level = None;
621        let new_spec = next.effective_subscription();
622
623        self.update_data_subscription(symbol, old_spec, new_spec)
624            .await?;
625
626        self.symbol_data_types.rcu(|m| {
627            if let Some(entry) = m.get_mut(symbol) {
628                entry.book_level = None;
629                if entry.is_empty() {
630                    m.remove(symbol);
631                }
632            }
633        });
634
635        Ok(())
636    }
637
638    /// Unsubscribes from quote data for a symbol.
639    ///
640    /// The underlying AX subscription is only removed when all data types
641    /// (quotes, trades, book) have been unsubscribed.
642    ///
643    /// # Errors
644    ///
645    /// Returns an error if the unsubscribe command cannot be sent.
646    pub async fn unsubscribe_quotes(&self, symbol: &str) -> AxWsResult<()> {
647        let _guard = self.subscribe_lock.lock().await;
648
649        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
650            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe quotes");
651            return Ok(());
652        };
653        let old_spec = current.effective_subscription();
654        let mut next = current.clone();
655        next.quotes = false;
656        let new_spec = next.effective_subscription();
657
658        self.update_data_subscription(symbol, old_spec, new_spec)
659            .await?;
660
661        self.symbol_data_types.rcu(|m| {
662            if let Some(entry) = m.get_mut(symbol) {
663                entry.quotes = false;
664                if entry.is_empty() {
665                    m.remove(symbol);
666                }
667            }
668        });
669
670        Ok(())
671    }
672
673    /// Unsubscribes from trade data for a symbol.
674    ///
675    /// The underlying AX subscription is only removed when all data types
676    /// (quotes, trades, book) have been unsubscribed.
677    ///
678    /// # Errors
679    ///
680    /// Returns an error if the unsubscribe command cannot be sent.
681    pub async fn unsubscribe_trades(&self, symbol: &str) -> AxWsResult<()> {
682        let _guard = self.subscribe_lock.lock().await;
683
684        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
685            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe trades");
686            return Ok(());
687        };
688        let old_spec = current.effective_subscription();
689        let mut next = current.clone();
690        next.trades = false;
691        let new_spec = next.effective_subscription();
692
693        self.update_data_subscription(symbol, old_spec, new_spec)
694            .await?;
695
696        self.symbol_data_types.rcu(|m| {
697            if let Some(entry) = m.get_mut(symbol) {
698                entry.trades = false;
699                if entry.is_empty() {
700                    m.remove(symbol);
701                }
702            }
703        });
704
705        Ok(())
706    }
707
708    /// Subscribes to mark prices for a symbol.
709    ///
710    /// Ensures at least an L1 subscription so that ticker messages
711    /// (which carry the mark price field) are received.
712    ///
713    /// # Errors
714    ///
715    /// Returns an error if the subscription command cannot be sent.
716    pub async fn subscribe_mark_prices(&self, symbol: &str) -> AxWsResult<()> {
717        let _guard = self.subscribe_lock.lock().await;
718
719        let current = self
720            .symbol_data_types
721            .load()
722            .get(symbol)
723            .cloned()
724            .unwrap_or_default();
725        let old_spec = current.effective_subscription();
726        let mut next = current.clone();
727        next.mark_prices = true;
728        let new_spec = next.effective_subscription();
729
730        self.update_data_subscription(symbol, old_spec, new_spec)
731            .await?;
732
733        self.symbol_data_types.rcu(|m| {
734            m.entry(symbol.to_string()).or_default().mark_prices = true;
735        });
736
737        Ok(())
738    }
739
740    /// Unsubscribes from mark prices for a symbol.
741    ///
742    /// The underlying AX subscription is only removed when all data types
743    /// have been unsubscribed.
744    ///
745    /// # Errors
746    ///
747    /// Returns an error if the unsubscribe command cannot be sent.
748    pub async fn unsubscribe_mark_prices(&self, symbol: &str) -> AxWsResult<()> {
749        let _guard = self.subscribe_lock.lock().await;
750
751        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
752            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe mark prices");
753            return Ok(());
754        };
755        let old_spec = current.effective_subscription();
756        let mut next = current.clone();
757        next.mark_prices = false;
758        let new_spec = next.effective_subscription();
759
760        self.update_data_subscription(symbol, old_spec, new_spec)
761            .await?;
762
763        self.symbol_data_types.rcu(|m| {
764            if let Some(entry) = m.get_mut(symbol) {
765                entry.mark_prices = false;
766                if entry.is_empty() {
767                    m.remove(symbol);
768                }
769            }
770        });
771
772        Ok(())
773    }
774
775    /// Subscribes to instrument status for a symbol.
776    ///
777    /// Ensures at least an L1 subscription so that ticker messages
778    /// (which carry the instrument state field) are received.
779    ///
780    /// # Errors
781    ///
782    /// Returns an error if the subscription command cannot be sent.
783    pub async fn subscribe_instrument_status(&self, symbol: &str) -> AxWsResult<()> {
784        let _guard = self.subscribe_lock.lock().await;
785
786        let current = self
787            .symbol_data_types
788            .load()
789            .get(symbol)
790            .cloned()
791            .unwrap_or_default();
792        let old_spec = current.effective_subscription();
793        let mut next = current.clone();
794        next.instrument_status = true;
795        let new_spec = next.effective_subscription();
796
797        self.update_data_subscription(symbol, old_spec, new_spec)
798            .await?;
799
800        self.symbol_data_types.rcu(|m| {
801            m.entry(symbol.to_string()).or_default().instrument_status = true;
802        });
803
804        Ok(())
805    }
806
807    /// Unsubscribes from instrument status for a symbol.
808    ///
809    /// The underlying AX subscription is only removed when all data types
810    /// have been unsubscribed.
811    ///
812    /// # Errors
813    ///
814    /// Returns an error if the unsubscribe command cannot be sent.
815    pub async fn unsubscribe_instrument_status(&self, symbol: &str) -> AxWsResult<()> {
816        let _guard = self.subscribe_lock.lock().await;
817
818        let Some(current) = self.symbol_data_types.load().get(symbol).cloned() else {
819            log::debug!("Symbol {symbol} not subscribed, skipping unsubscribe instrument status");
820            return Ok(());
821        };
822        let old_spec = current.effective_subscription();
823        let mut next = current.clone();
824        next.instrument_status = false;
825        let new_spec = next.effective_subscription();
826
827        self.update_data_subscription(symbol, old_spec, new_spec)
828            .await?;
829
830        self.symbol_data_types.rcu(|m| {
831            if let Some(entry) = m.get_mut(symbol) {
832                entry.instrument_status = false;
833                if entry.is_empty() {
834                    m.remove(symbol);
835                }
836            }
837        });
838
839        self.status_invalidations.lock().insert(Ustr::from(symbol));
840
841        Ok(())
842    }
843
844    async fn update_data_subscription(
845        &self,
846        symbol: &str,
847        old_spec: Option<AxMdSubscriptionSpec>,
848        new_spec: Option<AxMdSubscriptionSpec>,
849    ) -> AxWsResult<()> {
850        if old_spec == new_spec {
851            return Ok(());
852        }
853
854        match (old_spec, new_spec) {
855            (None, Some(spec)) => {
856                log::debug!("Subscribing {symbol} at {spec:?}");
857                self.send_subscribe(symbol, spec).await
858            }
859            (Some(old), None) => {
860                log::debug!("Unsubscribing {symbol} (no remaining data types)");
861                self.send_unsubscribe(symbol, old).await
862            }
863            (Some(old), Some(new)) => {
864                log::debug!("Resubscribing {symbol}: {old:?} -> {new:?}");
865                self.send_unsubscribe(symbol, old).await?;
866                if let Err(e) = self.send_subscribe(symbol, new).await {
867                    log::warn!("Resubscribe failed for {symbol} at {new:?}: {e}");
868                    if let Err(restore_err) = self.send_subscribe(symbol, old).await {
869                        log::error!(
870                            "Failed to restore {symbol} at {old:?}: {restore_err}, \
871                             reconnection required"
872                        );
873                        self.subscriptions.mark_subscribe(&old.topic(symbol));
874                    }
875                    return Err(e);
876                }
877                Ok(())
878            }
879            (None, None) => Ok(()),
880        }
881    }
882
883    async fn send_subscribe(&self, symbol: &str, spec: AxMdSubscriptionSpec) -> AxWsResult<()> {
884        let topic = spec.topic(symbol);
885        let request_id = self.next_request_id();
886
887        self.subscriptions.mark_subscribe(&topic);
888
889        if let Err(e) = self
890            .send_cmd(HandlerCommand::Subscribe {
891                request_id,
892                symbol: Ustr::from(symbol),
893                spec,
894            })
895            .await
896        {
897            self.subscriptions.mark_unsubscribe(&topic);
898            return Err(e);
899        }
900
901        Ok(())
902    }
903
904    async fn send_unsubscribe(&self, symbol: &str, spec: AxMdSubscriptionSpec) -> AxWsResult<()> {
905        let request_id = self.next_request_id();
906        let topic = spec.topic(symbol);
907        let was_pending = self
908            .subscriptions
909            .pending_subscribe_topics()
910            .contains(&topic);
911
912        self.subscriptions.mark_unsubscribe(&topic);
913
914        if let Err(e) = self
915            .send_cmd(HandlerCommand::Unsubscribe {
916                request_id,
917                symbol: Ustr::from(symbol),
918                topic: topic.clone(),
919            })
920            .await
921        {
922            self.restore_unsubscribe_state(&topic, was_pending);
923            return Err(e);
924        }
925
926        Ok(())
927    }
928
929    /// Subscribes to candle data for a symbol.
930    ///
931    /// Skips sending if already subscribed or subscription is pending.
932    ///
933    /// # Errors
934    ///
935    /// Returns an error if the subscription command cannot be sent.
936    pub async fn subscribe_candles(&self, symbol: &str, width: AxCandleWidth) -> AxWsResult<()> {
937        let _guard = self.subscribe_lock.lock().await;
938        let topic = format!("candles:{symbol}:{width:?}");
939
940        // Skip if already subscribed or pending
941        if self.is_subscribed_topic(&topic) {
942            log::debug!("Already subscribed to {topic}, skipping");
943            return Ok(());
944        }
945
946        let request_id = self.next_request_id();
947
948        // Mark pending BEFORE sending to prevent race conditions with concurrent subscribes
949        self.subscriptions.mark_subscribe(&topic);
950
951        if let Err(e) = self
952            .send_cmd(HandlerCommand::SubscribeCandles {
953                request_id,
954                symbol: Ustr::from(symbol),
955                width,
956            })
957            .await
958        {
959            // Rollback pending state on send failure
960            self.subscriptions.mark_unsubscribe(&topic);
961            return Err(e);
962        }
963
964        Ok(())
965    }
966
967    /// Unsubscribes from candle data for a symbol.
968    ///
969    /// # Errors
970    ///
971    /// Returns an error if the unsubscribe command cannot be sent.
972    pub async fn unsubscribe_candles(&self, symbol: &str, width: AxCandleWidth) -> AxWsResult<()> {
973        let _guard = self.subscribe_lock.lock().await;
974        let request_id = self.next_request_id();
975        let topic = format!("candles:{symbol}:{width:?}");
976        let was_pending = self
977            .subscriptions
978            .pending_subscribe_topics()
979            .contains(&topic);
980
981        if !self.is_subscribed_topic(&topic) {
982            log::debug!("Not subscribed to {topic}, skipping unsubscribe");
983            return Ok(());
984        }
985
986        self.subscriptions.mark_unsubscribe(&topic);
987
988        if let Err(e) = self
989            .send_cmd(HandlerCommand::UnsubscribeCandles {
990                request_id,
991                symbol: Ustr::from(symbol),
992                width,
993                topic: topic.clone(),
994            })
995            .await
996        {
997            self.restore_unsubscribe_state(&topic, was_pending);
998            return Err(e);
999        }
1000
1001        Ok(())
1002    }
1003
1004    fn restore_unsubscribe_state(&self, topic: &str, was_pending: bool) {
1005        self.subscriptions.confirm_unsubscribe(topic);
1006        self.subscriptions.mark_subscribe(topic);
1007        if !was_pending {
1008            self.subscriptions.confirm_subscribe(topic);
1009        }
1010    }
1011
1012    /// Returns a stream of WebSocket messages.
1013    ///
1014    /// # Panics
1015    ///
1016    /// Panics if called before `connect()` or if the stream has already been taken.
1017    pub fn stream(&mut self) -> impl futures_util::Stream<Item = AxDataWsMessage> + 'static {
1018        let rx = self
1019            .out_rx
1020            .take()
1021            .expect("Stream receiver already taken or client not connected - stream() can only be called once");
1022        let mut rx = Arc::try_unwrap(rx).expect(
1023            "Cannot take ownership of stream - client was cloned and other references exist",
1024        );
1025        async_stream::stream! {
1026            while let Some(msg) = rx.recv().await {
1027                yield msg;
1028            }
1029        }
1030    }
1031
1032    pub(crate) fn begin_shutdown(&self) {
1033        self.cancellation_token.load().cancel();
1034        self.signal.store(true, Ordering::Release);
1035    }
1036
1037    /// Disconnects the WebSocket connection gracefully.
1038    pub async fn disconnect(&self) {
1039        log::debug!("Disconnecting WebSocket");
1040        let _ = self.send_cmd(HandlerCommand::Disconnect).await;
1041    }
1042
1043    /// Closes the WebSocket connection and cleans up resources.
1044    ///
1045    /// # Errors
1046    ///
1047    /// Returns an error if the handler task fails or does not stop after abort.
1048    pub async fn close(&mut self) -> anyhow::Result<()> {
1049        let connect_lock = Arc::clone(&self.connect_lock);
1050        let _guard = connect_lock.lock().await;
1051        log::debug!("Closing WebSocket client");
1052
1053        // Send disconnect first to allow graceful cleanup before signal
1054        self.cancellation_token.load().cancel();
1055        let _ = self.send_cmd(HandlerCommand::Disconnect).await;
1056        tokio::time::sleep(Duration::from_millis(50)).await;
1057        self.signal.store(true, Ordering::Release);
1058
1059        let outcome = self
1060            .task_handle
1061            .finish(Duration::from_secs(2), Duration::from_secs(2))
1062            .await;
1063
1064        *self.reconnect_headers.lock() = None;
1065
1066        if let Some(control) = &self.socket_control {
1067            control.deregister();
1068        }
1069
1070        match outcome {
1071            None | Some(TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted) => Ok(()),
1072            Some(TaskJoinOutcome::Failed(error)) => Err(anyhow::anyhow!(
1073                "Architect AX data WebSocket handler failed: {error}"
1074            )),
1075            Some(TaskJoinOutcome::Incomplete) => Err(anyhow::anyhow!(
1076                "Architect AX data WebSocket handler did not stop after abort"
1077            )),
1078        }
1079    }
1080
1081    async fn send_cmd(&self, cmd: HandlerCommand) -> AxWsResult<()> {
1082        let guard = self.cmd_tx.read().await;
1083        guard
1084            .send(cmd)
1085            .map_err(|e| AxWsClientError::ChannelError(e.to_string()))
1086    }
1087}
1088
1089impl Drop for AxMdWebSocketClient {
1090    fn drop(&mut self) {
1091        if Arc::strong_count(&self.task_handle) == 1 && !self.task_handle.is_empty() {
1092            self.cancellation_token.load().cancel();
1093            self.signal.store(true, Ordering::Release);
1094            self.task_handle.abort();
1095
1096            if let Some(control) = &self.socket_control {
1097                control.deregister();
1098            }
1099        }
1100    }
1101}
1102
1103#[cfg(test)]
1104mod tests {
1105    use rstest::rstest;
1106
1107    use super::*;
1108
1109    #[tokio::test]
1110    async fn test_drop_aborts_handler_task() {
1111        let client = AxMdWebSocketClient::new(
1112            "ws://localhost:9999/md/ws".to_string(),
1113            "test_token".to_string(),
1114            30,
1115            TransportBackend::default(),
1116            None,
1117        );
1118        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
1119        let handle = tokio::spawn(async move {
1120            started_tx.send(()).expect("started receiver");
1121            std::future::pending::<()>().await;
1122        });
1123        let abort_handle = handle.abort_handle();
1124        client.task_handle.insert(handle);
1125        started_rx.await.expect("handler task started");
1126
1127        drop(client);
1128
1129        tokio::time::timeout(Duration::from_secs(1), async {
1130            while !abort_handle.is_finished() {
1131                tokio::task::yield_now().await;
1132            }
1133        })
1134        .await
1135        .expect("handler task aborted");
1136    }
1137
1138    #[rstest]
1139    fn test_effective_subscription_empty_returns_none() {
1140        let sdt = SymbolDataTypes::default();
1141        assert_eq!(sdt.effective_subscription(), None);
1142        assert!(sdt.is_empty());
1143    }
1144
1145    #[rstest]
1146    fn test_effective_subscription_book_level_takes_precedence() {
1147        let sdt = SymbolDataTypes {
1148            book_level: Some(AxMarketDataLevel::Level2),
1149            quotes: true,
1150            ..Default::default()
1151        };
1152        assert_eq!(
1153            sdt.effective_subscription(),
1154            Some(AxMdSubscriptionSpec::new(
1155                AxMarketDataLevel::Level2,
1156                Some(false),
1157                Some(false),
1158            ))
1159        );
1160        assert!(!sdt.is_empty());
1161    }
1162
1163    #[rstest]
1164    #[case(
1165        true,
1166        false,
1167        false,
1168        false,
1169        AxMarketDataLevel::Level1,
1170        Some(false),
1171        Some(false)
1172    )]
1173    #[case(false, true, false, false, AxMarketDataLevel::Trades, None, None)]
1174    #[case(
1175        false,
1176        false,
1177        true,
1178        false,
1179        AxMarketDataLevel::Level1,
1180        Some(false),
1181        Some(true)
1182    )]
1183    #[case(
1184        false,
1185        false,
1186        false,
1187        true,
1188        AxMarketDataLevel::Level1,
1189        Some(false),
1190        Some(true)
1191    )]
1192    fn test_effective_subscription_for_single_data_type(
1193        #[case] quotes: bool,
1194        #[case] trades: bool,
1195        #[case] mark_prices: bool,
1196        #[case] instrument_status: bool,
1197        #[case] level: AxMarketDataLevel,
1198        #[case] include_trades: Option<bool>,
1199        #[case] include_ticker: Option<bool>,
1200    ) {
1201        let sdt = SymbolDataTypes {
1202            quotes,
1203            trades,
1204            mark_prices,
1205            instrument_status,
1206            book_level: None,
1207        };
1208        assert_eq!(
1209            sdt.effective_subscription(),
1210            Some(AxMdSubscriptionSpec::new(
1211                level,
1212                include_trades,
1213                include_ticker,
1214            ))
1215        );
1216        assert!(!sdt.is_empty());
1217    }
1218
1219    #[rstest]
1220    #[case(false)]
1221    #[case(true)]
1222    #[tokio::test]
1223    async fn test_unsubscribe_send_failure_restores_subscription(#[case] was_pending: bool) {
1224        let client = AxMdWebSocketClient::new(
1225            "ws://localhost:9999/md/ws".to_string(),
1226            "test_token".to_string(),
1227            30,
1228            TransportBackend::default(),
1229            None,
1230        );
1231        let symbol = "EURUSD-PERP";
1232        let spec = AxMdSubscriptionSpec::new(AxMarketDataLevel::Level2, Some(false), Some(false));
1233        let topic = spec.topic(symbol);
1234        client.subscriptions.mark_subscribe(&topic);
1235        if !was_pending {
1236            client.subscriptions.confirm_subscribe(&topic);
1237        }
1238
1239        let error = client.send_unsubscribe(symbol, spec).await.unwrap_err();
1240
1241        assert_eq!(error.to_string(), "Channel error: channel closed");
1242        assert_eq!(client.subscription_count(), usize::from(!was_pending));
1243        assert_eq!(client.subscriptions.all_topics(), vec![topic]);
1244        assert_eq!(
1245            client.subscriptions.pending_subscribe_topics().len(),
1246            usize::from(was_pending)
1247        );
1248        assert!(client.subscriptions.pending_unsubscribe_topics().is_empty());
1249    }
1250
1251    #[rstest]
1252    #[case(false)]
1253    #[case(true)]
1254    #[tokio::test]
1255    async fn test_unsubscribe_candles_send_failure_restores_subscription(
1256        #[case] was_pending: bool,
1257    ) {
1258        let client = AxMdWebSocketClient::new(
1259            "ws://localhost:9999/md/ws".to_string(),
1260            "test_token".to_string(),
1261            30,
1262            TransportBackend::default(),
1263            None,
1264        );
1265        let symbol = "EURUSD-PERP";
1266        let width = AxCandleWidth::Minutes1;
1267        let topic = format!("candles:{symbol}:{width:?}");
1268        client.subscriptions.mark_subscribe(&topic);
1269        if !was_pending {
1270            client.subscriptions.confirm_subscribe(&topic);
1271        }
1272
1273        let error = client.unsubscribe_candles(symbol, width).await.unwrap_err();
1274
1275        assert_eq!(error.to_string(), "Channel error: channel closed");
1276        assert_eq!(client.subscription_count(), usize::from(!was_pending));
1277        assert_eq!(client.subscriptions.all_topics(), vec![topic]);
1278        assert_eq!(
1279            client.subscriptions.pending_subscribe_topics().len(),
1280            usize::from(was_pending)
1281        );
1282        assert!(client.subscriptions.pending_unsubscribe_topics().is_empty());
1283    }
1284
1285    #[rstest]
1286    #[tokio::test]
1287    async fn test_unsubscribe_candles_skips_untracked_topic() {
1288        let client = AxMdWebSocketClient::new(
1289            "ws://localhost:9999/md/ws".to_string(),
1290            "test_token".to_string(),
1291            30,
1292            TransportBackend::default(),
1293            None,
1294        );
1295
1296        client
1297            .unsubscribe_candles("EURUSD-PERP", AxCandleWidth::Minutes1)
1298            .await
1299            .unwrap();
1300
1301        assert_eq!(client.subscription_count(), 0);
1302        assert!(client.subscriptions.all_topics().is_empty());
1303        assert!(client.subscriptions.pending_subscribe_topics().is_empty());
1304        assert!(client.subscriptions.pending_unsubscribe_topics().is_empty());
1305    }
1306}
nautilus_architect_ax/websocket/data/client.rs

nautilus_architect_ax/websocket/data/
client.rs
