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

nautilus_tardis/machine/
client.rs
