rusty_feeder/exchange/binance/spot/

message_handler.rs

//! Binance Spot WebSocket message handler implementation

use crate::exchange::zerocopy_helpers::{deserialize_from_slice_borrowed, deserialize_from_vec};
use async_trait::async_trait;
use parking_lot::RwLock;
use quanta::Clock;
use rusty_common::json::{self, Value};
use rusty_common::websocket::{Message, MessageHandler, WebSocketError, WebSocketResult};
use simd_json::prelude::ValueObjectAccess;
use smartstring::alias::String;
use std::sync::Arc;
use tokio::sync::mpsc;

use super::data::{orderbook::OrderbookMessage, trade::TradeMessage};
use crate::exchange::binance::common::messages::BinanceWebsocketResponse;
use crate::provider::prelude::ConnectionStats;

/// Message handler for Binance Spot WebSocket connections
pub struct BinanceSpotMessageHandler {
    /// Clock for timing
    clock: Clock,

    /// Connection statistics
    stats: Arc<RwLock<ConnectionStats>>,

    /// Trade message sender
    trade_tx: Option<mpsc::Sender<TradeMessage>>,

    /// Orderbook message sender
    orderbook_tx: Option<mpsc::Sender<OrderbookMessage>>,

    /// Whether this is a combined stream
    is_combined: bool,

    /// Subscription message to send on connection
    subscription_message: Option<String>,
}

impl BinanceSpotMessageHandler {
    /// Create a new message handler for trade messages
    pub const fn new_trade_handler(
        clock: Clock,
        stats: Arc<RwLock<ConnectionStats>>,
        trade_tx: mpsc::Sender<TradeMessage>,
        is_combined: bool,
        subscription_message: Option<String>,
    ) -> Self {
        Self {
            clock,
            stats,
            trade_tx: Some(trade_tx),
            orderbook_tx: None,
            is_combined,
            subscription_message,
        }
    }

    /// Create a new message handler for orderbook messages
    pub const fn new_orderbook_handler(
        clock: Clock,
        stats: Arc<RwLock<ConnectionStats>>,
        orderbook_tx: mpsc::Sender<OrderbookMessage>,
        is_combined: bool,
        subscription_message: Option<String>,
    ) -> Self {
        Self {
            clock,
            stats,
            trade_tx: None,
            orderbook_tx: Some(orderbook_tx),
            is_combined,
            subscription_message,
        }
    }

    /// Update connection statistics
    #[inline]
    fn update_receive_stats(&self, message_size: usize) {
        let local_time = self.clock.raw();
        let mut s = self.stats.write();
        s.messages_received += 1;
        s.bytes_received += message_size as u64;
        s.last_message_time = local_time;
    }

    /// Process trade text message
    #[inline]
    async fn process_trade_text_message(&self, text: String) -> WebSocketResult<()> {
        let trade_tx = self
            .trade_tx
            .as_ref()
            .ok_or(WebSocketError::ConnectionClosed(
                "No trade sender configured".to_string(),
            ))?;

        // Zero-copy parsing: convert text to bytes and parse directly
        match deserialize_from_slice_borrowed::<BinanceWebsocketResponse<TradeMessage>>(
            text.as_bytes(),
        ) {
            Ok(BinanceWebsocketResponse::Data(trade_msg)) => {
                self.update_receive_stats(text.len());
                if let Err(e) = trade_tx.send(trade_msg).await {
                    log::error!("Failed to send TradeMessage: {e}");
                }
            }
            Ok(BinanceWebsocketResponse::Confirmation(conf_msg)) => {
                log::debug!("Received subscription confirmation: {conf_msg:?}");
            }
            Err(e_direct) => {
                if !self.is_combined {
                    log::error!("Failed to parse TradeMessage text: {e_direct}");
                    return Ok(());
                }

                let val = match deserialize_from_slice_borrowed::<Value>(text.as_bytes()) {
                    Ok(val) => val,
                    Err(e_combined) => {
                        log::error!(
                            "Failed to parse combined stream text message: {e_combined}, original error: {e_direct}"
                        );
                        return Ok(());
                    }
                };

                let Some(data) = val.get("data") else {
                    log::error!("Failed to parse trade text message: {e_direct}");
                    return Ok(());
                };

                let trade_msg =
                    match simd_json::serde::from_owned_value::<TradeMessage>(data.to_owned()) {
                        Ok(trade_msg) => trade_msg,
                        Err(e) => {
                            log::error!("Failed to parse trade text message: {e}");
                            return Ok(());
                        }
                    };

                self.update_receive_stats(text.len());

                if let Err(e) = trade_tx.try_send(trade_msg) {
                    log::error!("Failed to send trade message: {e}");
                }
            }
        }

        Ok(())
    }

