asterdex_sdk/ws/

client.rs

// US-010: WebSocketClient — manages WS connections with automatic reconnect
//
// Public API:
//   WebSocketClient::new(base_url) -> Result<Self, AsterDexError>
//   WebSocketClient::with_reconnect_config(self, cfg) -> Self
//   WebSocketClient::subscribe(&self, streams) -> Result<WebSocketStream, AsterDexError>
//   WebSocketClient::unsubscribe(stream, streams) -> Result<(), AsterDexError>
//   WebSocketClient::subscribe_user_data(&self, listen_key) -> Result<UserDataStream, AsterDexError>

use futures_util::{SinkExt, StreamExt};
use std::time::Duration;
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
use tokio_tungstenite::{connect_async, tungstenite::Message};
use url::Url;

use crate::rest::error::AsterDexError;
use crate::ws::reconnect::{ReconnectState, DEFAULT_INITIAL_BACKOFF, DEFAULT_MAX_BACKOFF};

/// Maximum number of streams allowed per single WebSocket connection (BR-006).
const MAX_STREAMS_PER_CONNECTION: usize = 200;

/// Channel buffer size. Large enough to absorb message bursts without backpressure
/// stalling the receive loop (which would delay pong responses and miss frames).
const CHANNEL_BUFFER: usize = 8_192;

/// Reconnect if no frame is received within this window.
/// Detects TCP half-open / silent dead connections that never send a close frame.
const WS_READ_TIMEOUT: Duration = Duration::from_secs(60);

/// Interval at which the client sends proactive ping frames.
/// Keeps the connection alive through idle-dropping proxies / NATs and provides
/// an independent liveness check even when the server sends no pings.
const WS_PING_INTERVAL: Duration = Duration::from_secs(20);

// ---------------------------------------------------------------------------
// Type aliases
// ---------------------------------------------------------------------------

type RawWsStream = tokio_tungstenite::WebSocketStream<
    tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
>;
type WsSink = futures_util::stream::SplitSink<RawWsStream, Message>;
type WsSource = futures_util::stream::SplitStream<RawWsStream>;

// ---------------------------------------------------------------------------
// ReconnectConfig — now Copy so it can be captured by spawned tasks without a
// manual field-by-field copy.
// ---------------------------------------------------------------------------

/// Configuration for the automatic reconnect engine (BR-007).
///
/// Default: 1 s initial backoff, 30 s max backoff, unlimited attempts.
#[derive(Clone, Copy)]
pub struct ReconnectConfig {
    /// Delay before the first reconnect attempt.
    pub initial_backoff: Duration,
    /// Upper bound on backoff delay.
    pub max_backoff: Duration,
    /// Maximum number of reconnect attempts. `None` means unlimited.
    pub max_attempts: Option<u32>,
}

impl Default for ReconnectConfig {
    fn default() -> Self {
        Self {
            initial_backoff: DEFAULT_INITIAL_BACKOFF,
            max_backoff: DEFAULT_MAX_BACKOFF,
            max_attempts: None,
        }
    }
}

// ---------------------------------------------------------------------------
// Internal helpers
// ---------------------------------------------------------------------------

/// Reason the inner receive loop exited.
#[derive(PartialEq, Eq)]
enum LoopExit {
    /// Connection lost or timed out — outer loop should reconnect.
    Reconnect,
    /// The `mpsc` receiver was dropped — the caller no longer needs this stream.
    ReceiverDropped,
}

/// Distinguishes market-data streams from user-data streams inside the unified loop.
enum StreamMode {
    /// Combined-stream endpoint — sends a SUBSCRIBE message after connect and
    /// unwraps the `{ "stream": "…", "data": {…} }` envelope.
    Market {
        streams: Vec<String>,
        stream_path: String,
    },
    /// User-data stream — server begins sending immediately, no SUBSCRIBE needed,
    /// raw JSON forwarded as-is.
    UserData { listen_key: String },
}

