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::string::urlencoding;
34use tokio_tungstenite::{
35    connect_async,
36    tungstenite::{self, 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 (ws_stream, ws_resp) = connect_async(url).await?;
129
130    handle_connection_response(&ws_resp)?;
131    log::debug!("Connected to {base_url}");
132
133    Ok(stream! {
134        let (writer, mut reader) = ws_stream.split();
135        let cancel = CancellationToken::new();
136        let heartbeat = heartbeat(writer, cancel.child_token());
137        pin_mut!(heartbeat);
138        let mut heartbeat_active = true;
139        let _cancel_heartbeat = cancel.drop_guard();
140
141        // Timeout awaiting the next record before checking signal
142        let timeout = Duration::from_millis(10);
143
144        log::debug!("Streaming from websocket...");
145
146        loop {
147            if signal.load(Ordering::Relaxed) {
148                log::debug!("Shutdown signal received");
149                break;
150            }
151
152            let result = tokio::select! {
153                result = tokio::time::timeout(timeout, reader.next()) => result,
154                () = &mut heartbeat, if heartbeat_active => {
155                    heartbeat_active = false;
156                    continue;
157                }
158            };
159            let msg = match result {
160                Ok(msg) => msg,
161                Err(_) => continue, // Timeout
162            };
163
164            match msg {
165                Some(Ok(msg)) => match msg {
166                    tungstenite::Message::Frame(_)
167                    | tungstenite::Message::Binary(_)
168                    | tungstenite::Message::Pong(_)
169                    | tungstenite::Message::Ping(_) => {
170                        log::trace!("Received {msg:?}");
171                    }
172                    tungstenite::Message::Close(Some(frame)) => {
173                        let reason = frame.reason.to_string();
174                        if frame.code == CloseCode::Normal {
175                            log::debug!("Connection closed normally: {reason}");
176                        } else {
177                            log::warn!(
178                                "Connection closed abnormally with code: {:?}, reason: {reason}", frame.code
179                            );
180                            yield Err(Error::ConnectionClosed { reason });
181                        }
182                        break;
183                    }
184                    tungstenite::Message::Close(None) => {
185                        log::warn!("Connection closed without a frame");
186                        yield Err(Error::ConnectionClosed {
187                            reason: "No close frame provided".to_string()
188                        });
189                        break;
190                    }
191                    tungstenite::Message::Text(msg) => {
192                        match serde_json::from_str::<WsMessage>(&msg) {
193                            Ok(parsed_msg) => yield Ok(parsed_msg),
194                            Err(e) => {
195                                log::error!("Failed to deserialize message: {msg}. Error: {e}");
196                                yield Err(Error::Deserialization(e));
197                            }
198                        }
199                    }
200                },
201                Some(Err(e)) => {
202                    log::warn!("WebSocket error: {e}");
203                    yield Err(Error::ConnectFailed(e));
204                    break;
205                }
206                None => {
207                    log::warn!("Connection closed unexpectedly");
208                    yield Err(Error::ConnectionClosed {
209                        reason: "Unexpected connection close".to_string(),
210                    });
211                    break;
212                }
213            }
214        }
215
216        log::debug!("Shutdown stream");
217    })
218}
219
220fn handle_connection_response(
221    ws_resp: &tungstenite::http::Response<Option<Vec<u8>>>,
222) -> Result<()> {
223    if ws_resp.status() != tungstenite::http::StatusCode::SWITCHING_PROTOCOLS {
224        return match ws_resp.body() {
225            Some(resp) => Err(Error::ConnectRejected {
226                status: ws_resp.status(),
227                reason: String::from_utf8_lossy(resp).to_string(),
228            }),
229            None => Err(Error::ConnectRejected {
230                status: ws_resp.status(),
231                reason: "Unknown reason".to_string(),
232            }),
233        };
234    }
235    Ok(())
236}
237
238async fn heartbeat<S>(mut sender: S, cancel: CancellationToken)
239where
240    S: Sink<tungstenite::Message> + Unpin,
241    S::Error: std::fmt::Display,
242{
243    let mut heartbeat_interval = tokio::time::interval(Duration::from_secs(10));
244
245    loop {
246        tokio::select! {
247            _ = heartbeat_interval.tick() => {}
248            () = cancel.cancelled() => break,
249        }
250
251        log::trace!("Sending PING");
252
253        tokio::select! {
254            result = sender.send(tungstenite::Message::Ping(vec![].into())) => {
255                if let Err(e) = result {
256                    log::debug!("Heartbeat send failed (connection closed): {e}");
257                    break;
258                }
259            }
260            () = cancel.cancelled() => break,
261        }
262    }
263
264    log::debug!("Heartbeat task exiting");
265}
266
267#[cfg(test)]
268mod tests {
269    use std::{
270        convert::Infallible,
271        pin::Pin,
272        task::{Context, Poll},
273        time::Duration,
274    };
275
276    use futures_util::Sink;
277    use rstest::rstest;
278    use tokio_tungstenite::tungstenite::Message;
279    use tokio_util::sync::CancellationToken;
280
281    use super::heartbeat;
282
283    struct StallSink;
284
285    impl Sink<Message> for StallSink {
286        type Error = Infallible;
287
288        fn poll_ready(
289            self: Pin<&mut Self>,
290            _cx: &mut Context<'_>,
291        ) -> Poll<Result<(), Self::Error>> {
292            Poll::Pending
293        }
294
295        fn start_send(self: Pin<&mut Self>, _item: Message) -> Result<(), Self::Error> {
296            Ok(())
297        }
298
299        fn poll_flush(
300            self: Pin<&mut Self>,
301            _cx: &mut Context<'_>,
302        ) -> Poll<Result<(), Self::Error>> {
303            Poll::Pending
304        }
305
306        fn poll_close(
307            self: Pin<&mut Self>,
308            _cx: &mut Context<'_>,
309        ) -> Poll<Result<(), Self::Error>> {
310            Poll::Ready(Ok(()))
311        }
312    }
313
314    #[rstest]
315    #[tokio::test]
316    async fn test_heartbeat_exits_on_cancel_during_stalled_send() {
317        let cancel = CancellationToken::new();
318        let task = tokio::spawn(heartbeat(StallSink, cancel.clone()));
319
320        tokio::task::yield_now().await;
321        cancel.cancel();
322
323        tokio::time::timeout(Duration::from_secs(1), task)
324            .await
325            .expect("heartbeat should exit after cancel during a stalled send")
326            .unwrap();
327    }
328}
nautilus_tardis/machine/mod.rs

nautilus_tardis/machine/
mod.rs