    /// Process trade binary message
    #[inline]
    async fn process_trade_binary_message(&self, bin: Vec<u8>) -> WebSocketResult<()> {
        let trade_tx = self
            .trade_tx
            .as_ref()
            .ok_or(WebSocketError::ConnectionClosed(
                "No trade sender configured".to_string(),
            ))?;

        let bin_len = bin.len();

        // For combined streams, we need to clone for fallback parsing
        // For direct streams, we can use zero-copy parsing
        if self.is_combined {
            // Clone only for combined streams that might need fallback
            match deserialize_from_vec::<TradeMessage>(bin.clone()) {
                Ok(trade_msg) => {
                    self.update_receive_stats(bin_len);
                    if let Err(e) = trade_tx.try_send(trade_msg) {
                        log::error!("Failed to send trade message: {e}");
                    }
                }
                Err(e_direct) => {
                    // Fallback to combined stream parsing
                    let val = match deserialize_from_vec::<Value>(bin) {
                        Ok(v) => v,
                        Err(e_value) => {
                            log::error!(
                                "Failed to parse combined stream binary message as Value: {e_value}, original error: {e_direct}"
                            );
                            return Ok(());
                        }
                    };

                    let Some(data) = val.get("data") else {
                        log::error!(
                            "Combined stream binary message is missing 'data' field, original error trying to parse as TradeMessage: {e_direct}"
                        );
                        return Ok(());
                    };

                    let trade_msg = match simd_json::serde::from_owned_value::<TradeMessage>(
                        data.to_owned(),
                    ) {
                        Ok(msg) => msg,
                        Err(e_data) => {
                            log::error!(
                                "Failed to parse 'data' field from combined stream binary message into TradeMessage: {e_data}, original error trying to parse root as TradeMessage: {e_direct}"
                            );
                            return Ok(());
                        }
                    };

                    self.update_receive_stats(bin_len);
                    if let Err(e_send) = trade_tx.try_send(trade_msg) {
                        log::error!(
                            "Failed to send trade message from combined binary stream: {e_send}"
                        );
                    }
                }
            }
        } else {
            // Zero-copy parsing for direct streams
            match deserialize_from_vec::<TradeMessage>(bin) {
                Ok(trade_msg) => {
                    self.update_receive_stats(bin_len);
                    if let Err(e) = trade_tx.try_send(trade_msg) {
                        log::error!("Failed to send trade message: {e}");
                    }
                }
                Err(e_direct) => {
                    log::error!("Failed to parse trade binary message: {e_direct}");
                    return Ok(());
                }
            }
        }

        Ok(())
    }

    /// Process orderbook text message
    #[inline]
    async fn process_orderbook_text_message(&self, text: String) -> WebSocketResult<()> {
        let orderbook_tx = self
            .orderbook_tx
            .as_ref()
            .ok_or(WebSocketError::ConnectionClosed(
                "No orderbook sender configured".to_string(),
            ))?;

        match json::parse_fast_owned::<BinanceWebsocketResponse<OrderbookMessage>>(text.clone()) {
            Ok(BinanceWebsocketResponse::Data(orderbook_msg)) => {
                self.update_receive_stats(text.len());
                if let Err(e) = orderbook_tx.send(orderbook_msg).await {
                    log::error!("Failed to send OrderbookMessage: {e}");
                }
            }
            Ok(BinanceWebsocketResponse::Confirmation(conf_msg)) => {
                log::debug!("Received subscription confirmation: {conf_msg:?}");
            }
            Err(e_direct) => {
                if !self.is_combined {
                    log::error!("Failed to parse OrderbookMessage text: {e_direct}");
                    return Ok(());
                }

                let val = match json::parse::<Value>(text.as_str()) {
                    Ok(v) => v,
                    Err(e_value) => {
                        log::error!(
                            "Failed to parse combined stream text message for orderbook as Value: {e_value}, original error: {e_direct}"
                        );
                        return Ok(());
                    }
                };

                let Some(data) = val.get("data") else {
                    log::error!(
                        "Combined stream text message for orderbook is missing 'data' field, original error trying to parse as OrderbookMessage: {e_direct}"
                    );
                    return Ok(());
                };

                let ob_msg = match simd_json::serde::from_owned_value::<OrderbookMessage>(
                    data.to_owned(),
                ) {
                    Ok(msg) => msg,
                    Err(e_data) => {
                        log::error!(
                            "Failed to parse 'data' field from combined stream text message for orderbook into OrderbookMessage: {e_data}, original error: {e_direct}"
                        );
                        return Ok(());
                    }
                };

                self.update_receive_stats(text.len());
                if let Err(e_send) = orderbook_tx.send(ob_msg).await {
                    log::error!(
                        "Failed to send orderbook message from combined text stream: {e_send}"
                    );
                }
            }
        }

        Ok(())
    }

