rusty_common/websocket/

client.rs

//! WebSocket client implementation
//!
//! Provides a high-performance WebSocket client using yawc.

use futures_util::{SinkExt, StreamExt};
use parking_lot::RwLock;
use std::sync::Arc;
use tokio::sync::{mpsc, watch};
use tokio::time::timeout;
use url::Url;
use yawc::{DeflateOptions, Options, WebSocket};

use super::{
    Message, MessageHandler, ReconnectStrategy, WebSocketConfig, WebSocketError, WebSocketResult,
};

/// WebSocket connection state
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ConnectionState {
    /// Not connected
    Disconnected,

    /// Connecting
    Connecting,

    /// Connected
    Connected,

    /// Reconnecting
    Reconnecting,

    /// Error state
    Error,

    /// Closed
    Closed,
}

/// WebSocket client statistics
#[derive(Debug, Default, Clone)]
pub struct ClientStats {
    /// Messages sent
    pub messages_sent: u64,

    /// Messages received
    pub messages_received: u64,

    /// Bytes sent
    pub bytes_sent: u64,

    /// Bytes received
    pub bytes_received: u64,

    /// Connection attempts
    pub connection_attempts: u32,

    /// Successful connections
    pub successful_connections: u32,

    /// Reconnection count
    pub reconnections: u32,

    /// Last ping timestamp (nanoseconds)
    pub last_ping_time: u64,

    /// Last pong timestamp (nanoseconds)
    pub last_pong_time: u64,
}

/// WebSocket stream type
pub type WebSocketStream = WebSocket;

/// WebSocket client
pub struct WebSocketClient {
    /// Configuration
    config: WebSocketConfig,

    /// Current connection state
    state: Arc<RwLock<ConnectionState>>,

    /// Statistics
    stats: Arc<RwLock<ClientStats>>,

    /// Reconnection strategy
    reconnect: ReconnectStrategy,

    /// Shutdown signal
    shutdown_tx: watch::Sender<bool>,
    shutdown_rx: watch::Receiver<bool>,

    /// Send channel
    send_tx: Option<mpsc::UnboundedSender<Message>>,
}

impl WebSocketClient {
    /// Create yawc Options from our WebSocketConfig
    fn create_yawc_options(&self) -> Options {
        let mut options = Options::default();

        // Configure compression
        if self.config.compression.enabled {
            // Set compression option (yawc currently only supports enabling/disabling)
            // Window bits configuration is not exposed in the current yawc API
            options.compression = Some(DeflateOptions::default());
        }

        // Configure size limits
        options.max_payload_read = Some(self.config.max_message_size);
        options.max_read_buffer = Some(self.config.max_message_size * 2); // Allow buffer for fragmented messages

        // Note: yawc 0.2 doesn't expose ping_interval and pong_timeout in Options.
        // Ping/pong is handled internally by yawc at the protocol level.
        // For exchanges that require custom application-level ping messages,
        // use the custom_ping_message field in WebSocketConfig.

        options
    }

    /// Create a new WebSocket client
    #[must_use]
    pub fn new(config: WebSocketConfig) -> Self {
        let (shutdown_tx, shutdown_rx) = watch::channel(false);

        Self {
            reconnect: ReconnectStrategy::new(config.reconnect.clone()),
            config,
            state: Arc::new(RwLock::new(ConnectionState::Disconnected)),
            stats: Arc::new(RwLock::new(ClientStats::default())),
            shutdown_tx,
            shutdown_rx,
            send_tx: None,
        }
    }

