nautilus_architect_ax/websocket/data/
handler.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
16//! Market data WebSocket message handler for Ax.
17
18use std::{
19    collections::VecDeque,
20    sync::{
21        Arc,
22        atomic::{AtomicBool, Ordering},
23    },
24};
25
26use ahash::AHashMap;
27use nautilus_network::websocket::{SubscriptionState, WebSocketClient};
28use tokio_tungstenite::tungstenite::Message;
29use ustr::Ustr;
30
31use super::AxMdSubscriptionSpec;
32use crate::{
33    common::enums::{AxCandleWidth, AxMdRequestType},
34    websocket::{
35        messages::{
36            AxDataWsMessage, AxMdMessage, AxMdSubscribe, AxMdSubscribeCandles, AxMdUnsubscribe,
37            AxMdUnsubscribeCandles,
38        },
39        parse::parse_md_message,
40    },
41};
42
43/// Commands sent from the outer client to the inner message handler.
44#[derive(Debug)]
45pub enum HandlerCommand {
46    /// Set the WebSocket client for this handler.
47    SetClient(WebSocketClient),
48    /// Disconnect the WebSocket connection.
49    Disconnect,
50    /// Replay all subscriptions after a reconnection.
51    ReplaySubscriptions,
52    /// Subscribe to market data for a symbol.
53    Subscribe {
54        /// Request ID for correlation.
55        request_id: i64,
56        /// Instrument symbol.
57        symbol: Ustr,
58        /// Market data subscription options.
59        spec: AxMdSubscriptionSpec,
60    },
61    /// Unsubscribe from market data for a symbol.
62    Unsubscribe {
63        /// Request ID for correlation.
64        request_id: i64,
65        /// Instrument symbol.
66        symbol: Ustr,
67        /// Subscription topic for state tracking.
68        topic: String,
69    },
70    /// Subscribe to candle data for a symbol.
71    SubscribeCandles {
72        /// Request ID for correlation.
73        request_id: i64,
74        /// Instrument symbol.
75        symbol: Ustr,
76        /// Candle width/interval.
77        width: AxCandleWidth,
78    },
79    /// Unsubscribe from candle data for a symbol.
80    UnsubscribeCandles {
81        /// Request ID for correlation.
82        request_id: i64,
83        /// Instrument symbol.
84        symbol: Ustr,
85        /// Candle width/interval.
86        width: AxCandleWidth,
87        /// Subscription topic for state tracking.
88        topic: String,
89    },
90}
91
92/// Market data feed handler that processes WebSocket messages.
93///
94/// Runs in a dedicated Tokio task and owns the WebSocket client exclusively.
95/// Emits raw venue types for downstream consumers to parse.
96pub(crate) struct AxMdWsFeedHandler {
97    signal: Arc<AtomicBool>,
98    inner: Option<WebSocketClient>,
99    cmd_rx: tokio::sync::mpsc::UnboundedReceiver<HandlerCommand>,
100    raw_rx: tokio::sync::mpsc::UnboundedReceiver<Message>,
101    subscriptions: SubscriptionState,
102    message_queue: VecDeque<AxDataWsMessage>,
103    replay_request_id: i64,
104    needs_subscription_replay: bool,
105    pending_subscription_requests: AHashMap<i64, PendingSubscriptionRequest>,
106}
107
108impl AxMdWsFeedHandler {
109    /// Creates a new [`AxMdWsFeedHandler`] instance.
110    #[must_use]
111    pub(crate) fn new(
112        signal: Arc<AtomicBool>,
113        cmd_rx: tokio::sync::mpsc::UnboundedReceiver<HandlerCommand>,
114        raw_rx: tokio::sync::mpsc::UnboundedReceiver<Message>,
115        subscriptions: SubscriptionState,
116    ) -> Self {
117        Self {
118            signal,
119            inner: None,
120            cmd_rx,
121            raw_rx,
122            subscriptions,
123            message_queue: VecDeque::new(),
124            replay_request_id: -1,
125            needs_subscription_replay: false,
126            pending_subscription_requests: AHashMap::new(),
127        }
128    }
129
130    fn next_replay_request_id(&mut self) -> i64 {
131        self.replay_request_id -= 1;
132        self.replay_request_id
133    }
134
135    async fn replay_subscriptions(&mut self) {
136        self.pending_subscription_requests.clear();
137        let topics = self.subscriptions.reset_after_reconnect();
138        if topics.is_empty() {
139            log::debug!("No subscriptions to replay after reconnect");
140            return;
141        }
142
143        log::debug!("Replaying {} subscriptions after reconnect", topics.len());
144
145        for topic in topics {
146            // Topic format: "symbol:Level:trades:ticker" or "candles:symbol:Width"
147            if let Some(rest) = topic.strip_prefix("candles:") {
148                if let Some((symbol, width_str)) = rest.rsplit_once(':') {
149                    if let Some(width) = Self::parse_candle_width(width_str) {
150                        let request_id = self.next_replay_request_id();
151                        log::debug!(
152                            "Replaying candle subscription: symbol={symbol}, width={width:?}"
153                        );
154                        self.pending_subscription_requests.insert(
155                            request_id,
156                            PendingSubscriptionRequest::Subscribe(topic.clone()),
157                        );
158                        self.send_subscribe_candles(request_id, Ustr::from(symbol), width)
159                            .await;
160                    } else {
161                        log::warn!("Failed to parse candle width from topic: {topic}");
162                    }
163                } else {
164                    log::warn!("Invalid candle topic format: {topic}");
165                }
166            } else if let Some((symbol, spec)) = AxMdSubscriptionSpec::parse_topic(&topic) {
167                let request_id = self.next_replay_request_id();
168                log::debug!("Replaying market data subscription: symbol={symbol}, spec={spec:?}");
169                self.pending_subscription_requests
170                    .insert(request_id, PendingSubscriptionRequest::Subscribe(topic));
171                self.send_subscribe(request_id, symbol, spec).await;
172            } else {
173                log::warn!("Failed to parse market data subscription topic: {topic}");
174            }
175        }
176
177        log::debug!("Subscription replay completed");
178    }
179
180    fn parse_candle_width(s: &str) -> Option<AxCandleWidth> {
181        match s {
182            "Seconds1" => Some(AxCandleWidth::Seconds1),
183            "Seconds5" => Some(AxCandleWidth::Seconds5),
184            "Minutes1" => Some(AxCandleWidth::Minutes1),
185            "Minutes5" => Some(AxCandleWidth::Minutes5),
186            "Minutes15" => Some(AxCandleWidth::Minutes15),
187            "Hours1" => Some(AxCandleWidth::Hours1),
188            "Days1" => Some(AxCandleWidth::Days1),
189            _ => None,
190        }
191    }
192
193    /// Returns the next message from the handler.
194    ///
195    /// This method blocks until a message is available or the handler is stopped.
196    pub(crate) async fn next(&mut self) -> Option<AxDataWsMessage> {
197        loop {
198            if self.needs_subscription_replay && self.message_queue.is_empty() {
199                self.needs_subscription_replay = false;
200                self.replay_subscriptions().await;
201            }
202
203            if let Some(msg) = self.message_queue.pop_front() {
204                return Some(msg);
205            }
206
207            tokio::select! {
208                Some(cmd) = self.cmd_rx.recv() => {
209                    self.handle_command(cmd).await;
210                }
211
212                () = tokio::time::sleep(std::time::Duration::from_millis(100)) => {
213                    if self.signal.load(Ordering::Acquire) {
214                        log::debug!("Stop signal received during idle period");
215                        return None;
216                    }
217                }
218
219                msg = self.raw_rx.recv() => {
220                    let msg = match msg {
221                        Some(msg) => msg,
222                        None => {
223                            log::debug!("WebSocket stream closed");
224                            return None;
225                        }
226                    };
227
228                    if let Message::Ping(data) = &msg {
229                        log::trace!("Received ping frame with {} bytes", data.len());
230
231                        if let Some(client) = &self.inner
232                            && let Err(e) = client.send_pong(data.to_vec()).await
233                        {
234                            log::warn!("Failed to send pong frame: {e}");
235                        }
236                        continue;
237                    }
238
239                    if let Some(message) = self.parse_raw_message(msg) {
240                        self.message_queue.push_back(message);
241                    }
242
243                    if self.signal.load(Ordering::Acquire) {
244                        log::debug!("Stop signal received");
245                        return None;
246                    }
247                }
248            }
249        }
250    }
251
252    async fn handle_command(&mut self, cmd: HandlerCommand) {
253        match cmd {
254            HandlerCommand::SetClient(client) => {
255                log::debug!("WebSocketClient received by handler");
256                self.inner = Some(client);
257            }
258            HandlerCommand::Disconnect => {
259                log::debug!("Disconnect command received");
260
261                if let Some(inner) = self.inner.take() {
262                    inner.disconnect().await;
263                }
264            }
265            HandlerCommand::ReplaySubscriptions => {
266                log::debug!("ReplaySubscriptions command received");
267                self.replay_subscriptions().await;
268            }
269            HandlerCommand::Subscribe {
270                request_id,
271                symbol,
272                spec,
273            } => {
274                log::debug!(
275                    "Subscribe command received: request_id={request_id}, symbol={symbol}, spec={spec:?}"
276                );
277                let topic = spec.topic(symbol.as_str());
278                self.pending_subscription_requests
279                    .insert(request_id, PendingSubscriptionRequest::Subscribe(topic));
280                self.send_subscribe(request_id, symbol, spec).await;
281            }
282            HandlerCommand::Unsubscribe {
283                request_id,
284                symbol,
285                topic,
286            } => {
287                log::debug!(
288                    "Unsubscribe command received: request_id={request_id}, symbol={symbol}"
289                );
290                self.pending_subscription_requests
291                    .insert(request_id, PendingSubscriptionRequest::Unsubscribe(topic));
292                self.send_unsubscribe(request_id, symbol).await;
293            }
294            HandlerCommand::SubscribeCandles {
295                request_id,
296                symbol,
297                width,
298            } => {
299                log::debug!(
300                    "SubscribeCandles command received: request_id={request_id}, symbol={symbol}, width={width:?}"
301                );
302                let topic = format!("candles:{symbol}:{width:?}");
303                self.pending_subscription_requests
304                    .insert(request_id, PendingSubscriptionRequest::Subscribe(topic));
305                self.send_subscribe_candles(request_id, symbol, width).await;
306            }
307            HandlerCommand::UnsubscribeCandles {
308                request_id,
309                symbol,
310                width,
311                topic,
312            } => {
313                log::debug!(
314                    "UnsubscribeCandles command received: request_id={request_id}, symbol={symbol}, width={width:?}"
315                );
316                self.pending_subscription_requests
317                    .insert(request_id, PendingSubscriptionRequest::Unsubscribe(topic));
318                self.message_queue
319                    .push_back(AxDataWsMessage::CandleUnsubscribed { symbol, width });
320                self.send_unsubscribe_candles(request_id, symbol, width)
321                    .await;
322            }
323        }
324    }
325
326    async fn send_subscribe(&mut self, request_id: i64, symbol: Ustr, spec: AxMdSubscriptionSpec) {
327        let msg = AxMdSubscribe {
328            rid: request_id,
329            msg_type: AxMdRequestType::Subscribe,
330            symbol,
331            level: spec.level,
332            trades: spec.trades,
333            ticker: spec.ticker,
334        };
335
336        if let Err(e) = self.send_json(&msg).await {
337            self.pending_subscription_requests.remove(&request_id);
338            log::error!("Failed to send subscribe message: {e}");
339        }
340    }
341
342    async fn send_unsubscribe(&mut self, request_id: i64, symbol: Ustr) {
343        let msg = AxMdUnsubscribe {
344            rid: request_id,
345            msg_type: AxMdRequestType::Unsubscribe,
346            symbol,
347        };
348
349        if let Err(e) = self.send_json(&msg).await {
350            self.pending_subscription_requests.remove(&request_id);
351            log::error!("Failed to send unsubscribe message: {e}");
352        }
353    }
354
355    async fn send_subscribe_candles(
356        &mut self,
357        request_id: i64,
358        symbol: Ustr,
359        width: AxCandleWidth,
360    ) {
361        let msg = AxMdSubscribeCandles {
362            rid: request_id,
363            msg_type: AxMdRequestType::SubscribeCandles,
364            symbol,
365            width,
366        };
367
368        if let Err(e) = self.send_json(&msg).await {
369            self.pending_subscription_requests.remove(&request_id);
370            log::error!("Failed to send subscribe_candles message: {e}");
371        }
372    }
373
374    async fn send_unsubscribe_candles(
375        &mut self,
376        request_id: i64,
377        symbol: Ustr,
378        width: AxCandleWidth,
379    ) {
380        let msg = AxMdUnsubscribeCandles {
381            rid: request_id,
382            msg_type: AxMdRequestType::UnsubscribeCandles,
383            symbol,
384            width,
385        };
386
387        if let Err(e) = self.send_json(&msg).await {
388            self.pending_subscription_requests.remove(&request_id);
389            log::error!("Failed to send unsubscribe_candles message: {e}");
390        }
391    }
392
393    async fn send_json<T: serde::Serialize>(&self, msg: &T) -> Result<(), String> {
394        let Some(inner) = &self.inner else {
395            return Err("No WebSocket client available".to_string());
396        };
397
398        let payload = serde_json::to_string(msg).map_err(|e| e.to_string())?;
399        log::trace!("Sending WebSocket payload ({} bytes)", payload.len());
400
401        inner
402            .send_text(payload, None)
403            .await
404            .map_err(|e| e.to_string())
405    }
406
407    fn parse_raw_message(&mut self, msg: Message) -> Option<AxDataWsMessage> {
408        match msg {
409            Message::Text(text) => {
410                if text == nautilus_network::RECONNECTED {
411                    log::info!("Received WebSocket reconnected signal");
412                    self.needs_subscription_replay = true;
413                    return Some(AxDataWsMessage::Reconnected);
414                }
415
416                log::trace!("Raw websocket message: {text}");
417
418                match parse_md_message(&text) {
419                    Ok(message) => self.handle_message(message),
420                    Err(e) => {
421                        log::error!("Failed to parse WebSocket message: {e}: {text}");
422                        None
423                    }
424                }
425            }
426            Message::Binary(data) => {
427                log::debug!("Received binary message with {} bytes", data.len());
428                None
429            }
430            Message::Close(_) => {
431                log::debug!("Received close message, waiting for reconnection");
432                None
433            }
434            _ => None,
435        }
436    }
437
438    fn handle_message(&mut self, message: AxMdMessage) -> Option<AxDataWsMessage> {
439        match &message {
440            AxMdMessage::Error(error) => {
441                let is_benign = error.message.contains("already subscribed")
442                    || error.message.contains("not subscribed");
443
444                if let Some(rid) = error.request_id
445                    && let Some(request) = self.pending_subscription_requests.remove(&rid)
446                {
447                    match request {
448                        PendingSubscriptionRequest::Subscribe(topic) => {
449                            self.subscriptions.mark_failure(&topic);
450                        }
451                        PendingSubscriptionRequest::Unsubscribe(topic) => {
452                            self.subscriptions.confirm_unsubscribe(&topic);
453                        }
454                    }
455                }
456
457                if is_benign {
458                    log::warn!("Subscription state: {}", error.message);
459                } else {
460                    log::error!("Received error from exchange: {}", error.message);
461                }
462            }
463            AxMdMessage::SubscriptionResponse(response) => {
464                let is_subscribe = response.result.subscribed.is_some()
465                    || response.result.subscribed_candle.is_some();
466                let is_unsubscribe = response.result.unsubscribed.is_some()
467                    || response.result.unsubscribed_candle.is_some();
468
469                if let Some(request) = self.pending_subscription_requests.remove(&response.rid) {
470                    match request {
471                        PendingSubscriptionRequest::Subscribe(topic) if is_subscribe => {
472                            self.subscriptions.confirm_subscribe(&topic);
473                        }
474                        PendingSubscriptionRequest::Unsubscribe(topic) if is_unsubscribe => {
475                            self.subscriptions.confirm_unsubscribe(&topic);
476                        }
477                        request => {
478                            log::warn!(
479                                "Unexpected subscription response for request: {request:?}, \
480                                 response={response:?}"
481                            );
482                        }
483                    }
484                }
485
486                if let Some(symbol) = &response.result.subscribed {
487                    log::debug!("Subscription confirmed for symbol: {symbol}");
488                } else if let Some(candle) = &response.result.subscribed_candle {
489                    log::debug!("Candle subscription confirmed: {candle}");
490                } else if let Some(symbol) = &response.result.unsubscribed {
491                    log::debug!("Unsubscription confirmed for symbol: {symbol}");
492                } else if let Some(candle) = &response.result.unsubscribed_candle {
493                    log::debug!("Candle unsubscription confirmed: {candle}");
494                }
495                return None;
496            }
497            _ => {}
498        }
499
500        Some(AxDataWsMessage::MdMessage(message))
501    }
502}
503
504#[derive(Debug)]
505enum PendingSubscriptionRequest {
506    Subscribe(String),
507    Unsubscribe(String),
508}
509
510#[cfg(test)]
511mod tests {
512    use nautilus_network::websocket::SubscriptionState;
513    use rstest::rstest;
514
515    use super::*;
516    use crate::websocket::messages::{AxMdSubscriptionResponse, AxMdSubscriptionResult, AxWsError};
517
518    const TOPIC: &str = "EURUSD-PERP:Level2:false:false";
519
520    #[rstest]
521    fn test_subscription_response_confirms_subscribe() {
522        let subscriptions = SubscriptionState::new(':');
523        subscriptions.mark_subscribe(TOPIC);
524        let mut handler = create_handler(subscriptions.clone());
525        handler
526            .pending_subscription_requests
527            .insert(1, PendingSubscriptionRequest::Subscribe(TOPIC.to_string()));
528
529        handler.handle_message(AxMdMessage::SubscriptionResponse(
530            AxMdSubscriptionResponse {
531                rid: 1,
532                result: AxMdSubscriptionResult {
533                    subscribed: Some("EURUSD-PERP".to_string()),
534                    subscribed_candle: None,
535                    unsubscribed: None,
536                    unsubscribed_candle: None,
537                },
538            },
539        ));
540
541        assert_eq!(subscriptions.len(), 1);
542        assert!(subscriptions.pending_subscribe_topics().is_empty());
543        assert!(subscriptions.pending_unsubscribe_topics().is_empty());
544    }
545
546    #[rstest]
547    fn test_subscription_error_keeps_topic_pending_for_replay() {
548        let subscriptions = SubscriptionState::new(':');
549        subscriptions.mark_subscribe(TOPIC);
550        subscriptions.confirm_subscribe(TOPIC);
551        let mut handler = create_handler(subscriptions.clone());
552        handler
553            .pending_subscription_requests
554            .insert(2, PendingSubscriptionRequest::Subscribe(TOPIC.to_string()));
555
556        handler.handle_message(AxMdMessage::Error(AxWsError {
557            code: Some("400".to_string()),
558            message: "subscription failed".to_string(),
559            request_id: Some(2),
560        }));
561
562        assert_eq!(subscriptions.len(), 0);
563        assert_eq!(
564            subscriptions.pending_subscribe_topics(),
565            vec![TOPIC.to_string()]
566        );
567        assert!(subscriptions.pending_unsubscribe_topics().is_empty());
568    }
569
570    #[rstest]
571    fn test_already_subscribed_keeps_topic_pending_for_replay() {
572        let subscriptions = SubscriptionState::new(':');
573        subscriptions.mark_subscribe(TOPIC);
574        let mut handler = create_handler(subscriptions.clone());
575        handler
576            .pending_subscription_requests
577            .insert(3, PendingSubscriptionRequest::Subscribe(TOPIC.to_string()));
578
579        handler.handle_message(AxMdMessage::Error(AxWsError {
580            code: Some("400".to_string()),
581            message: "already subscribed".to_string(),
582            request_id: Some(3),
583        }));
584
585        assert_eq!(subscriptions.len(), 0);
586        assert_eq!(
587            subscriptions.pending_subscribe_topics(),
588            vec![TOPIC.to_string()]
589        );
590        assert!(subscriptions.pending_unsubscribe_topics().is_empty());
591    }
592
593    #[rstest]
594    #[case("not subscribed")]
595    #[case("subscription failed")]
596    fn test_unsubscribe_error_confirms_removal(#[case] message: &str) {
597        let subscriptions = SubscriptionState::new(':');
598        subscriptions.mark_subscribe(TOPIC);
599        subscriptions.confirm_subscribe(TOPIC);
600        subscriptions.mark_unsubscribe(TOPIC);
601        let mut handler = create_handler(subscriptions.clone());
602        handler.pending_subscription_requests.insert(
603            4,
604            PendingSubscriptionRequest::Unsubscribe(TOPIC.to_string()),
605        );
606
607        handler.handle_message(AxMdMessage::Error(AxWsError {
608            code: Some("400".to_string()),
609            message: message.to_string(),
610            request_id: Some(4),
611        }));
612
613        assert_eq!(subscriptions.len(), 0);
614        assert!(subscriptions.pending_subscribe_topics().is_empty());
615        assert!(subscriptions.pending_unsubscribe_topics().is_empty());
616    }
617
618    #[rstest]
619    fn test_subscription_response_confirms_unsubscribe() {
620        let subscriptions = SubscriptionState::new(':');
621        subscriptions.mark_subscribe(TOPIC);
622        subscriptions.confirm_subscribe(TOPIC);
623        subscriptions.mark_unsubscribe(TOPIC);
624        let mut handler = create_handler(subscriptions.clone());
625        handler.pending_subscription_requests.insert(
626            5,
627            PendingSubscriptionRequest::Unsubscribe(TOPIC.to_string()),
628        );
629
630        handler.handle_message(AxMdMessage::SubscriptionResponse(
631            AxMdSubscriptionResponse {
632                rid: 5,
633                result: AxMdSubscriptionResult {
634                    subscribed: None,
635                    subscribed_candle: None,
636                    unsubscribed: Some("EURUSD-PERP".to_string()),
637                    unsubscribed_candle: None,
638                },
639            },
640        ));
641
642        assert_eq!(subscriptions.len(), 0);
643        assert!(subscriptions.pending_subscribe_topics().is_empty());
644        assert!(subscriptions.pending_unsubscribe_topics().is_empty());
645    }
646
647    #[rstest]
648    #[tokio::test]
649    async fn test_replay_clears_stale_requests_and_pending_unsubscribes() {
650        let subscriptions = SubscriptionState::new(':');
651        subscriptions.mark_subscribe(TOPIC);
652        subscriptions.confirm_subscribe(TOPIC);
653        subscriptions.mark_unsubscribe(TOPIC);
654        let mut handler = create_handler(subscriptions.clone());
655        handler.pending_subscription_requests.insert(
656            6,
657            PendingSubscriptionRequest::Unsubscribe(TOPIC.to_string()),
658        );
659
660        handler.replay_subscriptions().await;
661
662        assert_eq!(subscriptions.len(), 0);
663        assert!(subscriptions.all_topics().is_empty());
664        assert!(subscriptions.pending_subscribe_topics().is_empty());
665        assert!(subscriptions.pending_unsubscribe_topics().is_empty());
666        assert!(handler.pending_subscription_requests.is_empty());
667    }
668
669    fn create_handler(subscriptions: SubscriptionState) -> AxMdWsFeedHandler {
670        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
671        let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
672
673        AxMdWsFeedHandler::new(
674            Arc::new(AtomicBool::new(false)),
675            cmd_rx,
676            raw_rx,
677            subscriptions,
678        )
679    }
680}
nautilus_architect_ax/websocket/data/handler.rs

nautilus_architect_ax/websocket/data/
handler.rs