    /// Process orderbook binary message
    #[inline]
    async fn process_orderbook_binary_message(&self, bin: Vec<u8>) -> WebSocketResult<()> {
        let orderbook_tx = self
            .orderbook_tx
            .as_ref()
            .ok_or(WebSocketError::ConnectionClosed(
                "No orderbook sender configured".to_string(),
            ))?;

        let bin_len = bin.len();

        // For combined streams, we need to clone for fallback parsing
        // For direct streams, we can use zero-copy parsing
        if self.is_combined {
            // Clone only for combined streams that might need fallback
            match deserialize_from_vec::<OrderbookMessage>(bin.clone()) {
                Ok(ob_msg) => {
                    self.update_receive_stats(bin_len);
                    if let Err(e) = orderbook_tx.send(ob_msg).await {
                        log::error!("Failed to send orderbook message: {e}");
                    }
                }
                Err(e_direct) => {
                    // Fallback to combined stream parsing
                    let val = match deserialize_from_vec::<Value>(bin) {
                        Ok(v) => v,
                        Err(e_value) => {
                            log::error!(
                                "Failed to parse combined stream binary message for orderbook as Value: {e_value}, original error: {e_direct}"
                            );
                            return Ok(());
                        }
                    };

                    let Some(data) = val.get("data") else {
                        log::error!(
                            "Combined stream binary message for orderbook is missing 'data' field, original error trying to parse as OrderbookMessage: {e_direct}"
                        );
                        return Ok(());
                    };

                    let ob_msg = match simd_json::serde::from_owned_value::<OrderbookMessage>(
                        data.to_owned(),
                    ) {
                        Ok(msg) => msg,
                        Err(e_data) => {
                            log::error!(
                                "Failed to parse 'data' field from combined stream binary message for orderbook into OrderbookMessage: {e_data}, original error: {e_direct}"
                            );
                            return Ok(());
                        }
                    };

                    self.update_receive_stats(bin_len);
                    if let Err(e_send) = orderbook_tx.send(ob_msg).await {
                        log::error!(
                            "Failed to send orderbook message from combined binary stream: {e_send}"
                        );
                    }
                }
            }
        } else {
            // Zero-copy parsing for direct streams
            match deserialize_from_vec::<OrderbookMessage>(bin) {
                Ok(ob_msg) => {
                    self.update_receive_stats(bin_len);
                    if let Err(e) = orderbook_tx.send(ob_msg).await {
                        log::error!("Failed to send orderbook message: {e}");
                    }
                }
                Err(e_direct) => {
                    log::error!("Failed to parse orderbook binary message: {e_direct}");
                    return Ok(());
                }
            }
        }

        Ok(())
    }
}

#[async_trait]
impl MessageHandler for BinanceSpotMessageHandler {
    async fn on_connected(&mut self) -> WebSocketResult<()> {
        log::info!("Binance Spot WebSocket connected");
        self.stats.write().connected_time = self.clock.raw();

        // Subscription will be sent by the provider after connection
        Ok(())
    }

    async fn on_message(&mut self, message: Message) -> WebSocketResult<()> {
        match message {
            Message::Text(text) => {
                if self.trade_tx.is_some() {
                    self.process_trade_text_message(text).await?;
                } else if self.orderbook_tx.is_some() {
                    self.process_orderbook_text_message(text).await?;
                }
            }
            Message::Binary(bin) => {
                if self.trade_tx.is_some() {
                    self.process_trade_binary_message(bin).await?;
                } else if self.orderbook_tx.is_some() {
                    self.process_orderbook_binary_message(bin).await?;
                }
            }
            Message::Pong(_) => {
                // Update pong time
                self.stats.write().last_pong_time = self.clock.raw();
            }
            Message::Close(close_frame) => {
                log::info!("WebSocket connection closed: {close_frame:?}");
            }
            Message::Ping(payload) => {
                // Ping responses are handled automatically by WebSocketClient
                log::trace!("Received ping: {} bytes", payload.len());
            }
            Message::Frame(_) => {
                // Ignoring Frame messages
                log::trace!("Received unhandled Message::Frame");
            }
        }

        Ok(())
    }

    async fn on_disconnected(&mut self) -> WebSocketResult<()> {
        log::info!("Binance Spot WebSocket disconnected");
        Ok(())
    }
}