    /// Connect and run the client with a message handler
    pub async fn run<H: MessageHandler + 'static>(
        &mut self,
        mut handler: H,
    ) -> WebSocketResult<()> {
        let mut shutdown_rx = self.shutdown_rx.clone();

        loop {
            // Check for shutdown
            if *shutdown_rx.borrow() {
                break;
            }

            // Update state
            *self.state.write() = ConnectionState::Connecting;
            self.stats.write().connection_attempts += 1;

            // Try to connect
            match self.connect().await {
                Ok(ws) => {
                    // Update state
                    *self.state.write() = ConnectionState::Connected;
                    self.stats.write().successful_connections += 1;
                    self.reconnect.reset();

                    // Create a new channel for this connection (high-performance approach)
                    let (conn_send_tx, conn_send_rx) = mpsc::unbounded_channel();

                    // Store the sender in the client
                    self.send_tx = Some(conn_send_tx.clone());

                    // Set sender for handler so it can send messages
                    // Note: This clone is necessary as both the client and handler may need
                    // to send messages independently
                    handler.set_sender(conn_send_tx);

                    // Notify handler (after setting sender so it can send subscription messages)
                    if let Err(e) = handler.on_connected().await {
                        log::error!("Handler error on connected: {e}");
                    }

                    // Run the connection
                    if let Err(e) = self.run_connection(ws, &mut handler, conn_send_rx).await {
                        log::error!("Connection error: {e}");

                        // Notify handler
                        if let Err(e) = handler.on_error(e).await {
                            log::error!("Handler error on error: {e}");
                        }
                    }

                    // Clear send channel
                    self.send_tx = None;

                    // Notify handler
                    if let Err(e) = handler.on_disconnected().await {
                        log::error!("Handler error on disconnected: {e}");
                    }
                }
                Err(e) => {
                    log::error!("Failed to connect: {e}");
                    *self.state.write() = ConnectionState::Error;
                    self.send_tx = None;
                }
            }

            // Check if we should reconnect
            if !self.reconnect.should_reconnect() {
                break;
            }

            // Get reconnection delay
            if let Some(delay) = self.reconnect.next_delay() {
                *self.state.write() = ConnectionState::Reconnecting;
                self.stats.write().reconnections += 1;

                log::info!(
                    "Reconnecting in {:?} (attempt {})",
                    delay,
                    self.reconnect.attempts()
                );

                // Wait with cancellation
                tokio::select! {
                    _ = tokio::time::sleep(delay) => {}
                    _ = shutdown_rx.changed() => {
                        if *shutdown_rx.borrow() {
                            break;
                        }
                    }
                }
            } else {
                break;
            }
        }

        *self.state.write() = ConnectionState::Closed;
        Ok(())
    }

    /// Connect to the WebSocket server
    async fn connect(&self) -> WebSocketResult<WebSocketStream> {
        // Parse URL
        let url = Url::parse(&self.config.url)?;

        // Create yawc Options from our config
        let options = self.create_yawc_options();

        // Connect with timeout and options
        let ws = timeout(
            self.config.connect_timeout,
            WebSocket::connect(url).with_options(options),
        )
        .await
        .map_err(|_| WebSocketError::Timeout(self.config.connect_timeout.as_millis() as u64))?
        .map_err(WebSocketError::from)?;

        Ok(ws)
    }

    /// Run the WebSocket connection
    async fn run_connection<H: MessageHandler>(
        &self,
        mut ws: WebSocketStream,
        handler: &mut H,
        mut send_rx: mpsc::UnboundedReceiver<Message>,
    ) -> WebSocketResult<()> {
        let mut shutdown_rx = self.shutdown_rx.clone();

        loop {
            tokio::select! {
                // Shutdown signal
                _ = shutdown_rx.changed() => {
                    if *shutdown_rx.borrow() {
                        // Send close frame
                        ws.send(yawc::frame::FrameView::close(
                            yawc::close::CloseCode::Iana(1000),
                            "Client shutdown"
                        )).await?;
                        break;
                    }
                }

                // Send message (high-performance unbounded channel)
                Some(message) = send_rx.recv() => {
                    // Track bytes for stats
                    let bytes_sent = match &message {
                        Message::Text(text) => text.len() as u64,
                        Message::Binary(data) => data.len() as u64,
                        _ => 0,
                    };

                    // Convert to owned data for yawc's Bytes requirement
                    match message {
                        Message::Text(text) => {
                            // Convert SmartString to Vec<u8> then to Bytes
                            let data = text.as_bytes().to_vec();
                            ws.send(yawc::frame::FrameView::text(data)).await?;
                        }
                        Message::Binary(data) => {
                            // Pass owned Vec<u8> directly, it will convert to Bytes
                            ws.send(yawc::frame::FrameView::binary(data)).await?;
                        }
                        Message::Close(close_frame) => {
                            match close_frame {
                                Some((code, reason)) => {
                                    ws.send(yawc::frame::FrameView::close(
                                        yawc::close::CloseCode::Iana(code),
                                        &reason
                                    )).await?;
                                }
                                None => {
                                    ws.send(yawc::frame::FrameView::close(
                                        yawc::close::CloseCode::Iana(1000),
                                        "Normal closure"
                                    )).await?;
                                }
                            }
                        }
                        Message::Ping(data) => {
                            ws.send(yawc::frame::FrameView::ping(data)).await?;
                        }
                        Message::Pong(data) => {
                            ws.send(yawc::frame::FrameView::pong(data)).await?;
                        }
                        Message::Frame(frame) => {
                            ws.send(frame).await?;
                        }
                    }

                    // Update stats
                    let mut stats = self.stats.write();
                    stats.messages_sent += 1;
                    stats.bytes_sent += bytes_sent;
                }

                // Note: Ping/pong is handled automatically by yawc when Options are configured

                // Receive message
                frame = timeout(self.config.timeout, ws.next()) => {
                    match frame {
                        Ok(Some(frame)) => {
                            let message = Message::from_frame_view(frame);

                            // Update stats
                            {
                                let mut stats = self.stats.write();
                                stats.messages_received += 1;
                                match &message {
                                    Message::Text(text) => stats.bytes_received += text.len() as u64,
                                    Message::Binary(data) => stats.bytes_received += data.len() as u64,
                                    _ => {}
                                }
                            } // stats guard is dropped here

                            // Handle message
                            handler.on_message(message).await?;
                        }
                        Ok(None) => {
                            // Connection closed
                            break;
                        }
                        Err(_) => {
                            // Timeout
                            return Err(WebSocketError::Timeout(self.config.timeout.as_millis() as u64));
                        }
                    }
                }
            }
        }

        Ok(())
    }

    /// Send a message (high-performance, non-blocking)
    pub async fn send(&self, message: Message) -> WebSocketResult<()> {
        match &self.send_tx {
            Some(tx) => {
                // Use try_send for non-blocking operation
                tx.send(message).map_err(|_| WebSocketError::NotConnected)?;
                Ok(())
            }
            None => Err(WebSocketError::NotConnected),
        }
    }

    /// Try to send a message (non-blocking, returns false if channel is full)
    pub fn try_send(&self, message: Message) -> bool {
        match &self.send_tx {
            Some(tx) => {
                // Since we use unbounded channel, this should always succeed
                // unless the receiver is dropped (disconnected)
                tx.send(message).is_ok()
            }
            None => false,
        }
    }

    /// Get the current connection state
    pub fn state(&self) -> ConnectionState {
        *self.state.read()
    }

    /// Get statistics
    pub fn stats(&self) -> ClientStats {
        (*self.stats.read()).clone()
    }

    /// Shutdown the client
    pub fn shutdown(&self) {
        let _ = self.shutdown_tx.send(true);
    }
}

/// WebSocket client builder
pub struct WebSocketClientBuilder {
    config: WebSocketConfig,
}

impl WebSocketClientBuilder {
    /// Create a new builder
    #[must_use]
    pub const fn new(config: WebSocketConfig) -> Self {
        Self { config }
    }

    /// Build the client
    pub fn build(self) -> WebSocketClient {
        WebSocketClient::new(self.config)
    }
}