ws_reconnect_client/

stream.rs

use futures_util::{Stream, StreamExt, SinkExt, Future};
use serde::de::DeserializeOwned;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use tokio::sync::Mutex;
use tokio_tungstenite::tungstenite::Message;

use crate::{PingManager, Result, WebSocketError, WsReader, WsWriter};

/// A WebSocket message stream that automatically handles PING/PONG and yields parsed messages
///
/// This stream:
/// - Responds to server PINGs automatically with PONGs
/// - Sends periodic client PINGs to keep the connection alive
/// - Parses incoming text messages as JSON (supports both single messages and arrays)
/// - Filters out control messages (PING, PONG, empty strings)
/// - Only yields successfully parsed data messages of type T
pub struct MessageStream<T>
where
    T: DeserializeOwned,
{
    reader: WsReader,
    writer: Arc<Mutex<Option<WsWriter>>>,
    ping_manager: Option<PingManager>,
    // Buffer for parsed messages (one WS message can contain multiple T messages)
    message_buffer: Vec<T>,
    _phantom: std::marker::PhantomData<T>,
}

impl<T> MessageStream<T>
where
    T: DeserializeOwned,
{
    /// Create a new MessageStream
    ///
    /// # Arguments
    /// * `reader` - The WebSocket reader
    /// * `writer` - Shared WebSocket writer (Arc<Mutex<Option<WsWriter>>>)
    /// * `ping_interval_secs` - Optional interval for sending client pings (0 = disabled)
    pub fn new(reader: WsReader, writer: Arc<Mutex<Option<WsWriter>>>, ping_interval_secs: u64) -> Self {
        let ping_manager = if ping_interval_secs > 0 {
            Some(PingManager::new(ping_interval_secs))
        } else {
            None
        };

        Self {
            reader,
            writer,
            ping_manager,
            message_buffer: Vec::new(),
            _phantom: std::marker::PhantomData,
        }
    }

    /// Parse a text message as either a single message or array of messages
    fn parse_text_message(&self, text: &str) -> Result<Vec<T>> {
        // Skip control messages
        let trimmed = text.trim();
        if trimmed.is_empty() || trimmed == "PING" || trimmed == "PONG" {
            return Ok(Vec::new());
        }

        // Log raw message for debugging
        if trimmed.len() < 500 {
        } else {
        }

        // Try array first if starts with '[', then single message, then empty vec
        if trimmed.starts_with('[') {
            serde_json::from_str::<Vec<T>>(text).ok()
        } else {
            None
        }
        .or_else(|| serde_json::from_str::<T>(text).ok().map(|msg| vec![msg]))
        .ok_or(())
        .or(Ok(Vec::new()))
    }

    /// Handle an incoming message, possibly sending PONG in response to PING
    async fn handle_incoming_message(&mut self, msg: Message) -> Result<Vec<T>> {
        match msg {
            Message::Text(text) => self.parse_text_message(&text),

            Message::Binary(data) => {
                let text = String::from_utf8_lossy(&data);
                self.parse_text_message(&text)
            }

            Message::Ping(ping) => {
                // Respond to server ping with pong
                let mut writer_guard = self.writer.lock().await;
                if let Some(writer) = writer_guard.as_mut() {
                    writer
                        .send(Message::Pong(ping))
                        .await
                        .map_err(|_| WebSocketError::SendError)?;
                }
                Ok(Vec::new())
            }

            Message::Pong(_) => {
                // Pong received, connection is alive
                Ok(Vec::new())
            }

            Message::Close(_) => Err(WebSocketError::ConnectionClosed),

            Message::Frame(_) => Ok(Vec::new()),
        }
    }

    /// Send a periodic ping to keep the connection alive
    async fn send_ping(&mut self) -> Result<()> {
        let mut writer_guard = self.writer.lock().await;
        if let Some(writer) = writer_guard.as_mut() {
            writer
                .send(Message::Ping(vec![].into()))
                .await
                .map_err(|_| WebSocketError::SendError)?;
        }
        Ok(())
    }

    /// Get the next parsed message from the stream
    ///
    /// This method handles:
    /// - Incoming messages (responding to PINGs, parsing text)
    /// - Outgoing periodic PINGs (if enabled)
    /// - Returns the next successfully parsed message
    async fn next_message(&mut self) -> Option<Result<T>> {
        // Buffer for parsed messages (since one websocket message can contain multiple T messages)
        let mut buffer = Vec::new();

        loop {
            // If we have buffered messages, return the first one
            if !buffer.is_empty() {
                return Some(Ok(buffer.remove(0)));
            }

            tokio::select! {
                // Handle incoming messages
                msg = self.reader.next() => {
                    match msg {
                        Some(Ok(msg)) => {
                            match self.handle_incoming_message(msg).await {
                                Ok(messages) => {
                                    buffer = messages;
                                    // Continue loop to check buffer and return message if available
                                }
                                Err(e) => return Some(Err(e)),
                            }
                        }
                        Some(Err(e)) => {
                            return Some(Err(WebSocketError::Tungstenite(e)));
                        }
                        None => {
                            return None; // Stream ended
                        }
                    }
                }

                // Send periodic pings if enabled
                _ = async {
                    if let Some(ref mut pm) = self.ping_manager {
                        pm.wait_for_next_ping().await;
                    } else {
                        // Never completes if pings are disabled
                        std::future::pending::<()>().await;
                    }
                } => {
                    if let Err(e) = self.send_ping().await {
                        return Some(Err(e));
                    }
                }
            }
        }
    }
}

impl<T> Stream for MessageStream<T>
where
    T: DeserializeOwned + Unpin + 'static,
{
    type Item = Result<T>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        loop {
            // First check if we have buffered messages
            if !self.message_buffer.is_empty() {
                return Poll::Ready(Some(Ok(self.message_buffer.remove(0))));
            }

            // No buffered messages, poll next_message to fill the buffer
            let fut = self.next_message();
            tokio::pin!(fut);

            match fut.poll(cx) {
                Poll::Ready(Some(Ok(msg))) => {
                    // Got a message, return it immediately
                    return Poll::Ready(Some(Ok(msg)));
                }
                Poll::Ready(Some(Err(e))) => {
                    return Poll::Ready(Some(Err(e)));
                }
                Poll::Ready(None) => {
                    return Poll::Ready(None);
                }
                Poll::Pending => {
                    return Poll::Pending;
                }
            }
        }
    }
}