impl StreamMode {
    fn ws_url(&self, base_url: &Url) -> String {
        match self {
            StreamMode::Market { stream_path, .. } => {
                format!("{base_url}stream?streams={stream_path}")
            }
            StreamMode::UserData { listen_key } => {
                format!("{base_url}ws/{listen_key}")
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Public structs
// ---------------------------------------------------------------------------

/// Async WebSocket client — manages connection lifecycle and reconnection (BR-006, BR-007).
pub struct WebSocketClient {
    base_url: Url,
    reconnect_config: ReconnectConfig,
}

impl std::fmt::Debug for WebSocketClient {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("WebSocketClient")
            .field("base_url", &self.base_url.as_str())
            .finish()
    }
}

/// Handle for a market data stream.
///
/// Yields `Result<serde_json::Value, AsterDexError>` messages from the background receive loop.
/// Dropping this handle aborts the background task.
pub struct WebSocketStream {
    receiver: mpsc::Receiver<Result<serde_json::Value, AsterDexError>>,
    task_handle: JoinHandle<()>,
}

impl std::fmt::Debug for WebSocketStream {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("WebSocketStream").finish_non_exhaustive()
    }
}

impl WebSocketStream {
    /// Receive the next message from the stream.
    ///
    /// Returns `None` when the stream is closed (background task exited).
    pub async fn recv(&mut self) -> Option<Result<serde_json::Value, AsterDexError>> {
        self.receiver.recv().await
    }
}

impl Drop for WebSocketStream {
    fn drop(&mut self) {
        self.task_handle.abort();
    }
}

/// Handle for a user data stream. Same shape as [`WebSocketStream`] but for private events.
///
/// Dropping this handle aborts the background task.
pub struct UserDataStream {
    receiver: mpsc::Receiver<Result<serde_json::Value, AsterDexError>>,
    task_handle: JoinHandle<()>,
}

impl std::fmt::Debug for UserDataStream {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("UserDataStream").finish_non_exhaustive()
    }
}

impl UserDataStream {
    /// Receive the next user data event from the stream.
    ///
    /// Returns `None` when the stream is closed (background task ended).
    pub async fn recv(&mut self) -> Option<Result<serde_json::Value, AsterDexError>> {
        self.receiver.recv().await
    }

    /// Receive the next user data event as a typed `UserDataEvent`.
    pub async fn next_event(
        &mut self,
    ) -> Option<Result<crate::models::user_data::UserDataEvent, crate::rest::error::AsterDexError>>
    {
        let raw = self.receiver.recv().await?;
        match raw {
            Err(e) => Some(Err(e)),
            Ok(v) => Some(
                serde_json::from_value::<crate::models::user_data::UserDataEvent>(v).map_err(
                    |e| crate::rest::error::AsterDexError::SerdeError {
                        message: format!("failed to parse user data event: {e}"),
                    },
                ),
            ),
        }
    }
}

impl Drop for UserDataStream {
    fn drop(&mut self) {
        self.task_handle.abort();
    }
}

// ---------------------------------------------------------------------------
// WebSocketClient implementation
// ---------------------------------------------------------------------------

impl WebSocketClient {
    /// Construct a new WebSocket client.
    pub fn new(base_url: &str) -> Result<Self, AsterDexError> {
        let url = Self::normalize_ws_url(base_url)?;
        Ok(Self {
            base_url: url,
            reconnect_config: ReconnectConfig::default(),
        })
    }

    /// Builder method: set a custom reconnect configuration.
    pub fn with_reconnect_config(mut self, cfg: ReconnectConfig) -> Self {
        self.reconnect_config = cfg;
        self
    }

