Skip to main content

nautilus_tardis/
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//! Tardis data client for streaming replay or live data into the engine.
17
18use 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/// Tardis data client for streaming replay or live data into the platform.
72#[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    /// Creates a new [`TardisDataClient`] instance.
84    ///
85    /// # Errors
86    ///
87    /// Returns an error if the data event sender is not initialized.
88    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    /// Returns `true` if the client is configured for live streaming mode.
104    fn is_stream_mode(&self) -> bool {
105        self.config.options.is_empty() && !self.config.stream_options.is_empty()
106    }
107
108    /// Builds the WebSocket URL for connecting to the Tardis Machine Server.
109    ///
110    /// Ensures `derivative_ticker` is included in the data types for each
111    /// option set so that mark price, index price, and funding rate events
112    /// are available without requiring manual configuration.
113    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    /// Spawns the WebSocket message processing loop using an already-connected
144    /// stream. The initial handshake happens in `connect()` so callers get an
145    /// error if the first connection fails. In stream mode the spawned task
146    /// handles subsequent reconnections automatically.
147    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            // Process the initial (already-connected) stream
167            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            // Mark disconnected while reconnecting so health checks see the outage
183            connected.store(false, Ordering::Release);
184
185            // Reconnection loop (stream mode only)
186            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                // Reconnect WS first (critical path), then refresh instruments
203                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    /// Runs a single WebSocket session: starts heartbeat, processes messages,
257    /// and returns whether the caller should attempt reconnection.
258    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    /// Extracts and sends all data events from a `DerivativeTicker` message:
309    /// funding rate, mark price, and index price. Only emits events when values
310    /// change from the previous update. Returns `false` if the channel is broken
311    /// and the caller should exit the loop.
312    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    /// Processes WebSocket messages until the stream ends, an error occurs, or
351    /// the cancellation token fires. Returns `true` if the caller should attempt
352    /// reconnection (stream mode only).
353    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 // Tardis is multi-venue
449    }
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}