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

nautilus_tardis/
data.rs
