Skip to main content

nautilus_architect_ax/
data.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//! Live market data client implementation for the AX Exchange adapter.
17
18use std::{
19    future::Future,
20    sync::{
21        Arc,
22        atomic::{AtomicBool, Ordering},
23    },
24    time::Duration,
25};
26
27use ahash::{AHashMap, AHashSet};
28use anyhow::Context;
29use async_trait::async_trait;
30use futures_util::StreamExt;
31use jiff::{SignedDuration, Timestamp};
32use nautilus_common::{
33    clients::DataClient,
34    live::runner::get_data_event_sender,
35    messages::{
36        DataEvent, DataResponse,
37        data::{
38            BarsResponse, BookResponse, FundingRatesResponse, InstrumentResponse,
39            InstrumentsResponse, RequestBars, RequestBookSnapshot, RequestFundingRates,
40            RequestInstrument, RequestInstruments, RequestTrades, SubscribeBars,
41            SubscribeBookDeltas, SubscribeFundingRates, SubscribeIndexPrices, SubscribeInstrument,
42            SubscribeInstrumentClose, SubscribeInstrumentStatus, SubscribeInstruments,
43            SubscribeMarkPrices, SubscribeQuotes, SubscribeTrades, TradesResponse, UnsubscribeBars,
44            UnsubscribeBookDeltas, UnsubscribeFundingRates, UnsubscribeIndexPrices,
45            UnsubscribeInstrument, UnsubscribeInstrumentClose, UnsubscribeInstrumentStatus,
46            UnsubscribeInstruments, UnsubscribeMarkPrices, UnsubscribeQuotes, UnsubscribeTrades,
47        },
48    },
49};
50use nautilus_core::{
51    AtomicMap,
52    datetime::datetime_to_unix_nanos,
53    nanos::UnixNanos,
54    time::{AtomicTime, get_atomic_clock_realtime},
55};
56use nautilus_live::{
57    SocketControl,
58    task::{TaskGroup, TaskGroupGuard},
59};
60use nautilus_model::{
61    data::{Data, FundingRateUpdate, InstrumentStatus, MarkPriceUpdate},
62    enums::{BookType, MarketStatusAction},
63    identifiers::{ClientId, InstrumentId, Venue},
64    instruments::{Instrument, InstrumentAny},
65    types::Price,
66};
67use parking_lot::Mutex;
68use tokio_util::sync::CancellationToken;
69use ustr::Ustr;
70
71use crate::{
72    common::{
73        auth::run_auth_token_refresh,
74        consts::{AX_AUTH_TOKEN_TTL_SECS, AX_FUNDING_RATE_LOOKBACK_DAYS, AX_VENUE},
75        credential::Credential,
76        enums::{AxCandleWidth, AxInstrumentState, AxMarketDataLevel},
77        parse::{ax_timestamp_stn_to_unix_nanos, map_bar_spec_to_candle_width},
78    },
79    config::AxDataClientConfig,
80    http::client::AxHttpClient,
81    websocket::{
82        data::{
83            client::{AxMdWebSocketClient, AxWsClientError, SymbolDataTypes},
84            parse::{
85                parse_book_l1_quote, parse_book_l2_deltas, parse_book_l2_quote,
86                parse_book_l3_deltas, parse_book_l3_quote, parse_candle_bar, parse_trade_tick,
87            },
88        },
89        messages::{AxDataWsMessage, AxMdCandle, AxMdMessage},
90    },
91};
92
93/// AX Exchange data client for live market data streaming and historical data requests.
94///
95/// This client integrates with the Nautilus DataEngine to provide:
96/// - Real-time market data via WebSocket subscriptions
97/// - Historical data via REST API requests
98/// - Automatic instrument discovery and caching
99/// - Connection lifecycle management
100#[derive(Debug)]
101pub struct AxDataClient {
102    client_id: ClientId,
103    config: AxDataClientConfig,
104    http_client: AxHttpClient,
105    ws_client: AxMdWebSocketClient,
106    is_connected: Arc<AtomicBool>,
107    cancellation_token: CancellationToken,
108    session_tasks: TaskGroup,
109    pending_tasks: TaskGroup,
110    shutdown_errors: Vec<String>,
111    data_sender: tokio::sync::mpsc::UnboundedSender<DataEvent>,
112    instruments: Arc<AtomicMap<Ustr, InstrumentAny>>,
113    clock: &'static AtomicTime,
114    funding_rate_cancellations: AHashMap<InstrumentId, CancellationToken>,
115    funding_rate_cache: Arc<Mutex<AHashMap<InstrumentId, FundingRateUpdate>>>,
116}
117
118impl AxDataClient {
119    /// Creates a new [`AxDataClient`] instance.
120    ///
121    /// # Errors
122    ///
123    /// Returns an error if the data event sender cannot be obtained.
124    pub fn new(
125        client_id: ClientId,
126        config: AxDataClientConfig,
127        http_client: AxHttpClient,
128        ws_client: AxMdWebSocketClient,
129    ) -> anyhow::Result<Self> {
130        let clock = get_atomic_clock_realtime();
131        let data_sender = get_data_event_sender();
132        let ws_client = ws_client.with_socket_control(SocketControl::new(
133            client_id,
134            Some(*AX_VENUE),
135            "architect-ax-data-streams",
136        ));
137
138        // Share instruments cache with HTTP client
139        let instruments = http_client.instruments_cache.clone();
140
141        let session_tasks = TaskGroup::new();
142        let pending_tasks = TaskGroup::new();
143
144        Ok(Self {
145            client_id,
146            config,
147            http_client,
148            ws_client,
149            is_connected: Arc::new(AtomicBool::new(false)),
150            cancellation_token: CancellationToken::new(),
151            session_tasks,
152            pending_tasks,
153            shutdown_errors: Vec::new(),
154            data_sender,
155            instruments,
156            clock,
157            funding_rate_cancellations: AHashMap::new(),
158            funding_rate_cache: Arc::new(Mutex::new(AHashMap::new())),
159        })
160    }
161
162    /// Returns the venue for this data client.
163    #[must_use]
164    pub fn venue(&self) -> Venue {
165        *AX_VENUE
166    }
167
168    fn map_book_type_to_market_data_level(book_type: BookType) -> AxMarketDataLevel {
169        match book_type {
170            BookType::L3_MBO => AxMarketDataLevel::Level3,
171            BookType::L1_MBP | BookType::L2_MBP => AxMarketDataLevel::Level2,
172        }
173    }
174
175    /// Returns a reference to the instruments cache.
176    #[must_use]
177    pub fn instruments(&self) -> &Arc<AtomicMap<Ustr, InstrumentAny>> {
178        &self.instruments
179    }
180
181    /// Spawns a message handler task to forward WebSocket data to the DataEngine.
182    fn spawn_message_handler(&mut self) -> anyhow::Result<()> {
183        let stream = self.ws_client.stream();
184        let data_sender = self.data_sender.clone();
185        let cancellation_token = self.cancellation_token.clone();
186        let is_connected = Arc::clone(&self.is_connected);
187        let instruments = Arc::clone(&self.instruments);
188        let symbol_data_types = self.ws_client.symbol_data_types();
189        let status_invalidations = self.ws_client.status_invalidations();
190        let clock = self.clock;
191
192        self.session_tasks.spawn(async move {
193            tokio::pin!(stream);
194
195            let mut book_sequences: AHashMap<Ustr, u64> = AHashMap::new();
196            let mut candle_cache: AHashMap<(Ustr, AxCandleWidth), AxMdCandle> = AHashMap::new();
197            let mut instrument_states: AHashMap<Ustr, AxInstrumentState> = AHashMap::new();
198
199            loop {
200                tokio::select! {
201                    () = cancellation_token.cancelled() => {
202                        log::debug!("Message handler cancelled");
203                        break;
204                    }
205                    msg = stream.next() => {
206                        match msg {
207                            Some(ws_msg) => {
208                                drain_status_invalidations(
209                                    &status_invalidations,
210                                    &mut instrument_states,
211                                );
212
213                                handle_ws_message(
214                                    ws_msg,
215                                    &data_sender,
216                                    &instruments,
217                                    &symbol_data_types,
218                                    &mut book_sequences,
219                                    &mut candle_cache,
220                                    &mut instrument_states,
221                                    clock,
222                                );
223                            }
224                            None => {
225                                log::debug!("WebSocket stream ended");
226                                is_connected.store(false, Ordering::Release);
227                                break;
228                            }
229                        }
230                    }
231                }
232            }
233        })?;
234        Ok(())
235    }
236
237    fn spawn_instrument_refresh(&self) -> anyhow::Result<()> {
238        let minutes = self.config.update_instruments_interval_mins;
239        if minutes == 0 {
240            return Ok(());
241        }
242
243        let interval = Duration::from_secs(minutes.saturating_mul(60));
244        let cancellation = self.cancellation_token.clone();
245        let instruments_cache = Arc::clone(&self.instruments);
246        let http_client = self.http_client.clone();
247        let data_sender = self.data_sender.clone();
248        let client_id = self.client_id;
249
250        self.session_tasks.spawn(async move {
251            loop {
252                let sleep = tokio::time::sleep(interval);
253                tokio::pin!(sleep);
254                tokio::select! {
255                    () = cancellation.cancelled() => {
256                        log::debug!("Instrument refresh task cancelled");
257                        break;
258                    }
259                    () = &mut sleep => {
260                        match http_client.request_instruments(None, None).await {
261                            Ok(instruments) => {
262                                for inst in &instruments {
263                                    instruments_cache.insert(inst.symbol().inner(), inst.clone());
264
265                                    if let Err(e) = data_sender
266                                        .send(DataEvent::Instrument(inst.clone()))
267                                    {
268                                        log::warn!("Failed to send refreshed instrument: {e}");
269                                    }
270                                }
271                                http_client.cache_instruments(&instruments);
272                                log::debug!(
273                                    "Instruments refreshed: client_id={client_id}, count={}",
274                                    instruments.len(),
275                                );
276                            }
277                            Err(e) => {
278                                log::warn!("Failed to refresh instruments: client_id={client_id}, error={e:?}");
279                            }
280                        }
281                    }
282                }
283            }
284        })?;
285        Ok(())
286    }
287
288    #[expect(
289        clippy::unnecessary_wraps,
290        reason = "callers forward Result to trait methods"
291    )]
292    fn ws_symbol_op<F, Fut>(
293        &self,
294        instrument_id: InstrumentId,
295        op: F,
296        context: &'static str,
297    ) -> anyhow::Result<()>
298    where
299        F: FnOnce(AxMdWebSocketClient, String) -> Fut + Send + 'static,
300        Fut: Future<Output = Result<(), AxWsClientError>> + Send,
301    {
302        let symbol = instrument_id.symbol.to_string();
303        log::debug!("{context} for {symbol}");
304
305        let ws = self.ws_client.clone();
306        self.spawn_ws(
307            async move { op(ws, symbol).await.map_err(|e| anyhow::anyhow!(e)) },
308            context,
309        );
310
311        Ok(())
312    }
313
314    fn spawn_ws<F>(&self, fut: F, context: &'static str)
315    where
316        F: Future<Output = anyhow::Result<()>> + Send + 'static,
317    {
318        let future = async move {
319            if let Err(e) = fut.await {
320                log::error!("{context}: {e:?}");
321            }
322        };
323
324        if let Err(e) = self.pending_tasks.spawn(future) {
325            log::warn!("Skipping AX {context} after shutdown began: {e}");
326        }
327    }
328
329    fn spawn_task<F>(&self, fut: F)
330    where
331        F: Future<Output = ()> + Send + 'static,
332    {
333        if let Err(e) = self.pending_tasks.spawn(fut) {
334            log::warn!("Skipping AX data task after shutdown began: {e}");
335        }
336    }
337
338    fn abort_pending_tasks(&self) {
339        self.pending_tasks.begin_shutdown();
340    }
341
342    fn abort_all_tasks(&self) {
343        self.cancellation_token.cancel();
344        self.session_tasks.begin_shutdown();
345        self.abort_pending_tasks();
346        self.ws_client.begin_shutdown();
347
348        for cancellation in self.funding_rate_cancellations.values() {
349            cancellation.cancel();
350        }
351    }
352
353    async fn finish_all_tasks(&mut self) -> anyhow::Result<()> {
354        self.pending_tasks.begin_shutdown();
355        self.session_tasks.begin_shutdown();
356        let (pending_result, session_result) = tokio::join!(
357            self.pending_tasks
358                .finish_shutdown(Duration::from_secs(1), Duration::from_secs(2)),
359            self.session_tasks
360                .finish_shutdown(Duration::from_secs(1), Duration::from_secs(2)),
361        );
362        self.funding_rate_cancellations.clear();
363
364        pending_result.map_err(|e| anyhow::anyhow!("Failed to terminate AX data tasks: {e}"))?;
365        session_result
366            .map_err(|e| anyhow::anyhow!("Failed to terminate AX data session tasks: {e}"))?;
367        Ok(())
368    }
369
370    async fn teardown_partial_connect(&mut self) -> anyhow::Result<()> {
371        self.abort_all_tasks();
372
373        if let Err(e) = self.ws_client.close().await {
374            self.shutdown_errors.push(e.to_string());
375        }
376
377        if let Err(e) = self.finish_all_tasks().await {
378            self.shutdown_errors.push(e.to_string());
379        }
380        self.is_connected.store(false, Ordering::Release);
381
382        if !self.shutdown_errors.is_empty() {
383            anyhow::bail!(std::mem::take(&mut self.shutdown_errors).join("; "));
384        }
385        Ok(())
386    }
387}
388
389#[async_trait(?Send)]
390impl DataClient for AxDataClient {
391    fn client_id(&self) -> ClientId {
392        self.client_id
393    }
394
395    fn venue(&self) -> Option<Venue> {
396        Some(*AX_VENUE)
397    }
398
399    fn start(&mut self) -> anyhow::Result<()> {
400        log::debug!("Starting {}", self.client_id);
401        Ok(())
402    }
403
404    fn stop(&mut self) -> anyhow::Result<()> {
405        log::debug!("Stopping {}", self.client_id);
406
407        self.abort_all_tasks();
408        self.is_connected.store(false, Ordering::Release);
409        Ok(())
410    }
411
412    fn reset(&mut self) -> anyhow::Result<()> {
413        log::debug!("Resetting {}", self.client_id);
414
415        self.abort_all_tasks();
416        self.is_connected.store(false, Ordering::Release);
417        self.funding_rate_cache.lock().clear();
418        Ok(())
419    }
420
421    fn dispose(&mut self) -> anyhow::Result<()> {
422        log::debug!("Disposing {}", self.client_id);
423
424        self.abort_all_tasks();
425        self.is_connected.store(false, Ordering::Release);
426        Ok(())
427    }
428
429    fn is_connected(&self) -> bool {
430        self.is_connected.load(Ordering::Acquire)
431    }
432
433    fn is_disconnected(&self) -> bool {
434        !self.is_connected()
435    }
436
437    async fn connect(&mut self) -> anyhow::Result<()> {
438        if self.is_connected()
439            && !self.cancellation_token.is_cancelled()
440            && self.pending_tasks.is_open()
441            && self.session_tasks.is_open()
442        {
443            log::debug!("Already connected {}", self.client_id);
444            return Ok(());
445        }
446
447        log::info!("Connecting {}", self.client_id);
448
449        if self.cancellation_token.is_cancelled()
450            || !self.pending_tasks.is_open()
451            || !self.session_tasks.is_open()
452            || !self.funding_rate_cancellations.is_empty()
453        {
454            self.teardown_partial_connect().await?;
455            self.session_tasks
456                .start_generation()
457                .map_err(|e| anyhow::anyhow!("Failed to start AX data session generation: {e}"))?;
458            self.pending_tasks
459                .start_generation()
460                .map_err(|e| anyhow::anyhow!("Failed to start AX data task generation: {e}"))?;
461            self.cancellation_token = CancellationToken::new();
462        }
463        let cancellation_token = self.cancellation_token.clone();
464        let ws_client = self.ws_client.clone();
465        let setup_guard =
466            TaskGroupGuard::new(&[&self.session_tasks, &self.pending_tasks], move || {
467                cancellation_token.cancel();
468                ws_client.begin_shutdown();
469            });
470
471        let credential = if self.config.has_api_credentials() {
472            let credential =
473                Credential::resolve(self.config.api_key.clone(), self.config.api_secret.clone())
474                    .context("API credentials not configured")?;
475
476            let token = self
477                .http_client
478                .authenticate(
479                    credential.api_key(),
480                    credential.api_secret(),
481                    AX_AUTH_TOKEN_TTL_SECS,
482                )
483                .await
484                .context("Failed to authenticate with Ax")?;
485            log::debug!("Authenticated with Ax");
486            self.ws_client.set_auth_token(token);
487
488            // Only an authenticated client can read fee rates, and a data client may
489            // legitimately run without credentials.
490            self.http_client
491                .request_account_fees()
492                .await
493                .context("Failed to resolve Ax account fee rates")?;
494
495            Some(credential)
496        } else {
497            log::debug!("No Ax credentials configured, instruments will report zero fees");
498            None
499        };
500
501        let instruments = self
502            .http_client
503            .request_instruments(None, None)
504            .await
505            .context("Failed to fetch instruments")?;
506
507        for instrument in &instruments {
508            self.instruments
509                .insert(instrument.symbol().inner(), instrument.clone());
510
511            if let Err(e) = self
512                .data_sender
513                .send(DataEvent::Instrument(instrument.clone()))
514            {
515                log::warn!("Failed to send instrument: {e}");
516            }
517        }
518        self.http_client.cache_instruments(&instruments);
519        log::debug!(
520            "Cached {} instruments",
521            self.http_client.get_cached_symbols().len()
522        );
523
524        self.ws_client
525            .connect()
526            .await
527            .context("Failed to connect WebSocket")?;
528        log::debug!("WebSocket connected");
529
530        let session_result = async {
531            self.spawn_message_handler()?;
532            self.spawn_instrument_refresh()?;
533
534            if let Some(credential) = credential {
535                let ws_client = self.ws_client.clone();
536                self.session_tasks.spawn(run_auth_token_refresh(
537                    self.http_client.clone(),
538                    credential,
539                    move |token| ws_client.update_auth_token(token),
540                ))?;
541            }
542            Ok::<(), anyhow::Error>(())
543        }
544        .await;
545
546        if let Err(e) = session_result {
547            if let Err(teardown_error) = self.teardown_partial_connect().await {
548                return Err(e.context(format!("AX data startup teardown failed: {teardown_error}")));
549            }
550            return Err(e);
551        }
552
553        self.is_connected.store(true, Ordering::Release);
554        setup_guard.disarm();
555        log::info!("Connected {}", self.client_id);
556
557        Ok(())
558    }
559
560    async fn disconnect(&mut self) -> anyhow::Result<()> {
561        log::info!("Disconnecting {}", self.client_id);
562
563        self.abort_all_tasks();
564        let ws_result = self.ws_client.close().await;
565        let tasks_result = self.finish_all_tasks().await;
566        self.funding_rate_cache.lock().clear();
567
568        self.is_connected.store(false, Ordering::Release);
569        log::info!("Disconnected {}", self.client_id);
570
571        ws_result?;
572        tasks_result
573    }
574
575    fn subscribe_instruments(&mut self, _cmd: SubscribeInstruments) -> anyhow::Result<()> {
576        // AX does not have a real-time instruments channel; instruments are fetched via HTTP
577        log::debug!("Instruments subscription not applicable for AX (use request_instruments)");
578        Ok(())
579    }
580
581    fn subscribe_instrument(&mut self, _cmd: SubscribeInstrument) -> anyhow::Result<()> {
582        // AX does not have a real-time instrument channel; instruments are fetched via HTTP
583        log::debug!("Instrument subscription not applicable for AX (use request_instrument)");
584        Ok(())
585    }
586
587    fn subscribe_book_deltas(&mut self, cmd: SubscribeBookDeltas) -> anyhow::Result<()> {
588        let symbol = cmd.instrument_id.symbol.to_string();
589        let level = Self::map_book_type_to_market_data_level(cmd.book_type);
590        if cmd.book_type == BookType::L1_MBP {
591            log::warn!(
592                "Book type L1_MBP not supported by AX for deltas, downgrading {symbol} to LEVEL_2"
593            );
594        }
595        log::debug!("Subscribing to book deltas for {symbol} at {level:?}");
596
597        let ws = self.ws_client.clone();
598        self.spawn_ws(
599            async move {
600                ws.subscribe_book_deltas(&symbol, level)
601                    .await
602                    .map_err(|e| anyhow::anyhow!(e))
603            },
604            "subscribe book deltas",
605        );
606
607        Ok(())
608    }
609
610    fn subscribe_quotes(&mut self, cmd: SubscribeQuotes) -> anyhow::Result<()> {
611        self.ws_symbol_op(
612            cmd.instrument_id,
613            |ws, s| async move { ws.subscribe_quotes(&s).await },
614            "Subscribing to quotes",
615        )
616    }
617
618    fn subscribe_trades(&mut self, cmd: SubscribeTrades) -> anyhow::Result<()> {
619        self.ws_symbol_op(
620            cmd.instrument_id,
621            |ws, s| async move { ws.subscribe_trades(&s).await },
622            "Subscribing to trades",
623        )
624    }
625
626    fn subscribe_mark_prices(&mut self, cmd: SubscribeMarkPrices) -> anyhow::Result<()> {
627        self.ws_symbol_op(
628            cmd.instrument_id,
629            |ws, s| async move { ws.subscribe_mark_prices(&s).await },
630            "Subscribing to mark prices",
631        )
632    }
633
634    fn subscribe_index_prices(&mut self, _cmd: SubscribeIndexPrices) -> anyhow::Result<()> {
635        log::warn!("Index prices not supported by AX Exchange");
636        Ok(())
637    }
638
639    fn subscribe_bars(&mut self, cmd: SubscribeBars) -> anyhow::Result<()> {
640        let bar_type = cmd.bar_type;
641        let symbol = bar_type.instrument_id().symbol.to_string();
642        let width = map_bar_spec_to_candle_width(&bar_type.spec())?;
643        log::debug!("Subscribing to bars for {bar_type} (width: {width:?})");
644
645        let ws = self.ws_client.clone();
646        self.spawn_ws(
647            async move {
648                ws.subscribe_candles(&symbol, width)
649                    .await
650                    .map_err(|e| anyhow::anyhow!(e))
651            },
652            "subscribe bars",
653        );
654
655        Ok(())
656    }
657
658    fn subscribe_funding_rates(&mut self, cmd: SubscribeFundingRates) -> anyhow::Result<()> {
659        let poll_interval_mins = self.config.funding_rate_poll_interval_mins.max(1);
660
661        // Use 7-day lookback to capture latest rate across weekends/holidays
662        let lookback = SignedDuration::from_hours(24 * (AX_FUNDING_RATE_LOOKBACK_DAYS));
663
664        let instrument_id = cmd.instrument_id;
665
666        if self.funding_rate_cancellations.contains_key(&instrument_id) {
667            log::debug!("Already subscribed to funding rates for {instrument_id}");
668            return Ok(());
669        }
670
671        log::debug!("Subscribing to funding rates for {instrument_id} (HTTP polling)");
672
673        let http = self.http_client.clone();
674        let sender = self.data_sender.clone();
675        let symbol = instrument_id.symbol.inner();
676        let cancellation = self.cancellation_token.child_token();
677        let task_cancellation = cancellation.clone();
678        let cache = Arc::clone(&self.funding_rate_cache);
679        let clock = self.clock;
680
681        self.session_tasks.spawn(async move {
682            // First tick fires immediately for initial emission
683            let mut interval = tokio::time::interval(Duration::from_mins(poll_interval_mins));
684
685            loop {
686                tokio::select! {
687                    () = task_cancellation.cancelled() => {
688                        log::debug!("Funding rate polling cancelled for {symbol}");
689                        break;
690                    }
691                    _ = interval.tick() => {
692                        let now: Timestamp = clock.get_time_ns().into();
693                        let start = now - lookback;
694
695                        match http.request_funding_rates(instrument_id, Some(start), Some(now)).await {
696                            Ok(funding_rates) => {
697                                if funding_rates.is_empty() {
698                                    log::warn!(
699                                        "No funding rates returned for {symbol}"
700                                    );
701                                } else if let Some(update) = funding_rates.last() {
702                                    // Only emit if rate changed
703                                    let should_emit = cache.lock()
704                                        .get(&instrument_id) != Some(update);
705
706                                    if should_emit {
707                                        log::debug!(
708                                            "Funding rate for {symbol}: {}",
709                                            update.rate,
710                                        );
711                                        let update = *update;
712                                        cache.lock()
713                                            .insert(instrument_id, update);
714
715                                        if let Err(e) = sender.send(
716                                            DataEvent::FundingRate(update),
717                                        ) {
718                                            log::error!(
719                                                "Failed to send funding rate for {symbol}: {e}"
720                                            );
721                                        }
722                                    }
723                                }
724                            }
725                            Err(e) => {
726                                log::error!(
727                                    "Failed to poll funding rates for {symbol}: {e}"
728                                );
729                            }
730                        }
731                    }
732                }
733            }
734        })?;
735
736        self.funding_rate_cancellations
737            .insert(instrument_id, cancellation);
738        Ok(())
739    }
740
741    fn subscribe_instrument_status(
742        &mut self,
743        cmd: SubscribeInstrumentStatus,
744    ) -> anyhow::Result<()> {
745        self.ws_symbol_op(
746            cmd.instrument_id,
747            |ws, s| async move { ws.subscribe_instrument_status(&s).await },
748            "Subscribing to instrument status",
749        )
750    }
751
752    fn subscribe_instrument_close(&mut self, _cmd: SubscribeInstrumentClose) -> anyhow::Result<()> {
753        log::warn!("Instrument close not supported by AX Exchange");
754        Ok(())
755    }
756
757    fn unsubscribe_instruments(&mut self, _cmd: &UnsubscribeInstruments) -> anyhow::Result<()> {
758        Ok(())
759    }
760
761    fn unsubscribe_instrument(&mut self, _cmd: &UnsubscribeInstrument) -> anyhow::Result<()> {
762        Ok(())
763    }
764
765    fn unsubscribe_book_deltas(&mut self, cmd: &UnsubscribeBookDeltas) -> anyhow::Result<()> {
766        self.ws_symbol_op(
767            cmd.instrument_id,
768            |ws, s| async move { ws.unsubscribe_book_deltas(&s).await },
769            "Unsubscribing from book deltas",
770        )
771    }
772
773    fn unsubscribe_quotes(&mut self, cmd: &UnsubscribeQuotes) -> anyhow::Result<()> {
774        self.ws_symbol_op(
775            cmd.instrument_id,
776            |ws, s| async move { ws.unsubscribe_quotes(&s).await },
777            "Unsubscribing from quotes",
778        )
779    }
780
781    fn unsubscribe_trades(&mut self, cmd: &UnsubscribeTrades) -> anyhow::Result<()> {
782        self.ws_symbol_op(
783            cmd.instrument_id,
784            |ws, s| async move { ws.unsubscribe_trades(&s).await },
785            "Unsubscribing from trades",
786        )
787    }
788
789    fn unsubscribe_mark_prices(&mut self, cmd: &UnsubscribeMarkPrices) -> anyhow::Result<()> {
790        self.ws_symbol_op(
791            cmd.instrument_id,
792            |ws, s| async move { ws.unsubscribe_mark_prices(&s).await },
793            "Unsubscribing from mark prices",
794        )
795    }
796
797    fn unsubscribe_index_prices(&mut self, _cmd: &UnsubscribeIndexPrices) -> anyhow::Result<()> {
798        Ok(())
799    }
800
801    fn unsubscribe_bars(&mut self, cmd: &UnsubscribeBars) -> anyhow::Result<()> {
802        let bar_type = cmd.bar_type;
803        let symbol = bar_type.instrument_id().symbol.to_string();
804        let width = map_bar_spec_to_candle_width(&bar_type.spec())?;
805        log::debug!("Unsubscribing from bars for {bar_type}");
806
807        let ws = self.ws_client.clone();
808        self.spawn_ws(
809            async move {
810                ws.unsubscribe_candles(&symbol, width)
811                    .await
812                    .map_err(|e| anyhow::anyhow!(e))
813            },
814            "unsubscribe bars",
815        );
816
817        Ok(())
818    }
819
820    fn unsubscribe_funding_rates(&mut self, cmd: &UnsubscribeFundingRates) -> anyhow::Result<()> {
821        let instrument_id = cmd.instrument_id;
822
823        if let Some(cancellation) = self.funding_rate_cancellations.remove(&instrument_id) {
824            log::debug!("Unsubscribing from funding rates for {instrument_id}");
825            cancellation.cancel();
826            self.funding_rate_cache.lock().remove(&instrument_id);
827        } else {
828            log::debug!("Not subscribed to funding rates for {instrument_id}");
829        }
830
831        Ok(())
832    }
833
834    fn unsubscribe_instrument_status(
835        &mut self,
836        cmd: &UnsubscribeInstrumentStatus,
837    ) -> anyhow::Result<()> {
838        self.ws_symbol_op(
839            cmd.instrument_id,
840            |ws, s| async move { ws.unsubscribe_instrument_status(&s).await },
841            "Unsubscribing from instrument status",
842        )
843    }
844
845    fn unsubscribe_instrument_close(
846        &mut self,
847        _cmd: &UnsubscribeInstrumentClose,
848    ) -> anyhow::Result<()> {
849        Ok(())
850    }
851
852    fn request_instruments(&self, request: RequestInstruments) -> anyhow::Result<()> {
853        let http = self.http_client.clone();
854        let instruments_cache = Arc::clone(&self.instruments);
855        let sender = self.data_sender.clone();
856        let cancel = self.cancellation_token.clone();
857        let request_id = request.request_id;
858        let client_id = request.client_id.unwrap_or(self.client_id);
859        let venue = *AX_VENUE;
860        let start_nanos = datetime_to_unix_nanos(request.start);
861        let end_nanos = datetime_to_unix_nanos(request.end);
862        let params = request.params;
863        let clock = self.clock;
864
865        self.spawn_task(async move {
866            match http.request_instruments(None, None).await {
867                Ok(instruments) => {
868                    if cancel.is_cancelled() {
869                        return;
870                    }
871                    log::debug!("Fetched {} instruments from Ax", instruments.len());
872                    for inst in &instruments {
873                        instruments_cache.insert(inst.symbol().inner(), inst.clone());
874                    }
875                    http.cache_instruments(&instruments);
876
877                    let response = DataResponse::Instruments(InstrumentsResponse::new(
878                        request_id,
879                        client_id,
880                        venue,
881                        instruments,
882                        start_nanos,
883                        end_nanos,
884                        clock.get_time_ns(),
885                        params,
886                    ));
887
888                    if let Err(e) = sender.send(DataEvent::Response(response)) {
889                        log::error!("Failed to send instruments response: {e}");
890                    }
891                }
892                Err(e) => {
893                    log::error!("Failed to request instruments: {e}");
894                }
895            }
896        });
897
898        Ok(())
899    }
900
901    fn request_instrument(&self, request: RequestInstrument) -> anyhow::Result<()> {
902        let http = self.http_client.clone();
903        let instruments_cache = Arc::clone(&self.instruments);
904        let sender = self.data_sender.clone();
905        let cancel = self.cancellation_token.clone();
906        let request_id = request.request_id;
907        let client_id = request.client_id.unwrap_or(self.client_id);
908        let instrument_id = request.instrument_id;
909        let symbol = instrument_id.symbol.inner();
910        let start_nanos = datetime_to_unix_nanos(request.start);
911        let end_nanos = datetime_to_unix_nanos(request.end);
912        let params = request.params;
913        let clock = self.clock;
914
915        self.spawn_task(async move {
916            match http.request_instrument(symbol, None, None).await {
917                Ok(instrument) => {
918                    if cancel.is_cancelled() {
919                        return;
920                    }
921                    log::debug!("Fetched instrument {symbol} from Ax");
922                    instruments_cache.insert(symbol, instrument.clone());
923                    http.cache_instrument(instrument.clone());
924
925                    let response = DataResponse::Instrument(Box::new(InstrumentResponse::new(
926                        request_id,
927                        client_id,
928                        instrument_id,
929                        instrument,
930                        start_nanos,
931                        end_nanos,
932                        clock.get_time_ns(),
933                        params,
934                    )));
935
936                    if let Err(e) = sender.send(DataEvent::Response(response)) {
937                        log::error!("Failed to send instrument response: {e}");
938                    }
939                }
940                Err(e) => {
941                    log::error!("Failed to request instrument {symbol}: {e}");
942                }
943            }
944        });
945
946        Ok(())
947    }
948
949    fn request_book_snapshot(&self, request: RequestBookSnapshot) -> anyhow::Result<()> {
950        let http = self.http_client.clone();
951        let sender = self.data_sender.clone();
952        let cancel = self.cancellation_token.clone();
953        let request_id = request.request_id;
954        let client_id = request.client_id.unwrap_or(self.client_id);
955        let instrument_id = request.instrument_id;
956        let symbol = instrument_id.symbol.inner();
957        let depth = request.depth.map(|n| n.get());
958        let params = request.params;
959        let clock = self.clock;
960
961        self.spawn_task(async move {
962            match http.request_book_snapshot(symbol, depth).await {
963                Ok(book) => {
964                    if cancel.is_cancelled() {
965                        return;
966                    }
967                    log::debug!(
968                        "Fetched book snapshot for {symbol} ({} bids, {} asks)",
969                        book.bids(None).count(),
970                        book.asks(None).count(),
971                    );
972
973                    let response = DataResponse::Book(BookResponse::new(
974                        request_id,
975                        client_id,
976                        instrument_id,
977                        book,
978                        None,
979                        None,
980                        clock.get_time_ns(),
981                        params,
982                    ));
983
984                    if let Err(e) = sender.send(DataEvent::Response(response)) {
985                        log::error!("Failed to send book snapshot response: {e}");
986                    }
987                }
988                Err(e) => {
989                    log::error!("Failed to request book snapshot for {symbol}: {e}");
990                }
991            }
992        });
993
994        Ok(())
995    }
996
997    fn request_trades(&self, request: RequestTrades) -> anyhow::Result<()> {
998        let http = self.http_client.clone();
999        let sender = self.data_sender.clone();
1000        let cancel = self.cancellation_token.clone();
1001        let request_id = request.request_id;
1002        let client_id = request.client_id.unwrap_or(self.client_id);
1003        let instrument_id = request.instrument_id;
1004        let symbol = instrument_id.symbol.inner();
1005        let limit = request.limit.map(|n| n.get() as i32);
1006        let start_nanos = datetime_to_unix_nanos(request.start);
1007        let end_nanos = datetime_to_unix_nanos(request.end);
1008        let params = request.params;
1009        let clock = self.clock;
1010
1011        self.spawn_task(async move {
1012            match http
1013                .request_trade_ticks(symbol, limit, start_nanos, end_nanos)
1014                .await
1015            {
1016                Ok(ticks) => {
1017                    if cancel.is_cancelled() {
1018                        return;
1019                    }
1020                    log::debug!("Fetched {} trades for {symbol}", ticks.len());
1021
1022                    let response = DataResponse::Trades(TradesResponse::new(
1023                        request_id,
1024                        client_id,
1025                        instrument_id,
1026                        ticks,
1027                        start_nanos,
1028                        end_nanos,
1029                        clock.get_time_ns(),
1030                        params,
1031                    ));
1032
1033                    if let Err(e) = sender.send(DataEvent::Response(response)) {
1034                        log::error!("Failed to send trades response: {e}");
1035                    }
1036                }
1037                Err(e) => {
1038                    log::error!("Failed to request trades for {symbol}: {e}");
1039                }
1040            }
1041        });
1042
1043        Ok(())
1044    }
1045
1046    fn request_bars(&self, request: RequestBars) -> anyhow::Result<()> {
1047        let http = self.http_client.clone();
1048        let sender = self.data_sender.clone();
1049        let request_id = request.request_id;
1050        let client_id = request.client_id.unwrap_or(self.client_id);
1051        let bar_type = request.bar_type;
1052        let symbol = bar_type.instrument_id().symbol.inner();
1053        let start = request.start;
1054        let end = request.end;
1055        let start_nanos = datetime_to_unix_nanos(start);
1056        let end_nanos = datetime_to_unix_nanos(end);
1057        let params = request.params;
1058        let clock = self.clock;
1059        let width = match map_bar_spec_to_candle_width(&bar_type.spec()) {
1060            Ok(w) => w,
1061            Err(e) => {
1062                log::error!("Failed to map bar type {bar_type}: {e}");
1063                return Err(e);
1064            }
1065        };
1066
1067        let cancel = self.cancellation_token.clone();
1068
1069        self.spawn_task(async move {
1070            match http.request_bars(symbol, start, end, width).await {
1071                Ok(bars) => {
1072                    if cancel.is_cancelled() {
1073                        return;
1074                    }
1075                    log::debug!("Fetched {} bars for {symbol}", bars.len());
1076
1077                    let response = DataResponse::Bars(BarsResponse::new(
1078                        request_id,
1079                        client_id,
1080                        bar_type,
1081                        bars,
1082                        start_nanos,
1083                        end_nanos,
1084                        clock.get_time_ns(),
1085                        params,
1086                    ));
1087
1088                    if let Err(e) = sender.send(DataEvent::Response(response)) {
1089                        log::error!("Failed to send bars response: {e}");
1090                    }
1091                }
1092                Err(e) => {
1093                    log::error!("Failed to request bars for {symbol}: {e}");
1094                }
1095            }
1096        });
1097
1098        Ok(())
1099    }
1100
1101    fn request_funding_rates(&self, request: RequestFundingRates) -> anyhow::Result<()> {
1102        let http = self.http_client.clone();
1103        let sender = self.data_sender.clone();
1104        let cancel = self.cancellation_token.clone();
1105        let request_id = request.request_id;
1106        let client_id = request.client_id.unwrap_or(self.client_id);
1107        let instrument_id = request.instrument_id;
1108        let symbol = instrument_id.symbol.inner();
1109        let start = request.start;
1110        let end = request.end;
1111        let start_nanos = datetime_to_unix_nanos(start);
1112        let end_nanos = datetime_to_unix_nanos(end);
1113        let params = request.params;
1114        let clock = self.clock;
1115
1116        self.spawn_task(async move {
1117            match http.request_funding_rates(instrument_id, start, end).await {
1118                Ok(funding_rates) => {
1119                    if cancel.is_cancelled() {
1120                        return;
1121                    }
1122                    log::debug!("Fetched {} funding rates for {symbol}", funding_rates.len());
1123
1124                    let ts_init = clock.get_time_ns();
1125                    let response = DataResponse::FundingRates(FundingRatesResponse::new(
1126                        request_id,
1127                        client_id,
1128                        instrument_id,
1129                        funding_rates,
1130                        start_nanos,
1131                        end_nanos,
1132                        ts_init,
1133                        params,
1134                    ));
1135
1136                    if let Err(e) = sender.send(DataEvent::Response(response)) {
1137                        log::error!("Failed to send funding rates response: {e}");
1138                    }
1139                }
1140                Err(e) => {
1141                    log::error!("Failed to request funding rates for {symbol}: {e}");
1142                }
1143            }
1144        });
1145
1146        Ok(())
1147    }
1148}
1149
1150fn drain_status_invalidations(
1151    invalidations: &Arc<Mutex<AHashSet<Ustr>>>,
1152    instrument_states: &mut AHashMap<Ustr, AxInstrumentState>,
1153) {
1154    for symbol in invalidations.lock().drain() {
1155        instrument_states.remove(&symbol);
1156    }
1157}
1158
1159#[expect(clippy::too_many_arguments)]
1160fn handle_ws_message(
1161    msg: AxDataWsMessage,
1162    sender: &tokio::sync::mpsc::UnboundedSender<DataEvent>,
1163    instruments: &Arc<AtomicMap<Ustr, InstrumentAny>>,
1164    symbol_data_types: &Arc<AtomicMap<String, SymbolDataTypes>>,
1165    book_sequences: &mut AHashMap<Ustr, u64>,
1166    candle_cache: &mut AHashMap<(Ustr, AxCandleWidth), AxMdCandle>,
1167    instrument_states: &mut AHashMap<Ustr, AxInstrumentState>,
1168    clock: &'static AtomicTime,
1169) {
1170    match msg {
1171        AxDataWsMessage::Reconnected => {
1172            candle_cache.clear();
1173            instrument_states.clear();
1174            log::info!("WebSocket reconnected");
1175        }
1176        AxDataWsMessage::CandleUnsubscribed { symbol, width } => {
1177            candle_cache.remove(&(symbol, width));
1178        }
1179        AxDataWsMessage::MdMessage(md_msg) => {
1180            handle_md_message(
1181                md_msg,
1182                sender,
1183                instruments,
1184                symbol_data_types,
1185                book_sequences,
1186                candle_cache,
1187                instrument_states,
1188                clock,
1189            );
1190        }
1191    }
1192}
1193
1194#[expect(clippy::too_many_arguments)]
1195fn handle_md_message(
1196    message: AxMdMessage,
1197    sender: &tokio::sync::mpsc::UnboundedSender<DataEvent>,
1198    instruments: &Arc<AtomicMap<Ustr, InstrumentAny>>,
1199    symbol_data_types: &Arc<AtomicMap<String, SymbolDataTypes>>,
1200    book_sequences: &mut AHashMap<Ustr, u64>,
1201    candle_cache: &mut AHashMap<(Ustr, AxCandleWidth), AxMdCandle>,
1202    instrument_states: &mut AHashMap<Ustr, AxInstrumentState>,
1203    clock: &'static AtomicTime,
1204) {
1205    let ts_init = || -> UnixNanos { clock.get_time_ns() };
1206
1207    let instruments_snap = instruments.load();
1208    let sdt_snap = symbol_data_types.load();
1209
1210    match message {
1211        AxMdMessage::BookL1(book) => {
1212            let l1_subscribed = sdt_snap
1213                .get(book.s.as_str())
1214                .is_some_and(|e| e.quotes || e.book_level == Some(AxMarketDataLevel::Level1));
1215
1216            if !l1_subscribed {
1217                return;
1218            }
1219
1220            let Some(instrument) = instruments_snap.get(&book.s) else {
1221                log::error!(
1222                    "No instrument cached for symbol '{}' - cannot parse L1 book",
1223                    book.s
1224                );
1225                return;
1226            };
1227
1228            match parse_book_l1_quote(&book, instrument, ts_init()) {
1229                Ok(quote) => {
1230                    let _ = sender.send(DataEvent::Data(Data::Quote(quote)));
1231                }
1232                Err(e) => log::error!("Failed to parse L1 to QuoteTick: {e}"),
1233            }
1234        }
1235        AxMdMessage::BookL2(book) => {
1236            let symbol = book.s;
1237            let seq = book_sequences.entry(symbol).or_insert(0);
1238            *seq += 1;
1239            let sequence = *seq;
1240
1241            let Some(instrument) = instruments_snap.get(&symbol) else {
1242                log::error!("No instrument cached for symbol '{symbol}' - cannot parse L2 book");
1243                return;
1244            };
1245
1246            match parse_book_l2_deltas(&book, instrument, sequence, ts_init()) {
1247                Ok(deltas) => {
1248                    let _ = sender.send(DataEvent::Data(Data::Deltas(Box::new(deltas))));
1249                }
1250                Err(e) => log::error!("Failed to parse L2 to OrderBookDeltas: {e}"),
1251            }
1252
1253            let quotes_subscribed = sdt_snap
1254                .get(symbol.as_str())
1255                .is_some_and(|entry| entry.quotes);
1256
1257            if quotes_subscribed {
1258                match parse_book_l2_quote(&book, instrument, ts_init()) {
1259                    Ok(quote) => {
1260                        let _ = sender.send(DataEvent::Data(Data::Quote(quote)));
1261                    }
1262                    Err(e) => log::error!("Failed to parse L2 to QuoteTick: {e}"),
1263                }
1264            }
1265        }
1266        AxMdMessage::BookL3(book) => {
1267            let symbol = book.s;
1268            let seq = book_sequences.entry(symbol).or_insert(0);
1269            *seq += 1;
1270            let sequence = *seq;
1271
1272            let Some(instrument) = instruments_snap.get(&symbol) else {
1273                log::error!("No instrument cached for symbol '{symbol}' - cannot parse L3 book");
1274                return;
1275            };
1276
1277            match parse_book_l3_deltas(&book, instrument, sequence, ts_init()) {
1278                Ok(deltas) => {
1279                    let _ = sender.send(DataEvent::Data(Data::Deltas(Box::new(deltas))));
1280                }
1281                Err(e) => log::error!("Failed to parse L3 to OrderBookDeltas: {e}"),
1282            }
1283
1284            let quotes_subscribed = sdt_snap
1285                .get(symbol.as_str())
1286                .is_some_and(|entry| entry.quotes);
1287
1288            if quotes_subscribed {
1289                match parse_book_l3_quote(&book, instrument, ts_init()) {
1290                    Ok(quote) => {
1291                        let _ = sender.send(DataEvent::Data(Data::Quote(quote)));
1292                    }
1293                    Err(e) => log::error!("Failed to parse L3 to QuoteTick: {e}"),
1294                }
1295            }
1296        }
1297        AxMdMessage::Ticker(ticker) => {
1298            let Some(instrument) = instruments_snap.get(&ticker.s) else {
1299                log::debug!("No instrument cached for ticker symbol '{}'", ticker.s);
1300                return;
1301            };
1302
1303            let instrument_id = instrument.id();
1304            let price_precision = instrument.price_precision();
1305            let ts_event =
1306                ax_timestamp_stn_to_unix_nanos(ticker.ts, ticker.tn).unwrap_or_else(|_| ts_init());
1307            let ts_init = ts_init();
1308
1309            let mark_prices_subscribed = sdt_snap
1310                .get(ticker.s.as_str())
1311                .is_some_and(|e| e.mark_prices);
1312
1313            if mark_prices_subscribed && let Some(mark_price) = ticker.m {
1314                match Price::from_decimal_dp(mark_price, price_precision) {
1315                    Ok(price) => {
1316                        let update = MarkPriceUpdate::new(instrument_id, price, ts_event, ts_init);
1317                        let _ = sender.send(DataEvent::Data(Data::MarkPrice(update)));
1318                    }
1319                    Err(e) => {
1320                        log::error!("Failed to parse mark price for {}: {e}", ticker.s);
1321                    }
1322                }
1323            }
1324
1325            if let Some(state) = ticker.i {
1326                let status_subscribed = sdt_snap
1327                    .get(ticker.s.as_str())
1328                    .is_some_and(|e| e.instrument_status);
1329
1330                if status_subscribed {
1331                    let prev = instrument_states.insert(ticker.s, state);
1332                    if prev != Some(state) {
1333                        let action = MarketStatusAction::from(state);
1334                        let status = InstrumentStatus::new(
1335                            instrument_id,
1336                            action,
1337                            ts_event,
1338                            ts_init,
1339                            None,
1340                            None,
1341                            Some(state == AxInstrumentState::Open),
1342                            None,
1343                            None,
1344                        );
1345                        let _ = sender.send(DataEvent::InstrumentStatus(status));
1346                    }
1347                }
1348            }
1349        }
1350        AxMdMessage::Trade(trade) => {
1351            let trades_subscribed = sdt_snap.get(trade.s.as_str()).is_some_and(|e| e.trades);
1352
1353            if !trades_subscribed {
1354                return;
1355            }
1356
1357            let Some(instrument) = instruments_snap.get(&trade.s) else {
1358                log::error!(
1359                    "No instrument cached for symbol '{}' - cannot parse trade",
1360                    trade.s
1361                );
1362                return;
1363            };
1364
1365            match parse_trade_tick(&trade, instrument, ts_init()) {
1366                Ok(tick) => {
1367                    let _ = sender.send(DataEvent::Data(Data::Trade(tick)));
1368                }
1369                Err(e) => log::error!("Failed to parse trade to TradeTick: {e}"),
1370            }
1371        }
1372        AxMdMessage::Candle(candle) => {
1373            let cache_key = (candle.symbol, candle.width);
1374
1375            let closed_candle = if let Some(cached) = candle_cache.get(&cache_key) {
1376                if cached.ts == candle.ts {
1377                    None
1378                } else {
1379                    Some(cached.clone())
1380                }
1381            } else {
1382                None
1383            };
1384
1385            candle_cache.insert(cache_key, candle);
1386
1387            if let Some(closed) = closed_candle {
1388                let Some(instrument) = instruments_snap.get(&closed.symbol) else {
1389                    log::error!(
1390                        "No instrument cached for symbol '{}' - cannot parse candle",
1391                        closed.symbol
1392                    );
1393                    return;
1394                };
1395
1396                match parse_candle_bar(&closed, instrument, ts_init()) {
1397                    Ok(bar) => {
1398                        let _ = sender.send(DataEvent::Data(Data::Bar(bar)));
1399                    }
1400                    Err(e) => log::error!("Failed to parse candle to Bar: {e}"),
1401                }
1402            }
1403        }
1404        AxMdMessage::Heartbeat(_) => {
1405            log::trace!("Received heartbeat");
1406        }
1407        AxMdMessage::SubscriptionResponse(_) => {}
1408        AxMdMessage::Error(error) => {
1409            log::warn!("WebSocket error: {}", error.message);
1410        }
1411    }
1412}
1413
1414#[cfg(test)]
1415mod tests {
1416    use std::sync::Arc;
1417
1418    use ahash::{AHashMap, AHashSet};
1419    use nautilus_model::{
1420        data::InstrumentStatus,
1421        enums::AssetClass,
1422        identifiers::{InstrumentId, Symbol},
1423        instruments::PerpetualContract,
1424        types::{Currency, Price, Quantity},
1425    };
1426    use parking_lot::Mutex;
1427    use rstest::rstest;
1428    use rust_decimal::Decimal;
1429    use rust_decimal_macros::dec;
1430    use ustr::Ustr;
1431
1432    use super::*;
1433    use crate::websocket::{
1434        data::client::SymbolDataTypes,
1435        messages::{AxBookLevel, AxMdBookL2, AxMdMessage, AxMdTicker},
1436    };
1437
1438    #[rstest]
1439    fn test_drain_status_invalidations_removes_cached_state() {
1440        let invalidations = Arc::new(Mutex::new(AHashSet::new()));
1441        let mut states = AHashMap::new();
1442        let sym = Ustr::from("EURUSD-PERP");
1443
1444        states.insert(sym, AxInstrumentState::Open);
1445        invalidations.lock().insert(sym);
1446
1447        drain_status_invalidations(&invalidations, &mut states);
1448
1449        assert!(!states.contains_key(&sym));
1450        assert!(invalidations.lock().is_empty());
1451    }
1452
1453    #[rstest]
1454    fn test_drain_status_invalidations_no_op_when_empty() {
1455        let invalidations = Arc::new(Mutex::new(AHashSet::new()));
1456        let mut states = AHashMap::new();
1457        let sym = Ustr::from("EURUSD-PERP");
1458        states.insert(sym, AxInstrumentState::Open);
1459
1460        drain_status_invalidations(&invalidations, &mut states);
1461
1462        assert!(states.contains_key(&sym));
1463    }
1464
1465    fn ticker_test_instrument() -> InstrumentAny {
1466        let symbol = Symbol::new("EURUSD-PERP");
1467        let instrument = PerpetualContract::builder()
1468            .instrument_id(InstrumentId::new(symbol, *crate::common::consts::AX_VENUE))
1469            .raw_symbol(symbol)
1470            .underlying(Ustr::from("EURUSD"))
1471            .asset_class(AssetClass::FX)
1472            .quote_currency(Currency::USD())
1473            .settlement_currency(Currency::USD())
1474            .is_inverse(false)
1475            .price_precision(4)
1476            .size_precision(0)
1477            .price_increment(Price::from("0.0001"))
1478            .size_increment(Quantity::from("1"))
1479            .margin_init(Decimal::new(1, 2))
1480            .margin_maint(Decimal::new(5, 3))
1481            .maker_fee(Decimal::new(2, 4))
1482            .taker_fee(Decimal::new(5, 4))
1483            .ts_event(UnixNanos::default())
1484            .ts_init(UnixNanos::default())
1485            .build()
1486            .unwrap();
1487        InstrumentAny::PerpetualContract(instrument)
1488    }
1489
1490    fn ticker_message(state: AxInstrumentState) -> AxMdTicker {
1491        AxMdTicker {
1492            ts: 1_700_000_000,
1493            tn: 0,
1494            s: Ustr::from("EURUSD-PERP"),
1495            p: rust_decimal::Decimal::ZERO,
1496            q: 0,
1497            o: rust_decimal::Decimal::ZERO,
1498            l: rust_decimal::Decimal::ZERO,
1499            h: rust_decimal::Decimal::ZERO,
1500            v: 0,
1501            oi: None,
1502            m: None,
1503            i: Some(state),
1504            pl: None,
1505            pu: None,
1506            lsp: None,
1507        }
1508    }
1509
1510    fn collect_instrument_statuses(
1511        rx: &mut tokio::sync::mpsc::UnboundedReceiver<DataEvent>,
1512    ) -> Vec<InstrumentStatus> {
1513        let mut statuses = Vec::new();
1514
1515        while let Ok(event) = rx.try_recv() {
1516            if let DataEvent::InstrumentStatus(status) = event {
1517                statuses.push(status);
1518            }
1519        }
1520        statuses
1521    }
1522
1523    #[rstest]
1524    fn test_ticker_instrument_status_emitted_once_when_state_unchanged() {
1525        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
1526        let instruments = Arc::new(AtomicMap::new());
1527        instruments.insert(Ustr::from("EURUSD-PERP"), ticker_test_instrument());
1528
1529        let sdt = Arc::new(AtomicMap::new());
1530        sdt.insert(
1531            "EURUSD-PERP".to_string(),
1532            SymbolDataTypes {
1533                quotes: false,
1534                trades: false,
1535                mark_prices: false,
1536                instrument_status: true,
1537                book_level: None,
1538            },
1539        );
1540
1541        let mut book_sequences = AHashMap::new();
1542        let mut candle_cache = AHashMap::new();
1543        let mut instrument_states = AHashMap::new();
1544        let clock = get_atomic_clock_realtime();
1545
1546        let msg = AxMdMessage::Ticker(ticker_message(AxInstrumentState::Open));
1547        handle_md_message(
1548            msg.clone(),
1549            &tx,
1550            &instruments,
1551            &sdt,
1552            &mut book_sequences,
1553            &mut candle_cache,
1554            &mut instrument_states,
1555            clock,
1556        );
1557
1558        // Same state repeated: second call should not emit a second InstrumentStatus
1559        handle_md_message(
1560            msg,
1561            &tx,
1562            &instruments,
1563            &sdt,
1564            &mut book_sequences,
1565            &mut candle_cache,
1566            &mut instrument_states,
1567            clock,
1568        );
1569
1570        let statuses = collect_instrument_statuses(&mut rx);
1571        assert_eq!(
1572            statuses.len(),
1573            1,
1574            "expected a single emission, found {statuses:?}"
1575        );
1576        assert_eq!(statuses[0].is_trading, Some(true));
1577    }
1578
1579    #[rstest]
1580    fn test_ticker_instrument_status_emitted_on_transition() {
1581        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
1582        let instruments = Arc::new(AtomicMap::new());
1583        instruments.insert(Ustr::from("EURUSD-PERP"), ticker_test_instrument());
1584
1585        let sdt = Arc::new(AtomicMap::new());
1586        sdt.insert(
1587            "EURUSD-PERP".to_string(),
1588            SymbolDataTypes {
1589                quotes: false,
1590                trades: false,
1591                mark_prices: false,
1592                instrument_status: true,
1593                book_level: None,
1594            },
1595        );
1596
1597        let mut book_sequences = AHashMap::new();
1598        let mut candle_cache = AHashMap::new();
1599        let mut instrument_states = AHashMap::new();
1600        let clock = get_atomic_clock_realtime();
1601
1602        handle_md_message(
1603            AxMdMessage::Ticker(ticker_message(AxInstrumentState::Open)),
1604            &tx,
1605            &instruments,
1606            &sdt,
1607            &mut book_sequences,
1608            &mut candle_cache,
1609            &mut instrument_states,
1610            clock,
1611        );
1612        handle_md_message(
1613            AxMdMessage::Ticker(ticker_message(AxInstrumentState::Closed)),
1614            &tx,
1615            &instruments,
1616            &sdt,
1617            &mut book_sequences,
1618            &mut candle_cache,
1619            &mut instrument_states,
1620            clock,
1621        );
1622
1623        let statuses = collect_instrument_statuses(&mut rx);
1624        assert_eq!(statuses.len(), 2, "expected one emission per transition");
1625        assert_eq!(statuses[0].is_trading, Some(true));
1626        assert_eq!(statuses[1].is_trading, Some(false));
1627    }
1628
1629    #[rstest]
1630    fn test_ticker_instrument_status_skipped_when_not_subscribed() {
1631        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
1632        let instruments = Arc::new(AtomicMap::new());
1633        instruments.insert(Ustr::from("EURUSD-PERP"), ticker_test_instrument());
1634
1635        let sdt = Arc::new(AtomicMap::new());
1636        sdt.insert(
1637            "EURUSD-PERP".to_string(),
1638            SymbolDataTypes {
1639                quotes: false,
1640                trades: false,
1641                mark_prices: false,
1642                instrument_status: false,
1643                book_level: None,
1644            },
1645        );
1646
1647        let mut book_sequences = AHashMap::new();
1648        let mut candle_cache = AHashMap::new();
1649        let mut instrument_states = AHashMap::new();
1650        let clock = get_atomic_clock_realtime();
1651
1652        handle_md_message(
1653            AxMdMessage::Ticker(ticker_message(AxInstrumentState::Open)),
1654            &tx,
1655            &instruments,
1656            &sdt,
1657            &mut book_sequences,
1658            &mut candle_cache,
1659            &mut instrument_states,
1660            clock,
1661        );
1662
1663        let statuses = collect_instrument_statuses(&mut rx);
1664        assert!(statuses.is_empty());
1665    }
1666
1667    #[rstest]
1668    fn test_l2_book_emits_quote_when_quotes_subscribed() {
1669        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
1670        let instruments = Arc::new(AtomicMap::new());
1671        instruments.insert(Ustr::from("EURUSD-PERP"), ticker_test_instrument());
1672
1673        let sdt = Arc::new(AtomicMap::new());
1674        sdt.insert(
1675            "EURUSD-PERP".to_string(),
1676            SymbolDataTypes {
1677                quotes: true,
1678                book_level: Some(AxMarketDataLevel::Level2),
1679                ..Default::default()
1680            },
1681        );
1682
1683        let mut book_sequences = AHashMap::new();
1684        let mut candle_cache = AHashMap::new();
1685        let mut instrument_states = AHashMap::new();
1686        let clock = get_atomic_clock_realtime();
1687        let message = AxMdMessage::BookL2(AxMdBookL2 {
1688            ts: 1_700_000_000,
1689            tn: 123,
1690            s: Ustr::from("EURUSD-PERP"),
1691            b: vec![AxBookLevel {
1692                p: dec!(1.1441),
1693                q: 100,
1694            }],
1695            a: vec![AxBookLevel {
1696                p: dec!(1.1448),
1697                q: 200,
1698            }],
1699            st: true,
1700        });
1701
1702        handle_md_message(
1703            message,
1704            &tx,
1705            &instruments,
1706            &sdt,
1707            &mut book_sequences,
1708            &mut candle_cache,
1709            &mut instrument_states,
1710            clock,
1711        );
1712
1713        let events = std::iter::from_fn(|| rx.try_recv().ok()).collect::<Vec<_>>();
1714        let quote = events.iter().find_map(|event| match event {
1715            DataEvent::Data(Data::Quote(quote)) => Some(quote),
1716            _ => None,
1717        });
1718
1719        assert_eq!(
1720            quote.map(|quote| quote.bid_price),
1721            Some(Price::from("1.1441"))
1722        );
1723        assert_eq!(
1724            quote.map(|quote| quote.ask_price),
1725            Some(Price::from("1.1448"))
1726        );
1727    }
1728}