1use std::{
19 sync::{
20 Arc,
21 atomic::{AtomicBool, Ordering},
22 },
23 time::Duration,
24};
25
26use ahash::{AHashMap, AHashSet};
27use futures_util::{SinkExt, StreamExt};
28use nautilus_common::{
29 clients::DataClient,
30 live::runner::get_data_event_sender,
31 messages::{
32 DataEvent,
33 data::{
34 subscribe::{SubscribeFundingRates, SubscribeIndexPrices, SubscribeMarkPrices},
35 unsubscribe::{UnsubscribeFundingRates, UnsubscribeIndexPrices, UnsubscribeMarkPrices},
36 },
37 },
38};
39use nautilus_core::string::urlencoding;
40use nautilus_live::task::TaskGroup;
41use nautilus_model::{
42 data::Data,
43 identifiers::{ClientId, Venue},
44};
45use tokio_tungstenite::{connect_async, tungstenite};
46use tokio_util::sync::CancellationToken;
47
48use crate::{
49 common::{
50 consts::{
51 WS_HEARTBEAT_INTERVAL_SECS, WS_INITIAL_RECONNECT_DELAY_SECS,
52 WS_MAX_RECONNECT_DELAY_SECS,
53 },
54 enums::TardisDataType,
55 urls::resolve_ws_base_url,
56 },
57 config::{BookSnapshotOutput, TardisDataClientConfig},
58 http::TardisHttpClient,
59 machine::{
60 cache::DerivativeTickerCache,
61 client::determine_instrument_info,
62 message::WsMessage,
63 parse::{
64 parse_derivative_ticker_index_price, parse_derivative_ticker_mark_price,
65 parse_tardis_ws_message_data, parse_tardis_ws_message_funding_rate,
66 },
67 types::{TardisInstrumentKey, TardisInstrumentMiniInfo},
68 },
69};
70
71#[derive(Debug)]
73pub struct TardisDataClient {
74 client_id: ClientId,
75 config: TardisDataClientConfig,
76 is_connected: Arc<AtomicBool>,
77 cancellation_token: CancellationToken,
78 tasks: TaskGroup,
79 data_sender: tokio::sync::mpsc::UnboundedSender<DataEvent>,
80}
81
82impl TardisDataClient {
83 pub fn new(client_id: ClientId, config: TardisDataClientConfig) -> anyhow::Result<Self> {
89 let data_sender = get_data_event_sender();
90
91 let tasks = TaskGroup::new();
92
93 Ok(Self {
94 client_id,
95 config,
96 is_connected: Arc::new(AtomicBool::new(false)),
97 cancellation_token: tasks.cancellation_token(),
98 tasks,
99 data_sender,
100 })
101 }
102
103 fn is_stream_mode(&self) -> bool {
105 self.config.options.is_empty() && !self.config.stream_options.is_empty()
106 }
107
108 fn build_ws_url(&self, base_url: &str) -> anyhow::Result<String> {
114 let deriv = TardisDataType::DerivativeTicker.as_tardis_str();
115
116 if self.is_stream_mode() {
117 let mut options = self.config.stream_options.clone();
118 for opt in &mut options {
119 if !opt.data_types.iter().any(|dt| dt == deriv) {
120 opt.data_types.push(deriv.to_string());
121 }
122 }
123 let options_json = serde_json::to_string(&options)?;
124 Ok(format!(
125 "{base_url}/ws-stream-normalized?options={}",
126 urlencoding::encode(&options_json)
127 ))
128 } else {
129 let mut options = self.config.options.clone();
130 for opt in &mut options {
131 if !opt.data_types.iter().any(|dt| dt == deriv) {
132 opt.data_types.push(deriv.to_string());
133 }
134 }
135 let options_json = serde_json::to_string(&options)?;
136 Ok(format!(
137 "{base_url}/ws-replay-normalized?options={}",
138 urlencoding::encode(&options_json)
139 ))
140 }
141 }
142
143 fn spawn_ws_task(
148 &self,
149 ws_stream: tokio_tungstenite::WebSocketStream<
150 tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
151 >,
152 url: String,
153 instrument_map: AHashMap<TardisInstrumentKey, Arc<TardisInstrumentMiniInfo>>,
154 book_snapshot_output: BookSnapshotOutput,
155 extract_bbo_as_quotes: bool,
156 is_stream_mode: bool,
157 ) -> anyhow::Result<()> {
158 let sender = self.data_sender.clone();
159 let cancel = self.cancellation_token.clone();
160 let connected = self.is_connected.clone();
161
162 let future = async move {
163 let mut reconnect_delay = Duration::from_secs(WS_INITIAL_RECONNECT_DELAY_SECS);
164 let instrument_map = instrument_map;
165
166 let should_reconnect = Self::run_ws_session(
168 ws_stream,
169 &cancel,
170 &sender,
171 &instrument_map,
172 &book_snapshot_output,
173 extract_bbo_as_quotes,
174 )
175 .await;
176
177 if !should_reconnect || !is_stream_mode || cancel.is_cancelled() {
178 connected.store(false, Ordering::Release);
179 return;
180 }
181
182 connected.store(false, Ordering::Release);
184
185 loop {
187 log::warn!(
188 "Stream disconnected, reconnecting in {}s",
189 reconnect_delay.as_secs()
190 );
191
192 tokio::select! {
193 () = tokio::time::sleep(reconnect_delay) => {}
194 () = cancel.cancelled() => break,
195 }
196
197 reconnect_delay = std::cmp::min(
198 reconnect_delay * 2,
199 Duration::from_secs(WS_MAX_RECONNECT_DELAY_SECS),
200 );
201
202 let ws_result = tokio::select! {
204 result = connect_async(&url) => Some(result),
205 () = cancel.cancelled() => None,
206 };
207
208 let Some(ws_result) = ws_result else {
209 break;
210 };
211
212 match ws_result {
213 Ok((ws_stream, _)) => {
214 log::info!("Reconnected to Tardis Machine");
215 connected.store(true, Ordering::Release);
216 reconnect_delay = Duration::from_secs(WS_INITIAL_RECONNECT_DELAY_SECS);
217
218 let should_reconnect = Self::run_ws_session(
219 ws_stream,
220 &cancel,
221 &sender,
222 &instrument_map,
223 &book_snapshot_output,
224 extract_bbo_as_quotes,
225 )
226 .await;
227
228 if !should_reconnect || cancel.is_cancelled() {
229 break;
230 }
231
232 connected.store(false, Ordering::Release);
233 }
234 Err(e) => {
235 if cancel.is_cancelled() {
236 break;
237 }
238
239 log::warn!(
240 "Failed to reconnect to Tardis Machine: {e}, retrying in {}s",
241 reconnect_delay.as_secs()
242 );
243 }
244 }
245 }
246
247 connected.store(false, Ordering::Release);
248 };
249
250 self.tasks
251 .spawn(future)
252 .map_err(|e| anyhow::anyhow!("failed to register Tardis stream task: {e}"))?;
253 Ok(())
254 }
255
256 async fn run_ws_session(
259 ws_stream: tokio_tungstenite::WebSocketStream<
260 tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
261 >,
262 cancel: &CancellationToken,
263 sender: &tokio::sync::mpsc::UnboundedSender<DataEvent>,
264 instrument_map: &AHashMap<TardisInstrumentKey, Arc<TardisInstrumentMiniInfo>>,
265 book_snapshot_output: &BookSnapshotOutput,
266 extract_bbo_as_quotes: bool,
267 ) -> bool {
268 let (mut writer, mut reader) = ws_stream.split();
269
270 let heartbeat_token = cancel.child_token();
271 let heartbeat_signal = heartbeat_token.clone();
272
273 let heartbeat = async move {
274 let mut interval =
275 tokio::time::interval(Duration::from_secs(WS_HEARTBEAT_INTERVAL_SECS));
276 loop {
277 tokio::select! {
278 _ = interval.tick() => {
279 log::trace!("Sending PING");
280
281 if let Err(e) = writer.send(tungstenite::Message::Ping(vec![].into())).await {
282 log::debug!("Heartbeat send failed: {e}");
283 break;
284 }
285 }
286 () = heartbeat_signal.cancelled() => break,
287 }
288 }
289 };
290
291 let message_loop = async {
292 let should_reconnect = Self::run_ws_loop(
293 &mut reader,
294 cancel,
295 sender,
296 instrument_map,
297 book_snapshot_output,
298 extract_bbo_as_quotes,
299 )
300 .await;
301 heartbeat_token.cancel();
302 should_reconnect
303 };
304 let (should_reconnect, ()) = tokio::join!(message_loop, heartbeat);
305 should_reconnect
306 }
307
308 fn send_derivative_ticker_events(
313 ws_msg: &WsMessage,
314 info: &Arc<TardisInstrumentMiniInfo>,
315 sender: &tokio::sync::mpsc::UnboundedSender<DataEvent>,
316 cache: &mut DerivativeTickerCache,
317 ) -> bool {
318 if let Some(funding) = parse_tardis_ws_message_funding_rate(ws_msg.clone(), info)
319 && cache.should_emit_funding_rate(&funding)
320 && sender.send(DataEvent::FundingRate(funding)).is_err()
321 {
322 return false;
323 }
324
325 if let WsMessage::DerivativeTicker(msg) = ws_msg {
326 if let Ok(Some(mark_price)) =
327 parse_derivative_ticker_mark_price(msg, info.instrument_id, info.price_precision)
328 && cache.should_emit_mark_price(&mark_price)
329 && sender
330 .send(DataEvent::Data(Data::MarkPrice(mark_price)))
331 .is_err()
332 {
333 return false;
334 }
335
336 if let Ok(Some(index_price)) =
337 parse_derivative_ticker_index_price(msg, info.instrument_id, info.price_precision)
338 && cache.should_emit_index_price(&index_price)
339 && sender
340 .send(DataEvent::Data(Data::IndexPrice(index_price)))
341 .is_err()
342 {
343 return false;
344 }
345 }
346
347 true
348 }
349
350 async fn run_ws_loop(
354 reader: &mut futures_util::stream::SplitStream<
355 tokio_tungstenite::WebSocketStream<
356 tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
357 >,
358 >,
359 cancel: &CancellationToken,
360 sender: &tokio::sync::mpsc::UnboundedSender<DataEvent>,
361 instrument_map: &AHashMap<TardisInstrumentKey, Arc<TardisInstrumentMiniInfo>>,
362 book_snapshot_output: &BookSnapshotOutput,
363 extract_bbo_as_quotes: bool,
364 ) -> bool {
365 let mut ticker_cache = DerivativeTickerCache::default();
366
367 loop {
368 let msg = tokio::select! {
369 msg = reader.next() => msg,
370 () = cancel.cancelled() => {
371 log::debug!("Stream task cancelled");
372 return false;
373 }
374 };
375
376 match msg {
377 Some(Ok(tungstenite::Message::Text(text))) => {
378 match serde_json::from_str::<WsMessage>(&text) {
379 Ok(ws_msg) => {
380 if matches!(ws_msg, WsMessage::Disconnect(_)) {
381 log::debug!("Received disconnect message");
382 continue;
383 }
384
385 let info = determine_instrument_info(&ws_msg, instrument_map);
386
387 if let Some(info) = info {
388 if matches!(ws_msg, WsMessage::DerivativeTicker(_)) {
389 if !Self::send_derivative_ticker_events(
390 &ws_msg,
391 &info,
392 sender,
393 &mut ticker_cache,
394 ) {
395 return false;
396 }
397 } else {
398 let data = parse_tardis_ws_message_data(
399 ws_msg,
400 &info,
401 book_snapshot_output,
402 extract_bbo_as_quotes,
403 );
404
405 for data in data {
406 if let Err(e) = sender.send(DataEvent::Data(data)) {
407 log::error!("Failed to send data event: {e}");
408 return false;
409 }
410 }
411 }
412 }
413 }
414 Err(e) => {
415 log::error!("Failed to deserialize message: {e}");
416 }
417 }
418 }
419 Some(Ok(tungstenite::Message::Close(frame))) => {
420 if let Some(frame) = frame {
421 log::debug!("WebSocket closed: {} {}", frame.code, frame.reason);
422 } else {
423 log::debug!("WebSocket closed");
424 }
425 return true;
426 }
427 Some(Ok(_)) => {}
428 Some(Err(e)) => {
429 log::warn!("WebSocket error: {e}");
430 return true;
431 }
432 None => {
433 log::debug!("Stream ended");
434 return true;
435 }
436 }
437 }
438 }
439}
440
441#[async_trait::async_trait(?Send)]
442impl DataClient for TardisDataClient {
443 fn client_id(&self) -> ClientId {
444 self.client_id
445 }
446
447 fn venue(&self) -> Option<Venue> {
448 None }
450
451 fn start(&mut self) -> anyhow::Result<()> {
452 log::info!("Starting {}", self.client_id);
453 Ok(())
454 }
455
456 fn stop(&mut self) -> anyhow::Result<()> {
457 log::info!("Stopping {}", self.client_id);
458 self.tasks.begin_shutdown();
459 self.is_connected.store(false, Ordering::Release);
460 Ok(())
461 }
462
463 fn reset(&mut self) -> anyhow::Result<()> {
464 self.tasks.begin_shutdown();
465 self.is_connected.store(false, Ordering::Release);
466 Ok(())
467 }
468
469 fn dispose(&mut self) -> anyhow::Result<()> {
470 self.stop()
471 }
472
473 fn is_connected(&self) -> bool {
474 self.is_connected.load(Ordering::Acquire)
475 }
476
477 fn is_disconnected(&self) -> bool {
478 !self.is_connected()
479 }
480
481 fn subscribe_mark_prices(&mut self, _cmd: SubscribeMarkPrices) -> anyhow::Result<()> {
482 Ok(())
483 }
484
485 fn subscribe_index_prices(&mut self, _cmd: SubscribeIndexPrices) -> anyhow::Result<()> {
486 Ok(())
487 }
488
489 fn subscribe_funding_rates(&mut self, _cmd: SubscribeFundingRates) -> anyhow::Result<()> {
490 Ok(())
491 }
492
493 fn unsubscribe_mark_prices(&mut self, _cmd: &UnsubscribeMarkPrices) -> anyhow::Result<()> {
494 Ok(())
495 }
496
497 fn unsubscribe_index_prices(&mut self, _cmd: &UnsubscribeIndexPrices) -> anyhow::Result<()> {
498 Ok(())
499 }
500
501 fn unsubscribe_funding_rates(&mut self, _cmd: &UnsubscribeFundingRates) -> anyhow::Result<()> {
502 Ok(())
503 }
504
505 async fn connect(&mut self) -> anyhow::Result<()> {
506 if self.is_connected() && self.tasks.is_open() {
507 return Ok(());
508 }
509
510 if self.config.options.is_empty() && self.config.stream_options.is_empty() {
511 anyhow::bail!("Either replay `options` or `stream_options` must be provided");
512 }
513
514 if !self.tasks.is_open() {
515 self.tasks
516 .finish_shutdown(Duration::from_secs(1), Duration::from_secs(2))
517 .await
518 .map_err(|e| anyhow::anyhow!("Failed to terminate Tardis tasks: {e}"))?;
519 self.tasks
520 .start_generation()
521 .map_err(|e| anyhow::anyhow!("Failed to start Tardis task generation: {e}"))?;
522 self.cancellation_token = self.tasks.cancellation_token();
523 }
524
525 let is_stream_mode = self.is_stream_mode();
526 let book_snapshot_output = self.config.book_snapshot_output.clone();
527 let extract_bbo_as_quotes = self.config.extract_bbo_as_quotes;
528
529 let http_client = TardisHttpClient::new(
530 self.config.api_key.as_deref(),
531 None,
532 None,
533 self.config.normalize_symbols,
534 self.config.proxy_url.clone(),
535 )?;
536
537 let exchanges: AHashSet<_> = if is_stream_mode {
538 self.config
539 .stream_options
540 .iter()
541 .map(|opt| opt.exchange)
542 .collect()
543 } else {
544 self.config.options.iter().map(|opt| opt.exchange).collect()
545 };
546
547 let base_url = resolve_ws_base_url(self.config.tardis_ws_url.as_deref())?;
548 let (instrument_map, instruments) = http_client
549 .bootstrap_instruments(&exchanges)
550 .await
551 .map_err(|e| anyhow::anyhow!("Failed to bootstrap instruments: {e}"))?;
552
553 for instrument in instruments {
554 if let Err(e) = self.data_sender.send(DataEvent::Instrument(instrument)) {
555 log::error!("Failed to send instrument event: {e}");
556 }
557 }
558
559 let url = self.build_ws_url(&base_url)?;
560
561 let mode_label = if is_stream_mode { "stream" } else { "replay" };
562 log::info!("Connecting to Tardis Machine {mode_label}");
563 log::debug!("URL: {url}");
564
565 let (ws_stream, _) = connect_async(&url)
566 .await
567 .map_err(|e| anyhow::anyhow!("Failed to connect to Tardis Machine: {e}"))?;
568
569 log::info!("Connected to Tardis Machine");
570
571 if let Err(e) = self.spawn_ws_task(
572 ws_stream,
573 url,
574 instrument_map,
575 book_snapshot_output,
576 extract_bbo_as_quotes,
577 is_stream_mode,
578 ) {
579 self.tasks.begin_shutdown();
580 if let Err(teardown_error) = self
581 .tasks
582 .finish_shutdown(Duration::from_secs(1), Duration::from_secs(2))
583 .await
584 {
585 return Err(e.context(format!("Tardis startup teardown failed: {teardown_error}")));
586 }
587 return Err(e);
588 }
589 self.is_connected.store(true, Ordering::Release);
590
591 log::info!("Connected: {}", self.client_id);
592 Ok(())
593 }
594
595 async fn disconnect(&mut self) -> anyhow::Result<()> {
596 self.tasks.begin_shutdown();
597 let had_tasks = !self.tasks.is_empty();
598 let tasks_result = self
599 .tasks
600 .finish_shutdown(Duration::from_secs(1), Duration::from_secs(2))
601 .await
602 .map_err(|e| anyhow::anyhow!("Failed to terminate Tardis tasks: {e}"));
603
604 if had_tasks {
605 log::info!("Disconnected: {}", self.client_id);
606 }
607
608 self.is_connected.store(false, Ordering::Release);
609
610 tasks_result
611 }
612}
613
614#[cfg(test)]
615mod tests {
616 use jiff::civil::Date;
617 use nautilus_common::live::runner::set_data_event_sender;
618 use rstest::rstest;
619
620 use super::*;
621 use crate::{
622 common::{consts::TARDIS_CLIENT_ID, enums::TardisExchange},
623 config::TardisDataClientConfig,
624 machine::types::ReplayNormalizedRequestOptions,
625 };
626
627 fn setup_test_env() {
628 use std::cell::OnceCell;
629
630 thread_local! {
631 static INIT: OnceCell<()> = const { OnceCell::new() };
632 }
633
634 INIT.with(|cell| {
635 cell.get_or_init(|| {
636 let (sender, _receiver) = tokio::sync::mpsc::unbounded_channel::<DataEvent>();
637 set_data_event_sender(sender);
638 });
639 });
640 }
641
642 #[rstest]
643 fn test_build_ws_url_injects_derivative_ticker() {
644 setup_test_env();
645
646 let config = TardisDataClientConfig {
647 options: vec![ReplayNormalizedRequestOptions {
648 exchange: TardisExchange::BinanceFutures,
649 symbols: Some(vec!["BTCUSDT".to_string()]),
650 from: Date::new(2024, 1, 1).unwrap(),
651 to: Date::new(2024, 1, 2).unwrap(),
652 data_types: vec!["trade".to_string()],
653 with_disconnect_messages: Some(false),
654 }],
655 ..Default::default()
656 };
657
658 let client = TardisDataClient::new(*TARDIS_CLIENT_ID, config).unwrap();
659 let url = client.build_ws_url("ws://localhost:8001").unwrap();
660
661 assert!(
662 url.contains("derivative_ticker"),
663 "URL should contain derivative_ticker but was: {url}"
664 );
665 assert!(url.contains("trade"), "URL should still contain trade");
666 }
667
668 #[rstest]
669 fn test_build_ws_url_does_not_duplicate_derivative_ticker() {
670 setup_test_env();
671
672 let config = TardisDataClientConfig {
673 options: vec![ReplayNormalizedRequestOptions {
674 exchange: TardisExchange::BinanceFutures,
675 symbols: Some(vec!["BTCUSDT".to_string()]),
676 from: Date::new(2024, 1, 1).unwrap(),
677 to: Date::new(2024, 1, 2).unwrap(),
678 data_types: vec!["trade".to_string(), "derivative_ticker".to_string()],
679 with_disconnect_messages: Some(false),
680 }],
681 ..Default::default()
682 };
683
684 let client = TardisDataClient::new(*TARDIS_CLIENT_ID, config).unwrap();
685 let ws_url = client.build_ws_url("ws://localhost:8001").unwrap();
686
687 let decoded = urlencoding::decode(ws_url.split("options=").nth(1).unwrap()).unwrap();
688 let count = decoded.matches("derivative_ticker").count();
689 assert_eq!(count, 1, "derivative_ticker should appear exactly once");
690 }
691
692 #[rstest]
693 fn test_stop_marks_client_disconnected_synchronously() {
694 setup_test_env();
695
696 let mut client =
697 TardisDataClient::new(*TARDIS_CLIENT_ID, TardisDataClientConfig::default()).unwrap();
698 client.is_connected.store(true, Ordering::Release);
699
700 client.stop().unwrap();
701
702 assert!(client.is_disconnected());
703 }
704
705 #[rstest]
706 fn test_reset_marks_client_disconnected_synchronously() {
707 setup_test_env();
708
709 let mut client =
710 TardisDataClient::new(*TARDIS_CLIENT_ID, TardisDataClientConfig::default()).unwrap();
711 client.is_connected.store(true, Ordering::Release);
712
713 client.reset().unwrap();
714
715 assert!(client.is_disconnected());
716 }
717}