nautilus_tardis/machine/
mod.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
16pub mod cache;
17pub mod client;
18pub mod message;
19pub mod parse;
20pub mod types;
21
22use std::{
23    sync::{
24        Arc,
25        atomic::{AtomicBool, Ordering},
26    },
27    time::Duration,
28};
29
30use async_stream::stream;
31use futures_util::{Sink, SinkExt, Stream, StreamExt, pin_mut};
32use message::WsMessage;
33use nautilus_core::{consts::NAUTILUS_USER_AGENT, string::urlencoding};
34use tokio_tungstenite::{
35    connect_async,
36    tungstenite::{self, client::IntoClientRequest, protocol::frame::coding::CloseCode},
37};
38use tokio_util::sync::CancellationToken;
39use types::{ReplayNormalizedRequestOptions, StreamNormalizedRequestOptions};
40
41pub use crate::machine::client::TardisMachineClient;
42
43pub type Result<T> = std::result::Result<T, Error>;
44
45/// The error that could happen while interacting with Tardis Machine Server.
46#[derive(Debug, thiserror::Error)]
47pub enum Error {
48    /// An error that could happen when an empty options array was given.
49    #[error("Options cannot be empty")]
50    EmptyOptions,
51    /// An error when failed to connect to Tardis' websocket connection.
52    #[error("Failed to connect: {0}")]
53    ConnectFailed(#[from] tungstenite::Error),
54    /// An error when WS connection to the machine server was rejected.
55    #[error("Connection rejected: {reason}")]
56    ConnectRejected {
57        /// The status code for the initial WS connection.
58        status: tungstenite::http::StatusCode,
59        /// The reason why the connection was rejected.
60        reason: String,
61    },
62    /// An error where the websocket connection was closed unexpectedly by Tardis.
63    #[error("Connection closed: {reason}")]
64    ConnectionClosed {
65        /// The reason why the connection was closed.
66        reason: String,
67    },
68    /// An error when deserializing the response from Tardis.
69    #[error("Failed to deserialize message: {0}")]
70    Deserialization(#[from] serde_json::Error),
71}
72
73/// Connects to the Tardis Machine WS replay endpoint and returns a stream of WebSocket messages.
74///
75/// # Errors
76///
77/// Returns `Error::EmptyOptions` if no options provided,
78/// or `Error::ConnectFailed`/`Error::ConnectRejected` if connection fails.
79pub async fn replay_normalized(
80    base_url: &str,
81    options: Vec<ReplayNormalizedRequestOptions>,
82    signal: Arc<AtomicBool>,
83) -> Result<impl Stream<Item = Result<WsMessage>>> {
84    if options.is_empty() {
85        return Err(Error::EmptyOptions);
86    }
87
88    let path = format!("{base_url}/ws-replay-normalized?options=");
89    let options = serde_json::to_string(&options)?;
90
91    let plain_url = format!("{path}{options}");
92    log::debug!("Connecting to {plain_url}");
93
94    let url = format!("{path}{}", urlencoding::encode(&options));
95    stream_from_websocket(base_url, url, signal).await
96}
97
98/// Connects to the Tardis Machine WS streaming endpoint and returns a stream of WebSocket messages.
99///
100/// # Errors
101///
102/// Returns `Error::EmptyOptions` if no options provided,
103/// or `Error::ConnectFailed`/`Error::ConnectRejected` if connection fails.
104pub async fn stream_normalized(
105    base_url: &str,
106    options: Vec<StreamNormalizedRequestOptions>,
107    signal: Arc<AtomicBool>,
108) -> Result<impl Stream<Item = Result<WsMessage>>> {
109    if options.is_empty() {
110        return Err(Error::EmptyOptions);
111    }
112
113    let path = format!("{base_url}/ws-stream-normalized?options=");
114    let options = serde_json::to_string(&options)?;
115
116    let plain_url = format!("{path}{options}");
117    log::debug!("Connecting to {plain_url}");
118
119    let url = format!("{path}{}", urlencoding::encode(&options));
120    stream_from_websocket(base_url, url, signal).await
121}
122
123async fn stream_from_websocket(
124    base_url: &str,
125    url: String,
126    signal: Arc<AtomicBool>,
127) -> Result<impl Stream<Item = Result<WsMessage>>> {
128    let mut request = url.into_client_request()?;
129    request.headers_mut().insert(
130        tungstenite::http::header::USER_AGENT,
131        tungstenite::http::HeaderValue::from_static(NAUTILUS_USER_AGENT),
132    );
133
134    let (ws_stream, ws_resp) = connect_async(request).await?;
135
136    handle_connection_response(&ws_resp)?;
137    log::debug!("Connected to {base_url}");
138
139    Ok(stream! {
140        let (writer, mut reader) = ws_stream.split();
141        let cancel = CancellationToken::new();
142        let heartbeat = heartbeat(writer, cancel.child_token());
143        pin_mut!(heartbeat);
144        let mut heartbeat_active = true;
145        let _cancel_heartbeat = cancel.drop_guard();
146
147        // Timeout awaiting the next record before checking signal
148        let timeout = Duration::from_millis(10);
149
150        log::debug!("Streaming from websocket...");
151
152        loop {
153            if signal.load(Ordering::Relaxed) {
154                log::debug!("Shutdown signal received");
155                break;
156            }
157
158            let result = tokio::select! {
159                result = tokio::time::timeout(timeout, reader.next()) => result,
160                () = &mut heartbeat, if heartbeat_active => {
161                    heartbeat_active = false;
162                    continue;
163                }
164            };
165            let msg = match result {
166                Ok(msg) => msg,
167                Err(_) => continue, // Timeout
168            };
169
170            match msg {
171                Some(Ok(msg)) => match msg {
172                    tungstenite::Message::Frame(_)
173                    | tungstenite::Message::Binary(_)
174                    | tungstenite::Message::Pong(_)
175                    | tungstenite::Message::Ping(_) => {
176                        log::trace!("Received {msg:?}");
177                    }
178                    tungstenite::Message::Close(Some(frame)) => {
179                        let reason = frame.reason.to_string();
180                        if frame.code == CloseCode::Normal {
181                            log::debug!("Connection closed normally: {reason}");
182                        } else {
183                            log::warn!(
184                                "Connection closed abnormally with code: {:?}, reason: {reason}", frame.code
185                            );
186                            yield Err(Error::ConnectionClosed { reason });
187                        }
188                        break;
189                    }
190                    tungstenite::Message::Close(None) => {
191                        log::warn!("Connection closed without a frame");
192                        yield Err(Error::ConnectionClosed {
193                            reason: "No close frame provided".to_string()
194                        });
195                        break;
196                    }
197                    tungstenite::Message::Text(msg) => {
198                        match serde_json::from_str::<WsMessage>(&msg) {
199                            Ok(parsed_msg) => yield Ok(parsed_msg),
200                            Err(e) => {
201                                log::error!("Failed to deserialize message: {msg}. Error: {e}");
202                                yield Err(Error::Deserialization(e));
203                            }
204                        }
205                    }
206                },
207                Some(Err(e)) => {
208                    log::warn!("WebSocket error: {e}");
209                    yield Err(Error::ConnectFailed(e));
210                    break;
211                }
212                None => {
213                    log::warn!("Connection closed unexpectedly");
214                    yield Err(Error::ConnectionClosed {
215                        reason: "Unexpected connection close".to_string(),
216                    });
217                    break;
218                }
219            }
220        }
221
222        log::debug!("Shutdown stream");
223    })
224}
225
226fn handle_connection_response(
227    ws_resp: &tungstenite::http::Response<Option<Vec<u8>>>,
228) -> Result<()> {
229    if ws_resp.status() != tungstenite::http::StatusCode::SWITCHING_PROTOCOLS {
230        return match ws_resp.body() {
231            Some(resp) => Err(Error::ConnectRejected {
232                status: ws_resp.status(),
233                reason: String::from_utf8_lossy(resp).to_string(),
234            }),
235            None => Err(Error::ConnectRejected {
236                status: ws_resp.status(),
237                reason: "Unknown reason".to_string(),
238            }),
239        };
240    }
241    Ok(())
242}
243
244async fn heartbeat<S>(mut sender: S, cancel: CancellationToken)
245where
246    S: Sink<tungstenite::Message> + Unpin,
247    S::Error: std::fmt::Display,
248{
249    let mut heartbeat_interval = tokio::time::interval(Duration::from_secs(10));
250
251    loop {
252        tokio::select! {
253            _ = heartbeat_interval.tick() => {}
254            () = cancel.cancelled() => break,
255        }
256
257        log::trace!("Sending PING");
258
259        tokio::select! {
260            result = sender.send(tungstenite::Message::Ping(vec![].into())) => {
261                if let Err(e) = result {
262                    log::debug!("Heartbeat send failed (connection closed): {e}");
263                    break;
264                }
265            }
266            () = cancel.cancelled() => break,
267        }
268    }
269
270    log::debug!("Heartbeat task exiting");
271}
272
273#[cfg(test)]
274mod tests {
275    use std::{
276        convert::Infallible,
277        pin::Pin,
278        task::{Context, Poll},
279        time::Duration,
280    };
281
282    use futures_util::Sink;
283    use rstest::rstest;
284    use tokio_tungstenite::tungstenite::Message;
285    use tokio_util::sync::CancellationToken;
286
287    use super::heartbeat;
288
289    struct StallSink;
290
291    impl Sink<Message> for StallSink {
292        type Error = Infallible;
293
294        fn poll_ready(
295            self: Pin<&mut Self>,
296            _cx: &mut Context<'_>,
297        ) -> Poll<Result<(), Self::Error>> {
298            Poll::Pending
299        }
300
301        fn start_send(self: Pin<&mut Self>, _item: Message) -> Result<(), Self::Error> {
302            Ok(())
303        }
304
305        fn poll_flush(
306            self: Pin<&mut Self>,
307            _cx: &mut Context<'_>,
308        ) -> Poll<Result<(), Self::Error>> {
309            Poll::Pending
310        }
311
312        fn poll_close(
313            self: Pin<&mut Self>,
314            _cx: &mut Context<'_>,
315        ) -> Poll<Result<(), Self::Error>> {
316            Poll::Ready(Ok(()))
317        }
318    }
319
320    #[rstest]
321    #[tokio::test]
322    async fn test_heartbeat_exits_on_cancel_during_stalled_send() {
323        let cancel = CancellationToken::new();
324        let task = tokio::spawn(heartbeat(StallSink, cancel.clone()));
325
326        tokio::task::yield_now().await;
327        cancel.cancel();
328
329        tokio::time::timeout(Duration::from_secs(1), task)
330            .await
331            .expect("heartbeat should exit after cancel during a stalled send")
332            .unwrap();
333    }
334}
nautilus_tardis/machine/mod.rs

nautilus_tardis/machine/
mod.rs
