1use 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#[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 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 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 #[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 #[must_use]
183 pub fn instruments(&self) -> &Arc<AtomicMap<Ustr, InstrumentAny>> {
184 &self.instruments
185 }
186
187 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 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 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 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 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 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 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 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 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
1464fn 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 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}