    /// Subscribe to market data streams (combined stream endpoint).
    ///
    /// # Errors
    /// - `WebSocketError("at least one stream name required")` if `streams` is empty.
    /// - `WebSocketError("max 200 streams per connection")` if > 200 streams.
    /// - `WebSocketError("connection failed: …")` if the initial connection fails.
    pub async fn subscribe(
        &self,
        streams: Vec<&str>,
    ) -> Result<WebSocketStream, AsterDexError> {
        if streams.is_empty() {
            return Err(AsterDexError::WebSocketError {
                message: "at least one stream name required".to_string(),
            });
        }
        if streams.len() > MAX_STREAMS_PER_CONNECTION {
            return Err(AsterDexError::WebSocketError {
                message: format!("max {} streams per connection", MAX_STREAMS_PER_CONNECTION),
            });
        }

        let stream_path = streams.join("/");
        let ws_url = format!("{}stream?streams={}", self.base_url, stream_path);

        let (ws_stream, _response) =
            connect_async(&ws_url)
                .await
                .map_err(|e| AsterDexError::WebSocketError {
                    message: format!("connection failed: {e}"),
                })?;

        let (tx, rx) = mpsc::channel(CHANNEL_BUFFER);
        let owned_streams: Vec<String> = streams.iter().map(|s| s.to_string()).collect();
        let mode = StreamMode::Market {
            streams: owned_streams,
            stream_path,
        };
        let base_url = self.base_url.clone();
        let config = self.reconnect_config; // Copy

        let task_handle = tokio::spawn(async move {
            Self::run_stream_loop(base_url, mode, tx, config, Some(ws_stream)).await;
        });

        Ok(WebSocketStream {
            receiver: rx,
            task_handle,
        })
    }

    /// Unsubscribe from specific streams on an active connection.
    ///
    /// Not yet supported — drop the [`WebSocketStream`] to close the connection.
    pub async fn unsubscribe(
        _stream: &mut WebSocketStream,
        _streams: Vec<&str>,
    ) -> Result<(), AsterDexError> {
        Err(AsterDexError::WebSocketError {
            message: "unsubscribe not yet supported -- drop stream to close".to_string(),
        })
    }

    /// Subscribe to the user data stream using a listen key.
    ///
    /// # Errors
    /// - `WebSocketError("user data stream connection failed: …")` if the initial connection fails.
    pub async fn subscribe_user_data(
        &self,
        listen_key: &str,
    ) -> Result<UserDataStream, AsterDexError> {
        let ws_url = format!("{}ws/{}", self.base_url, listen_key);

        let (ws_stream, _response) =
            connect_async(&ws_url)
                .await
                .map_err(|e| AsterDexError::WebSocketError {
                    message: format!("user data stream connection failed: {e}"),
                })?;

        let (tx, rx) = mpsc::channel(CHANNEL_BUFFER);
        let mode = StreamMode::UserData {
            listen_key: listen_key.to_string(),
        };
        let base_url = self.base_url.clone();
        let config = self.reconnect_config; // Copy

        let task_handle = tokio::spawn(async move {
            Self::run_stream_loop(base_url, mode, tx, config, Some(ws_stream)).await;
        });

        Ok(UserDataStream {
            receiver: rx,
            task_handle,
        })
    }

    // -------------------------------------------------------------------------
    // Internal: unified reconnect loop (replaces separate market / user-data loops)
    // -------------------------------------------------------------------------

