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

nautilus_tardis/machine/
client.rs
