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