    async fn run_stream_loop(
        base_url: Url,
        mode: StreamMode,
        tx: mpsc::Sender<Result<serde_json::Value, AsterDexError>>,
        config: ReconnectConfig,
        initial_connection: Option<RawWsStream>,
    ) {
        let mut reconnect = ReconnectState::new(
            config.initial_backoff,
            config.max_backoff,
            config.max_attempts,
        );
        let mut cached_conn = initial_connection;

        loop {
            let ws_url = mode.ws_url(&base_url);

            let ws_result = if let Some(ws) = cached_conn.take() {
                Ok(ws)
            } else {
                connect_async(&ws_url).await.map(|(ws, _)| ws)
            };

            match ws_result {
                Ok(ws_stream) => {
                    reconnect.reset();
                    tracing::info!(url = %ws_url, "WebSocket connected");

                    let (mut sink, mut source) = ws_stream.split();

                    // Post-connect setup: market streams must send a SUBSCRIBE frame.
                    let setup_ok = match &mode {
                        StreamMode::Market { streams, .. } => {
                            let sub_msg = serde_json::json!({
                                "method": "SUBSCRIBE",
                                "params": streams,
                                "id": 1,
                            });
                            if sink.send(Message::Text(sub_msg.to_string())).await.is_err() {
                                tracing::warn!("Failed to send SUBSCRIBE — reconnecting");
                                false
                            } else {
                                true
                            }
                        }
                        StreamMode::UserData { .. } => true,
                    };

                    if setup_ok {
                        let exit =
                            Self::receive_loop(&mode, &mut sink, &mut source, &tx).await;
                        let _ = sink.close().await;
                        if exit == LoopExit::ReceiverDropped {
                            return;
                        }
                    } else {
                        let _ = sink.close().await;
                    }
                }
                Err(e) => {
                    tracing::warn!(error = ?e, "WebSocket connection failed");
                }
            }

            // Exponential backoff before next reconnect attempt.
            match reconnect.next_backoff() {
                Some(delay) => {
                    tracing::info!(
                        attempt = reconnect.current_attempt,
                        backoff_secs = delay.as_secs(),
                        "WebSocket reconnect attempt",
                    );
                    tokio::time::sleep(delay).await;
                }
                None => {
                    tracing::warn!("WebSocket max reconnect attempts reached");
                    let _ = tx
                        .send(Err(AsterDexError::WebSocketError {
                            message: "max reconnect attempts reached".to_string(),
                        }))
                        .await;
                    return;
                }
            }
        }
    }

    // -------------------------------------------------------------------------
    // Internal: inner receive loop — runs while the connection is healthy.
    //
    // Uses tokio::select! with three arms:
    //   1. Incoming WS frame (source.next()) — resets the read-timeout deadline.
    //   2. Periodic heartbeat tick — sends a proactive Ping to keep the
    //      connection alive and detect silent dead connections.
    //   3. Read-timeout deadline — fires when no frame arrives within
    //      WS_READ_TIMEOUT; triggers a reconnect.
    //
    // Data messages use try_send (non-blocking): if the consumer is slow and
    // the channel is full, the message is dropped with a warning rather than
    // blocking this loop — which would delay Pong responses and starve the
    // timeout deadline.
    // -------------------------------------------------------------------------

    async fn receive_loop(
        mode: &StreamMode,
        sink: &mut WsSink,
        source: &mut WsSource,
        tx: &mpsc::Sender<Result<serde_json::Value, AsterDexError>>,
    ) -> LoopExit {
        use tokio::sync::mpsc::error::TrySendError;

        // First heartbeat fires after WS_PING_INTERVAL, not immediately.
        let mut ping_interval = tokio::time::interval_at(
            tokio::time::Instant::now() + WS_PING_INTERVAL,
            WS_PING_INTERVAL,
        );
        ping_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);

        // Read-timeout future — reset each time any frame arrives.
        let deadline = tokio::time::sleep(WS_READ_TIMEOUT);
        tokio::pin!(deadline);

