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, sender::EventSender},
31    messages::{
32        DataEvent,
33        data::{
34            subscribe::{SubscribeFundingRates, SubscribeIndexPrices, SubscribeMarkPrices},
35            unsubscribe::{UnsubscribeFundingRates, UnsubscribeIndexPrices, UnsubscribeMarkPrices},
36        },
37    },
38};
39use nautilus_core::{consts::NAUTILUS_USER_AGENT, string::urlencoding};
40use nautilus_live::task::TaskGroup;
41use nautilus_model::{
42    data::Data,
43    identifiers::{ClientId, Venue},
44};
45use nautilus_network::url_origin;
46use tokio_tungstenite::{
47    connect_async,
48    tungstenite::{
49        self,
50        client::IntoClientRequest,
51        http::{HeaderValue, header::USER_AGENT},
52    },
53};
54use tokio_util::sync::CancellationToken;
55
56use crate::{
57    common::{
58        consts::{
59            WS_HEARTBEAT_INTERVAL_SECS, WS_INITIAL_RECONNECT_DELAY_SECS,
60            WS_MAX_RECONNECT_DELAY_SECS,
61        },
62        enums::TardisDataType,
63        urls::resolve_ws_base_url,
64    },
65    config::{BookSnapshotOutput, TardisDataClientConfig},
66    http::TardisHttpClient,
67    machine::{
68        cache::DerivativeTickerCache,
69        client::determine_instrument_info,
70        is_unsupported_streaming_error,
71        message::{WsMessage, decode_ws_message},
72        parse::{
73            parse_derivative_ticker_index_price, parse_derivative_ticker_mark_price,
74            parse_tardis_ws_message_data, parse_tardis_ws_message_funding_rate,
75        },
76        types::{TardisInstrumentKey, TardisInstrumentMiniInfo},
77        validate_stream_options,
78    },
79};
80
81/// Tardis data client for streaming replay or live data into the platform.
82#[derive(Debug)]
83pub struct TardisDataClient {
84    client_id: ClientId,
85    config: TardisDataClientConfig,
86    is_connected: Arc<AtomicBool>,
87    cancellation_token: CancellationToken,
88    tasks: TaskGroup,
89    data_sender: EventSender<DataEvent>,
90}
91
92impl TardisDataClient {
93    /// Creates a new [`TardisDataClient`] instance.
94    ///
95    /// # Errors
96    ///
97    /// Returns an error if the data event sender is not initialized.
98    pub fn new(client_id: ClientId, config: TardisDataClientConfig) -> anyhow::Result<Self> {
99        let data_sender = get_data_event_sender();
100
101        let tasks = TaskGroup::new();
102
103        Ok(Self {
104            client_id,
105            config,
106            is_connected: Arc::new(AtomicBool::new(false)),
107            cancellation_token: tasks.cancellation_token(),
108            tasks,
109            data_sender,
110        })
111    }
112
113    /// Returns `true` if the client is configured for live streaming mode.
114    fn is_stream_mode(&self) -> bool {
115        self.config.options.is_empty() && !self.config.stream_options.is_empty()
116    }
117
118    /// Builds the WebSocket URL for connecting to the Tardis Machine Server.
119    ///
120    /// Ensures `derivative_ticker` is included in the data types for each
121    /// option set so that mark price, index price, and funding rate events
122    /// are available without requiring manual configuration.
123    fn build_ws_url(&self, base_url: &str) -> anyhow::Result<String> {
124        let deriv = TardisDataType::DerivativeTicker.as_tardis_str();
125
126        if self.is_stream_mode() {
127            let mut options = self.config.stream_options.clone();
128            for opt in &mut options {
129                if !opt.data_types.iter().any(|dt| dt == deriv) {
130                    opt.data_types.push(deriv.to_string());
131                }
132            }
133            let options_json = serde_json::to_string(&options)?;
134            Ok(format!(
135                "{base_url}/ws-stream-normalized?options={}",
136                urlencoding::encode(&options_json)
137            ))
138        } else {
139            let mut options = self.config.options.clone();
140            for opt in &mut options {
141                if !opt.data_types.iter().any(|dt| dt == deriv) {
142                    opt.data_types.push(deriv.to_string());
143                }
144            }
145            let options_json = serde_json::to_string(&options)?;
146            Ok(format!(
147                "{base_url}/ws-replay-normalized?options={}",
148                urlencoding::encode(&options_json)
149            ))
150        }
151    }
152
153    /// Spawns the WebSocket message processing loop using an already-connected
154    /// stream. The initial handshake happens in `connect()` so callers get an
155    /// error if the first connection fails. In stream mode the spawned task
156    /// handles subsequent reconnections automatically.
157    fn spawn_ws_task(
158        &self,
159        ws_stream: tokio_tungstenite::WebSocketStream<
160            tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
161        >,
162        url: String,
163        instrument_map: AHashMap<TardisInstrumentKey, Arc<TardisInstrumentMiniInfo>>,
164        book_snapshot_output: BookSnapshotOutput,
165        extract_bbo_as_quotes: bool,
166        is_stream_mode: bool,
167    ) -> anyhow::Result<()> {
168        let sender = self.data_sender.clone();
169        let cancel = self.cancellation_token.clone();
170        let connected = self.is_connected.clone();
171
172        let future = async move {
173            let mut reconnect_delay = Duration::from_secs(WS_INITIAL_RECONNECT_DELAY_SECS);
174            let instrument_map = instrument_map;
175
176            // Process the initial (already-connected) stream
177            let should_reconnect = Self::run_ws_session(
178                ws_stream,
179                &cancel,
180                &sender,
181                &instrument_map,
182                &book_snapshot_output,
183                extract_bbo_as_quotes,
184            )
185            .await;
186
187            if !should_reconnect || !is_stream_mode || cancel.is_cancelled() {
188                connected.store(false, Ordering::Release);
189                return;
190            }
191
192            // Mark disconnected while reconnecting so health checks see the outage
193            connected.store(false, Ordering::Release);
194
195            // Reconnection loop (stream mode only)
196            loop {
197                log::warn!(
198                    "Stream disconnected, reconnecting in {}s",
199                    reconnect_delay.as_secs()
200                );
201
202                tokio::select! {
203                    () = tokio::time::sleep(reconnect_delay) => {}
204                    () = cancel.cancelled() => break,
205                }
206
207                reconnect_delay = std::cmp::min(
208                    reconnect_delay * 2,
209                    Duration::from_secs(WS_MAX_RECONNECT_DELAY_SECS),
210                );
211
212                // Reconnect WS first (critical path), then refresh instruments
213                let ws_request = match build_ws_request(&url) {
214                    Ok(request) => request,
215                    Err(_) => {
216                        log::error!(
217                            "Failed to build Tardis Machine request: {}",
218                            url_origin(&url)
219                        );
220                        break;
221                    }
222                };
223
224                let ws_result = tokio::select! {
225                    result = connect_async(ws_request) => Some(result),
226                    () = cancel.cancelled() => None,
227                };
228
229                let Some(ws_result) = ws_result else {
230                    break;
231                };
232
233                match ws_result {
234                    Ok((ws_stream, _)) => {
235                        log::info!("Reconnected to Tardis Machine");
236                        connected.store(true, Ordering::Release);
237                        reconnect_delay = Duration::from_secs(WS_INITIAL_RECONNECT_DELAY_SECS);
238
239                        let should_reconnect = Self::run_ws_session(
240                            ws_stream,
241                            &cancel,
242                            &sender,
243                            &instrument_map,
244                            &book_snapshot_output,
245                            extract_bbo_as_quotes,
246                        )
247                        .await;
248
249                        if !should_reconnect || cancel.is_cancelled() {
250                            break;
251                        }
252
253                        connected.store(false, Ordering::Release);
254                    }
255                    Err(_) => {
256                        if cancel.is_cancelled() {
257                            break;
258                        }
259
260                        log::warn!(
261                            "Failed to reconnect to Tardis Machine: {}, retrying in {}s",
262                            url_origin(&url),
263                            reconnect_delay.as_secs()
264                        );
265                    }
266                }
267            }
268
269            connected.store(false, Ordering::Release);
270        };
271
272        self.tasks
273            .spawn(future)
274            .map_err(|e| anyhow::anyhow!("failed to register Tardis stream task: {e}"))?;
275        Ok(())
276    }
277
278    /// Runs a single WebSocket session: starts heartbeat, processes messages,
279    /// and returns whether the caller should attempt reconnection.
280    async fn run_ws_session(
281        ws_stream: tokio_tungstenite::WebSocketStream<
282            tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
283        >,
284        cancel: &CancellationToken,
285        sender: &EventSender<DataEvent>,
286        instrument_map: &AHashMap<TardisInstrumentKey, Arc<TardisInstrumentMiniInfo>>,
287        book_snapshot_output: &BookSnapshotOutput,
288        extract_bbo_as_quotes: bool,
289    ) -> bool {
290        let (mut writer, mut reader) = ws_stream.split();
291
292        let heartbeat_token = cancel.child_token();
293        let heartbeat_signal = heartbeat_token.clone();
294
295        let heartbeat = async move {
296            let mut interval =
297                tokio::time::interval(Duration::from_secs(WS_HEARTBEAT_INTERVAL_SECS));
298            loop {
299                tokio::select! {
300                    _ = interval.tick() => {
301                        log::trace!("Sending PING");
302
303                        if writer.send(tungstenite::Message::Ping(vec![].into())).await.is_err() {
304                            log::debug!("Heartbeat send failed");
305                            break;
306                        }
307                    }
308                    () = heartbeat_signal.cancelled() => break,
309                }
310            }
311        };
312
313        let message_loop = async {
314            let should_reconnect = Self::run_ws_loop(
315                &mut reader,
316                cancel,
317                sender,
318                instrument_map,
319                book_snapshot_output,
320                extract_bbo_as_quotes,
321            )
322            .await;
323            heartbeat_token.cancel();
324            should_reconnect
325        };
326        let (should_reconnect, ()) = tokio::join!(message_loop, heartbeat);
327        should_reconnect
328    }
329
330    /// Extracts and sends all data events from a `DerivativeTicker` message:
331    /// funding rate, mark price, and index price. Only emits events when values
332    /// change from the previous update. Returns `false` if the channel is broken
333    /// and the caller should exit the loop.
334    fn send_derivative_ticker_events(
335        ws_msg: &WsMessage,
336        info: &Arc<TardisInstrumentMiniInfo>,
337        sender: &EventSender<DataEvent>,
338        cache: &mut DerivativeTickerCache,
339    ) -> bool {
340        if let Some(funding) = parse_tardis_ws_message_funding_rate(ws_msg.clone(), info)
341            && cache.should_emit_funding_rate(&funding)
342            && sender.send(DataEvent::FundingRate(funding)).is_err()
343        {
344            return false;
345        }
346
347        if let WsMessage::DerivativeTicker(msg) = ws_msg {
348            if let Ok(Some(mark_price)) =
349                parse_derivative_ticker_mark_price(msg, info.instrument_id, info.price_precision)
350                && cache.should_emit_mark_price(&mark_price)
351                && sender
352                    .send(DataEvent::Data(Data::MarkPrice(mark_price)))
353                    .is_err()
354            {
355                return false;
356            }
357
358            if let Ok(Some(index_price)) =
359                parse_derivative_ticker_index_price(msg, info.instrument_id, info.price_precision)
360                && cache.should_emit_index_price(&index_price)
361                && sender
362                    .send(DataEvent::Data(Data::IndexPrice(index_price)))
363                    .is_err()
364            {
365                return false;
366            }
367        }
368
369        true
370    }
371
372    /// Processes WebSocket messages until the stream ends, an error occurs, or
373    /// the cancellation token fires. Returns `true` if the caller should attempt
374    /// reconnection (stream mode only).
375    async fn run_ws_loop(
376        reader: &mut futures_util::stream::SplitStream<
377            tokio_tungstenite::WebSocketStream<
378                tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
379            >,
380        >,
381        cancel: &CancellationToken,
382        sender: &EventSender<DataEvent>,
383        instrument_map: &AHashMap<TardisInstrumentKey, Arc<TardisInstrumentMiniInfo>>,
384        book_snapshot_output: &BookSnapshotOutput,
385        extract_bbo_as_quotes: bool,
386    ) -> bool {
387        let mut ticker_cache = DerivativeTickerCache::default();
388
389        loop {
390            let msg = tokio::select! {
391                msg = reader.next() => msg,
392                () = cancel.cancelled() => {
393                    log::debug!("Stream task cancelled");
394                    return false;
395                }
396            };
397
398            match msg {
399                Some(Ok(tungstenite::Message::Text(text))) => match decode_ws_message(&text) {
400                    Ok(ws_msg) => {
401                        if matches!(ws_msg, WsMessage::Disconnect(_)) {
402                            log::debug!("Received disconnect message");
403                            continue;
404                        }
405
406                        let info = determine_instrument_info(&ws_msg, instrument_map);
407
408                        if let Some(info) = info {
409                            if matches!(ws_msg, WsMessage::DerivativeTicker(_)) {
410                                if !Self::send_derivative_ticker_events(
411                                    &ws_msg,
412                                    &info,
413                                    sender,
414                                    &mut ticker_cache,
415                                ) {
416                                    return false;
417                                }
418                            } else {
419                                let data = parse_tardis_ws_message_data(
420                                    ws_msg,
421                                    &info,
422                                    book_snapshot_output,
423                                    extract_bbo_as_quotes,
424                                );
425
426                                for data in data {
427                                    if let Err(e) = sender.send(DataEvent::Data(data)) {
428                                        log::error!("Failed to send data event: {e}");
429                                        return false;
430                                    }
431                                }
432                            }
433                        }
434                    }
435                    Err(e) => {
436                        log::error!(
437                            "Failed to deserialize message at line {}, column {}",
438                            e.line(),
439                            e.column()
440                        );
441                    }
442                },
443                Some(Ok(tungstenite::Message::Close(Some(frame))))
444                    if is_unsupported_streaming_error(&frame.reason) =>
445                {
446                    log::error!("Tardis Machine rejected streaming");
447                    return false;
448                }
449                Some(Ok(tungstenite::Message::Close(frame))) => {
450                    if let Some(frame) = frame {
451                        log::debug!("WebSocket closed: {}", frame.code);
452                    } else {
453                        log::debug!("WebSocket closed");
454                    }
455                    return true;
456                }
457                Some(Ok(_)) => {}
458                Some(Err(_)) => {
459                    log::warn!("WebSocket error");
460                    return true;
461                }
462                None => {
463                    log::debug!("Stream ended");
464                    return true;
465                }
466            }
467        }
468    }
469}
470
471#[async_trait::async_trait(?Send)]
472impl DataClient for TardisDataClient {
473    fn client_id(&self) -> ClientId {
474        self.client_id
475    }
476
477    fn venue(&self) -> Option<Venue> {
478        None // Tardis is multi-venue
479    }
480
481    fn start(&mut self) -> anyhow::Result<()> {
482        log::info!("Starting {}", self.client_id);
483        Ok(())
484    }
485
486    fn stop(&mut self) -> anyhow::Result<()> {
487        log::info!("Stopping {}", self.client_id);
488        self.tasks.begin_shutdown();
489        self.is_connected.store(false, Ordering::Release);
490        Ok(())
491    }
492
493    fn reset(&mut self) -> anyhow::Result<()> {
494        self.tasks.begin_shutdown();
495        self.is_connected.store(false, Ordering::Release);
496        Ok(())
497    }
498
499    fn dispose(&mut self) -> anyhow::Result<()> {
500        self.stop()
501    }
502
503    fn is_connected(&self) -> bool {
504        self.is_connected.load(Ordering::Acquire)
505    }
506
507    fn is_disconnected(&self) -> bool {
508        !self.is_connected()
509    }
510
511    fn subscribe_mark_prices(&mut self, _cmd: SubscribeMarkPrices) -> anyhow::Result<()> {
512        Ok(())
513    }
514
515    fn subscribe_index_prices(&mut self, _cmd: SubscribeIndexPrices) -> anyhow::Result<()> {
516        Ok(())
517    }
518
519    fn subscribe_funding_rates(&mut self, _cmd: SubscribeFundingRates) -> anyhow::Result<()> {
520        Ok(())
521    }
522
523    fn unsubscribe_mark_prices(&mut self, _cmd: &UnsubscribeMarkPrices) -> anyhow::Result<()> {
524        Ok(())
525    }
526
527    fn unsubscribe_index_prices(&mut self, _cmd: &UnsubscribeIndexPrices) -> anyhow::Result<()> {
528        Ok(())
529    }
530
531    fn unsubscribe_funding_rates(&mut self, _cmd: &UnsubscribeFundingRates) -> anyhow::Result<()> {
532        Ok(())
533    }
534
535    async fn connect(&mut self) -> anyhow::Result<()> {
536        if self.is_connected() && self.tasks.is_open() {
537            return Ok(());
538        }
539
540        match (
541            self.config.options.as_slice(),
542            self.config.stream_options.as_slice(),
543        ) {
544            ([], []) => {
545                anyhow::bail!("Either replay `options` or `stream_options` must be provided")
546            }
547            ([], stream_options) => validate_stream_options(stream_options)?,
548            _ => {}
549        }
550
551        if !self.tasks.is_open() {
552            self.tasks
553                .finish_shutdown(Duration::from_secs(1), Duration::from_secs(2))
554                .await
555                .map_err(|e| anyhow::anyhow!("Failed to terminate Tardis tasks: {e}"))?;
556            self.tasks
557                .start_generation()
558                .map_err(|e| anyhow::anyhow!("Failed to start Tardis task generation: {e}"))?;
559            self.cancellation_token = self.tasks.cancellation_token();
560        }
561
562        let is_stream_mode = self.is_stream_mode();
563        let book_snapshot_output = self.config.book_snapshot_output.clone();
564        let extract_bbo_as_quotes = self.config.extract_bbo_as_quotes;
565
566        let http_client = TardisHttpClient::new(
567            self.config
568                .api_key
569                .as_ref()
570                .map(|value| value.expose_secret()),
571            self.config
572                .tardis_http_url
573                .as_ref()
574                .map(|value| value.expose_secret()),
575            None,
576            self.config.normalize_symbols,
577            self.config
578                .proxy_url
579                .as_ref()
580                .map(|value| value.expose_secret().to_owned()),
581        )?;
582
583        let exchanges: AHashSet<_> = if is_stream_mode {
584            self.config
585                .stream_options
586                .iter()
587                .map(|opt| opt.exchange)
588                .collect()
589        } else {
590            self.config.options.iter().map(|opt| opt.exchange).collect()
591        };
592
593        let base_url = resolve_ws_base_url(
594            self.config
595                .tardis_ws_url
596                .as_ref()
597                .map(|value| value.expose_secret()),
598        )?;
599        let (instrument_map, instruments) = http_client
600            .bootstrap_instruments(&exchanges)
601            .await
602            .map_err(|e| anyhow::anyhow!("Failed to bootstrap instruments: {e}"))?;
603
604        for instrument in instruments {
605            if let Err(e) = self.data_sender.send(DataEvent::Instrument(instrument)) {
606                log::error!("Failed to send instrument event: {e}");
607            }
608        }
609
610        let url = self.build_ws_url(&base_url)?;
611
612        let mode_label = if is_stream_mode { "stream" } else { "replay" };
613        log::info!("Connecting to Tardis Machine {mode_label}");
614        log::debug!("URL: {}", url_origin(&url));
615
616        let (ws_stream, _) = connect_async(build_ws_request(&url)?)
617            .await
618            .map_err(|e| anyhow::anyhow!("Failed to connect to Tardis Machine: {e}"))?;
619
620        log::info!("Connected to Tardis Machine");
621
622        self.is_connected.store(true, Ordering::Release);
623
624        if let Err(e) = self.spawn_ws_task(
625            ws_stream,
626            url,
627            instrument_map,
628            book_snapshot_output,
629            extract_bbo_as_quotes,
630            is_stream_mode,
631        ) {
632            self.is_connected.store(false, Ordering::Release);
633            self.tasks.begin_shutdown();
634            if let Err(teardown_error) = self
635                .tasks
636                .finish_shutdown(Duration::from_secs(1), Duration::from_secs(2))
637                .await
638            {
639                return Err(e.context(format!("Tardis startup teardown failed: {teardown_error}")));
640            }
641            return Err(e);
642        }
643
644        log::info!("Connected: {}", self.client_id);
645        Ok(())
646    }
647
648    async fn disconnect(&mut self) -> anyhow::Result<()> {
649        self.tasks.begin_shutdown();
650        let had_tasks = !self.tasks.is_empty();
651        let tasks_result = self
652            .tasks
653            .finish_shutdown(Duration::from_secs(1), Duration::from_secs(2))
654            .await
655            .map_err(|e| anyhow::anyhow!("Failed to terminate Tardis tasks: {e}"));
656
657        if had_tasks {
658            log::info!("Disconnected: {}", self.client_id);
659        }
660
661        self.is_connected.store(false, Ordering::Release);
662
663        tasks_result
664    }
665}
666
667/// Builds a Tardis Machine WebSocket handshake request carrying the Nautilus
668/// user agent.
669fn build_ws_request(url: &str) -> anyhow::Result<tungstenite::http::Request<()>> {
670    let mut request = url
671        .into_client_request()
672        .map_err(|e| anyhow::anyhow!("Failed to build Tardis Machine WebSocket request: {e}"))?;
673    request
674        .headers_mut()
675        .insert(USER_AGENT, HeaderValue::from_static(NAUTILUS_USER_AGENT));
676    Ok(request)
677}
678
679#[cfg(test)]
680mod tests {
681    use jiff::civil::Date;
682    use nautilus_common::live::runner::set_data_event_sender;
683    use rstest::rstest;
684
685    use super::*;
686    use crate::{
687        common::{consts::TARDIS_CLIENT_ID, enums::TardisExchange},
688        config::TardisDataClientConfig,
689        machine::types::{ReplayNormalizedRequestOptions, StreamNormalizedRequestOptions},
690    };
691
692    #[rstest]
693    #[case(vec![], "Either replay `options` or `stream_options` must be provided")]
694    #[case(
695        vec![TardisExchange::Bitmex],
696        "Unsupported Tardis streaming exchange: bitmex (historical-only; use replay options)"
697    )]
698    #[case(
699        vec![TardisExchange::Deribit, TardisExchange::Bitmex],
700        "Unsupported Tardis streaming exchange: bitmex (historical-only; use replay options)"
701    )]
702    #[tokio::test]
703    async fn test_connect_rejects_invalid_stream_options_before_network_access(
704        #[case] exchanges: Vec<TardisExchange>,
705        #[case] expected_error: &str,
706    ) {
707        setup_test_env();
708
709        let config = TardisDataClientConfig {
710            api_key: Some("test-key".into()),
711            tardis_http_url: Some("http://127.0.0.1:0".into()),
712            tardis_ws_url: Some("ws://127.0.0.1:0".into()),
713            stream_options: exchanges
714                .into_iter()
715                .map(|exchange| StreamNormalizedRequestOptions {
716                    exchange,
717                    symbols: None,
718                    data_types: vec!["trade".to_string()],
719                    with_disconnect_messages: None,
720                    timeout_interval_ms: None,
721                })
722                .collect(),
723            ..Default::default()
724        };
725
726        let mut client = TardisDataClient::new(*TARDIS_CLIENT_ID, config).unwrap();
727
728        let error = client.connect().await.unwrap_err();
729
730        assert_eq!(error.to_string(), expected_error);
731        assert!(client.is_disconnected());
732    }
733
734    #[rstest]
735    #[case(
736        "Error: Real-time streaming is not supported for exchange bitmex",
737        false
738    )]
739    #[case(
740        "Error: Real-time streaming is not supported for exchange coinflex",
741        false
742    )]
743    #[case("Too many subsequent errors when connecting to deribit WS API", true)]
744    #[tokio::test]
745    async fn test_ws_session_close_reconnect_policy(
746        #[case] reason: &'static str,
747        #[case] expected_reconnect: bool,
748    ) {
749        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
750        let address = listener.local_addr().unwrap();
751
752        let server = tokio::spawn(async move {
753            let (socket, _) = listener.accept().await.unwrap();
754            let mut websocket = tokio_tungstenite::accept_async(socket).await.unwrap();
755            websocket
756                .close(Some(tungstenite::protocol::CloseFrame {
757                    code: tungstenite::protocol::frame::coding::CloseCode::Error,
758                    reason: reason.into(),
759                }))
760                .await
761                .unwrap();
762        });
763
764        let (websocket, _) = connect_async(format!("ws://{address}")).await.unwrap();
765        let (sender, _receiver) = tokio::sync::mpsc::unbounded_channel();
766        let reconnect = tokio::time::timeout(
767            Duration::from_secs(2),
768            TardisDataClient::run_ws_session(
769                websocket,
770                &CancellationToken::new(),
771                &EventSender::from(sender),
772                &AHashMap::new(),
773                &BookSnapshotOutput::Deltas,
774                false,
775            ),
776        )
777        .await
778        .unwrap();
779        server.await.unwrap();
780
781        assert_eq!(reconnect, expected_reconnect);
782    }
783
784    #[tokio::test]
785    async fn test_unsupported_exchange_finishes_stream_task_without_reconnecting() {
786        setup_test_env();
787
788        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
789        let url = format!("ws://{}", listener.local_addr().unwrap());
790
791        let server = tokio::spawn(async move {
792            let (socket, _) = listener.accept().await.unwrap();
793            let mut websocket = tokio_tungstenite::accept_async(socket).await.unwrap();
794            websocket
795                .close(Some(tungstenite::protocol::CloseFrame {
796                    code: tungstenite::protocol::frame::coding::CloseCode::Error,
797                    reason: "Error: Real-time streaming is not supported for exchange bitmex"
798                        .into(),
799                }))
800                .await
801                .unwrap();
802
803            listener
804        });
805
806        let (websocket, _) = connect_async(&url).await.unwrap();
807        let client =
808            TardisDataClient::new(*TARDIS_CLIENT_ID, TardisDataClientConfig::default()).unwrap();
809        client.is_connected.store(true, Ordering::Release);
810        client
811            .spawn_ws_task(
812                websocket,
813                url,
814                AHashMap::new(),
815                BookSnapshotOutput::Deltas,
816                false,
817                true,
818            )
819            .unwrap();
820
821        tokio::time::timeout(Duration::from_secs(2), async {
822            while !client.tasks.all_finished() {
823                tokio::task::yield_now().await;
824            }
825        })
826        .await
827        .expect("Unsupported exchange must finish the task without reconnecting");
828
829        let listener = server.await.unwrap();
830
831        assert!(client.is_disconnected());
832        assert!(
833            tokio::time::timeout(Duration::from_millis(50), listener.accept())
834                .await
835                .is_err()
836        );
837    }
838
839    #[rstest]
840    fn test_build_ws_request_rejects_invalid_url() {
841        assert!(build_ws_request("not a url").is_err());
842    }
843
844    fn setup_test_env() {
845        use std::cell::OnceCell;
846
847        thread_local! {
848            static INIT: OnceCell<()> = const { OnceCell::new() };
849        }
850
851        INIT.with(|cell| {
852            cell.get_or_init(|| {
853                let (sender, _receiver) = tokio::sync::mpsc::unbounded_channel::<DataEvent>();
854                set_data_event_sender(sender);
855            });
856        });
857    }
858
859    #[rstest]
860    fn test_build_ws_url_injects_derivative_ticker() {
861        setup_test_env();
862
863        let config = TardisDataClientConfig {
864            options: vec![ReplayNormalizedRequestOptions {
865                exchange: TardisExchange::BinanceFutures,
866                symbols: Some(vec!["BTCUSDT".to_string()]),
867                from: Date::new(2024, 1, 1).unwrap(),
868                to: Date::new(2024, 1, 2).unwrap(),
869                data_types: vec!["trade".to_string()],
870                with_disconnect_messages: Some(false),
871            }],
872            ..Default::default()
873        };
874
875        let client = TardisDataClient::new(*TARDIS_CLIENT_ID, config).unwrap();
876        let url = client.build_ws_url("ws://localhost:8001").unwrap();
877
878        assert!(
879            url.contains("derivative_ticker"),
880            "URL should contain derivative_ticker but was: {url}"
881        );
882        assert!(url.contains("trade"), "URL should still contain trade");
883    }
884
885    #[rstest]
886    fn test_build_ws_url_does_not_duplicate_derivative_ticker() {
887        setup_test_env();
888
889        let config = TardisDataClientConfig {
890            options: vec![ReplayNormalizedRequestOptions {
891                exchange: TardisExchange::BinanceFutures,
892                symbols: Some(vec!["BTCUSDT".to_string()]),
893                from: Date::new(2024, 1, 1).unwrap(),
894                to: Date::new(2024, 1, 2).unwrap(),
895                data_types: vec!["trade".to_string(), "derivative_ticker".to_string()],
896                with_disconnect_messages: Some(false),
897            }],
898            ..Default::default()
899        };
900
901        let client = TardisDataClient::new(*TARDIS_CLIENT_ID, config).unwrap();
902        let ws_url = client.build_ws_url("ws://localhost:8001").unwrap();
903
904        let decoded = urlencoding::decode(ws_url.split("options=").nth(1).unwrap()).unwrap();
905        let count = decoded.matches("derivative_ticker").count();
906        assert_eq!(count, 1, "derivative_ticker should appear exactly once");
907    }
908
909    #[rstest]
910    fn test_stop_marks_client_disconnected_synchronously() {
911        setup_test_env();
912
913        let mut client =
914            TardisDataClient::new(*TARDIS_CLIENT_ID, TardisDataClientConfig::default()).unwrap();
915        client.is_connected.store(true, Ordering::Release);
916
917        client.stop().unwrap();
918
919        assert!(client.is_disconnected());
920    }
921
922    #[rstest]
923    fn test_reset_marks_client_disconnected_synchronously() {
924        setup_test_env();
925
926        let mut client =
927            TardisDataClient::new(*TARDIS_CLIENT_ID, TardisDataClientConfig::default()).unwrap();
928        client.is_connected.store(true, Ordering::Release);
929
930        client.reset().unwrap();
931
932        assert!(client.is_disconnected());
933    }
934}
nautilus_tardis/data.rs

nautilus_tardis/
data.rs
