nautilus_tardis/machine/
client.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
16use std::{
17    fmt::Debug,
18    sync::{
19        Arc,
20        atomic::{AtomicBool, Ordering},
21    },
22};
23
24use ahash::AHashMap;
25use futures_util::{Stream, StreamExt, pin_mut};
26use nautilus_model::data::Data;
27use nautilus_network::url_origin;
28
29use super::{
30    Error,
31    message::WsMessage,
32    replay_normalized, stream_normalized,
33    types::{
34        ReplayNormalizedRequestOptions, StreamNormalizedRequestOptions, TardisInstrumentKey,
35        TardisInstrumentMiniInfo,
36    },
37};
38use crate::{
39    common::urls::resolve_ws_base_url, config::BookSnapshotOutput,
40    machine::parse::parse_tardis_ws_message_data,
41};
42
43/// Provides a client for connecting to a [Tardis Machine Server](https://docs.tardis.dev/api/tardis-machine).
44#[cfg_attr(
45    feature = "python",
46    pyo3::pyclass(module = "nautilus_trader.adapters.tardis", from_py_object)
47)]
48#[cfg_attr(
49    feature = "python",
50    pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.adapters.tardis")
51)]
52#[derive(Clone)]
53pub struct TardisMachineClient {
54    pub base_url: String,
55    pub replay_signal: Arc<AtomicBool>,
56    pub stream_signal: Arc<AtomicBool>,
57    pub instruments: AHashMap<TardisInstrumentKey, Arc<TardisInstrumentMiniInfo>>,
58    pub normalize_symbols: bool,
59    pub book_snapshot_output: BookSnapshotOutput,
60    pub extract_bbo_as_quotes: bool,
61}
62
63impl Debug for TardisMachineClient {
64    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
65        f.debug_struct(stringify!(TardisMachineClient))
66            .field("base_url", &url_origin(&self.base_url))
67            .field("replay_signal", &self.replay_signal)
68            .field("stream_signal", &self.stream_signal)
69            .field("instruments", &self.instruments)
70            .field("normalize_symbols", &self.normalize_symbols)
71            .field("book_snapshot_output", &self.book_snapshot_output)
72            .field("extract_bbo_as_quotes", &self.extract_bbo_as_quotes)
73            .finish()
74    }
75}
76
77impl TardisMachineClient {
78    /// Creates a new [`TardisMachineClient`] instance.
79    ///
80    /// # Errors
81    ///
82    /// Returns an error if `base_url` is not provided and `TARDIS_MACHINE_WS_URL` env var is missing.
83    pub fn new(
84        base_url: Option<&str>,
85        normalize_symbols: bool,
86        book_snapshot_output: BookSnapshotOutput,
87    ) -> anyhow::Result<Self> {
88        let base_url = resolve_ws_base_url(base_url)?;
89
90        Ok(Self {
91            base_url,
92            replay_signal: Arc::new(AtomicBool::new(false)),
93            stream_signal: Arc::new(AtomicBool::new(false)),
94            instruments: AHashMap::new(),
95            normalize_symbols,
96            book_snapshot_output,
97            extract_bbo_as_quotes: false,
98        })
99    }
100
101    pub fn add_instrument_info(&mut self, info: TardisInstrumentMiniInfo) {
102        let key = info.as_tardis_instrument_key();
103        self.instruments.insert(key, Arc::new(info));
104    }
105
106    /// Returns `true` if `close()` has been called.
107    ///
108    /// This checks that both replay and stream signals have been set,
109    /// which only occurs when `close()` is explicitly called.
110    #[must_use]
111    pub fn is_closed(&self) -> bool {
112        // Use Acquire ordering to synchronize with Release stores in close()
113        self.replay_signal.load(Ordering::Acquire) && self.stream_signal.load(Ordering::Acquire)
114    }
115
116    pub fn close(&mut self) {
117        log::debug!("Closing");
118
119        // Use Release ordering to ensure visibility to Acquire loads in is_closed()
120        self.replay_signal.store(true, Ordering::Release);
121        self.stream_signal.store(true, Ordering::Release);
122
123        log::debug!("Closed");
124    }
125
126    /// Connects to the Tardis Machine replay WebSocket and yields parsed `Data` items.
127    ///
128    /// # Errors
129    ///
130    /// Returns an error if the WebSocket connection cannot be established.
131    pub async fn replay(
132        &self,
133        options: Vec<ReplayNormalizedRequestOptions>,
134    ) -> Result<impl Stream<Item = Result<Data, Error>>, Error> {
135        let stream = replay_normalized(&self.base_url, options, self.replay_signal.clone()).await?;
136
137        // We use Box::pin to heap-allocate the stream and ensure it implements
138        // Unpin for safe async handling across lifetimes.
139        Ok(handle_ws_stream(
140            Box::pin(stream),
141            None,
142            Some(self.instruments.clone()),
143            self.book_snapshot_output.clone(),
144            self.extract_bbo_as_quotes,
145        ))
146    }
147
148    /// Connects to the Tardis Machine stream WebSocket for a single instrument and yields parsed `Data` items.
149    ///
150    /// # Errors
151    ///
152    /// Returns an error if the WebSocket connection cannot be established.
153    pub async fn stream(
154        &self,
155        instrument: TardisInstrumentMiniInfo,
156        options: Vec<StreamNormalizedRequestOptions>,
157    ) -> Result<impl Stream<Item = Result<Data, Error>>, Error> {
158        let stream = stream_normalized(&self.base_url, options, self.stream_signal.clone()).await?;
159
160        // We use Box::pin to heap-allocate the stream and ensure it implements
161        // Unpin for safe async handling across lifetimes.
162        Ok(handle_ws_stream(
163            Box::pin(stream),
164            Some(Arc::new(instrument)),
165            None,
166            self.book_snapshot_output.clone(),
167            self.extract_bbo_as_quotes,
168        ))
169    }
170}
171
172fn handle_ws_stream<S>(
173    stream: S,
174    instrument: Option<Arc<TardisInstrumentMiniInfo>>,
175    instrument_map: Option<AHashMap<TardisInstrumentKey, Arc<TardisInstrumentMiniInfo>>>,
176    book_snapshot_output: BookSnapshotOutput,
177    extract_bbo_as_quotes: bool,
178) -> impl Stream<Item = Result<Data, Error>>
179where
180    S: Stream<Item = Result<WsMessage, Error>> + Unpin,
181{
182    assert!(
183        instrument.is_some() || instrument_map.is_some(),
184        "Either `instrument` or `instrument_map` must be provided"
185    );
186
187    async_stream::stream! {
188        pin_mut!(stream);
189
190        while let Some(result) = stream.next().await {
191            match result {
192                Ok(msg) => {
193                    if matches!(msg, WsMessage::Disconnect(_)) {
194                        log::debug!("Received disconnect message: {msg:?}");
195                        continue;
196                    }
197
198                    let info = instrument.clone().or_else(|| {
199                        instrument_map
200                            .as_ref()
201                            .and_then(|map| determine_instrument_info(&msg, map))
202                    });
203
204                    if let Some(info) = info {
205                        for data in parse_tardis_ws_message_data(
206                            msg,
207                            &info,
208                            &book_snapshot_output,
209                            extract_bbo_as_quotes,
210                        ) {
211                            yield Ok(data);
212                        }
213                    } else {
214                        log::debug!("Missing instrument info for message: {msg:?}");
215                        yield Err(Error::ConnectionClosed {
216                            reason: super::MISSING_INSTRUMENT_INFO.to_string()
217                        });
218                        break;
219                    }
220                }
221                Err(e) => {
222                    yield Err(e);
223                    break;
224                }
225            }
226        }
227    }
228}
229
230pub fn determine_instrument_info(
231    msg: &WsMessage,
232    instrument_map: &AHashMap<TardisInstrumentKey, Arc<TardisInstrumentMiniInfo>>,
233) -> Option<Arc<TardisInstrumentMiniInfo>> {
234    let key = match msg {
235        WsMessage::BookChange(msg) => TardisInstrumentKey::new(msg.symbol, msg.exchange),
236        WsMessage::BookSnapshot(msg) => TardisInstrumentKey::new(msg.symbol, msg.exchange),
237        WsMessage::Trade(msg) => TardisInstrumentKey::new(msg.symbol, msg.exchange),
238        WsMessage::TradeBar(msg) => TardisInstrumentKey::new(msg.symbol, msg.exchange),
239        WsMessage::DerivativeTicker(msg) => TardisInstrumentKey::new(msg.symbol, msg.exchange),
240        WsMessage::OptionSummary(msg) => TardisInstrumentKey::new(msg.symbol, msg.exchange),
241        WsMessage::Disconnect(_) => return None,
242    };
243
244    if let Some(inst) = instrument_map.get(&key) {
245        Some(inst.clone())
246    } else {
247        log::error!("Instrument definition info not available for {key:?}");
248        None
249    }
250}
251
252#[cfg(test)]
253mod tests {
254    use rstest::rstest;
255
256    use super::*;
257
258    #[rstest::rstest]
259    #[case("wss://stream.example.com/feed", "wss://stream.example.com")]
260    #[case(
261        "wss://SYN-user:SYN-password@stream.example.com:9443/SYN-path?token=SYN-query#SYN-fragment",
262        "wss://stream.example.com:9443"
263    )]
264    fn test_debug_preserves_url_origin(#[case] url: &str, #[case] origin: &str) {
265        let client =
266            TardisMachineClient::new(Some(url), true, BookSnapshotOutput::default()).unwrap();
267        let debug = format!("{client:?}");
268        assert!(debug.contains(&format!("base_url: {origin:?}")));
269        assert!(!debug.contains("SYN-"));
270        assert_eq!(client.base_url, url);
271    }
272
273    #[rstest]
274    fn test_is_closed_initial_state() {
275        let client = TardisMachineClient::new(
276            Some("ws://localhost:8001"),
277            false,
278            BookSnapshotOutput::Deltas,
279        )
280        .unwrap();
281        // Initially neither signal is set, so is_closed should be false
282        assert!(!client.is_closed());
283    }
284
285    #[rstest]
286    fn test_is_closed_after_close() {
287        let mut client = TardisMachineClient::new(
288            Some("ws://localhost:8001"),
289            false,
290            BookSnapshotOutput::Deltas,
291        )
292        .unwrap();
293        client.close();
294        // After close(), both signals are set, so is_closed should be true
295        assert!(client.is_closed());
296    }
297
298    #[rstest]
299    fn test_is_closed_partial_signal() {
300        let client = TardisMachineClient::new(
301            Some("ws://localhost:8001"),
302            false,
303            BookSnapshotOutput::Deltas,
304        )
305        .unwrap();
306        // Set only one signal - is_closed should still be false
307        // (since close() wasn't called, which sets both)
308        client.replay_signal.store(true, Ordering::Release);
309        assert!(!client.is_closed());
310
311        client.stream_signal.store(true, Ordering::Release);
312        // Now both are set, so is_closed should be true
313        assert!(client.is_closed());
314    }
315}
nautilus_tardis/machine/client.rs

nautilus_tardis/machine/
client.rs