        loop {
            tokio::select! {
                // Arm 1: read timeout — no frame received in WS_READ_TIMEOUT seconds.
                _ = &mut deadline => {
                    tracing::warn!(
                        timeout_secs = WS_READ_TIMEOUT.as_secs(),
                        "WebSocket read timeout — reconnecting"
                    );
                    return LoopExit::Reconnect;
                }

                // Arm 2: heartbeat — send a proactive Ping.
                _ = ping_interval.tick() => {
                    tracing::trace!("Sending heartbeat Ping");
                    if sink.send(Message::Ping(vec![])).await.is_err() {
                        tracing::warn!("Heartbeat Ping failed — reconnecting");
                        return LoopExit::Reconnect;
                    }
                }

                // Arm 3: incoming frame from the server.
                msg = source.next() => {
                    // Any received frame resets the read-timeout deadline.
                    deadline
                        .as_mut()
                        .reset(tokio::time::Instant::now() + WS_READ_TIMEOUT);

                    match msg {
                        Some(Ok(Message::Text(text))) => {
                            match serde_json::from_str::<serde_json::Value>(&text) {
                                Err(e) => {
                                    tracing::warn!(error = ?e, "Failed to parse WS message");
                                }
                                Ok(mut v) => {
                                    let data = match mode {
                                        StreamMode::Market { .. } => {
                                            // Unwrap combined-stream envelope:
                                            // { "stream": "…", "data": {…} }
                                            if let Some(d) = v.get_mut("data") {
                                                d.take()
                                            } else {
                                                v
                                            }
                                        }
                                        StreamMode::UserData { .. } => v,
                                    };
                                    match tx.try_send(Ok(data)) {
                                        Ok(()) => {}
                                        Err(TrySendError::Full(_)) => {
                                            tracing::warn!(
                                                "WS receive channel full — dropping message"
                                            );
                                        }
                                        Err(TrySendError::Closed(_)) => {
                                            return LoopExit::ReceiverDropped;
                                        }
                                    }
                                }
                            }
                        }

                        Some(Ok(Message::Ping(data))) => {
                            // Respond immediately — never blocked by the data channel.
                            // A short timeout prevents this arm from stalling the loop
                            // if the TCP send buffer is momentarily full.
                            tracing::trace!("Received server Ping — sending Pong");
                            let _ = tokio::time::timeout(
                                Duration::from_secs(5),
                                sink.send(Message::Pong(data)),
                            )
                            .await;
                        }

                        Some(Ok(Message::Close(_))) | None => {
                            tracing::warn!("WebSocket disconnected (close frame or EOF)");
                            return LoopExit::Reconnect;
                        }

                        Some(Err(e)) => {
                            tracing::warn!(error = ?e, "WebSocket error");
                            return LoopExit::Reconnect;
                        }

                        _ => {
                            // Pong, Binary, Frame — ignored
                        }
                    }
                }
            }
        }
    }

    /// Normalize a WebSocket base URL: ensure trailing slash so path concatenation works.
    fn normalize_ws_url(url: &str) -> Result<Url, AsterDexError> {
        let url = if url.ends_with('/') {
            url.to_string()
        } else {
            format!("{url}/")
        };
        Url::parse(&url).map_err(|e| AsterDexError::ConfigError {
            message: format!("invalid WebSocket base URL: {e}"),
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    // US-010: subscribe with 201 streams returns WebSocketError (validation before connect)
    #[tokio::test]
    async fn max_streams_exceeded_returns_error() {
        let client =
            WebSocketClient::new("wss://fstream.asterdex-testnet.com").expect("valid URL");
        let streams: Vec<&str> = (0..201).map(|_| "btcusdt@aggTrade").collect();
        let result = client.subscribe(streams).await;
        assert!(
            matches!(result, Err(AsterDexError::WebSocketError { ref message }) if message.contains("max 200 streams")),
            "expected WebSocketError about max streams, got: {result:?}"
        );
    }

    // US-010: subscribe with empty vec returns WebSocketError
    #[tokio::test]
    async fn empty_streams_returns_error() {
        let client =
            WebSocketClient::new("wss://fstream.asterdex-testnet.com").expect("valid URL");
        let result = client.subscribe(vec![]).await;
        assert!(
            matches!(result, Err(AsterDexError::WebSocketError { ref message }) if message.contains("at least one stream")),
            "expected WebSocketError about empty streams, got: {result:?}"
        );
    }

    // US-010: invalid URL returns ConfigError
    #[tokio::test]
    async fn new_invalid_url_returns_config_error() {
        let result = WebSocketClient::new("not a url :::");
        assert!(
            matches!(result, Err(AsterDexError::ConfigError { .. })),
            "expected ConfigError, got: {result:?}"
        );
    }

    // ReconnectConfig is now Copy — verify it can be copied without a manual field clone.
    #[test]
    fn reconnect_config_is_copy() {
        let cfg = ReconnectConfig::default();
        let _copy = cfg; // moves by copy
        let _another = cfg; // still accessible — proves Copy
    }
}