volt_client_grpc/

volt_connection.rs

//! Volt connection management.
//!
//! Handles persistent connections to the Volt with automatic reconnection
//! and ping/pong keep-alive.

use crate::config::VoltClientConfig;
use crate::error::{Result, VoltError};
use crate::proto::{ConnectHello, ConnectRequest};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{mpsc, Mutex, RwLock};
use tokio::time::{interval, Instant};

/// Connection state
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ConnectionState {
    /// Not connected
    Disconnected,
    /// Connecting
    Connecting,
    /// Connected and active
    Connected,
    /// Reconnecting after disconnect
    Reconnecting,
}

/// Events emitted by the connection
#[derive(Debug, Clone)]
pub enum ConnectionEvent {
    /// Connection established with connection ID
    Connected(String),
    /// Connection lost
    Disconnected,
    /// Ping received with timestamp
    Ping(u64),
    /// Error occurred
    Error(String),
    /// Event received from server
    Event(serde_json::Value),
    /// Invoke request received
    InvokeRequest(serde_json::Value),
}

/// Channel for sending requests on the connection
pub type RequestSender = mpsc::Sender<ConnectRequest>;
/// Channel for receiving connection events
pub type EventReceiver = mpsc::Receiver<ConnectionEvent>;

/// Manages a persistent connection to a Volt server
pub struct VoltConnection {
    /// Current connection state
    state: Arc<RwLock<ConnectionState>>,
    /// Connection ID assigned by the server
    connection_id: Arc<RwLock<String>>,
    /// Configuration
    _config: VoltClientConfig,
    /// Whether auto-reconnect is enabled
    auto_retry: bool,
    /// Ping interval in milliseconds
    ping_interval: Duration,
    /// Reconnect interval in milliseconds
    reconnect_interval: Duration,
    /// Timeout interval in milliseconds  
    timeout_interval: Duration,
    /// Flag to signal shutdown
    dying: Arc<RwLock<bool>>,
    /// Event sender
    event_tx: Option<mpsc::Sender<ConnectionEvent>>,
    /// Request sender for the active connection
    request_tx: Arc<Mutex<Option<RequestSender>>>,
    /// Last ping timestamp
    last_ping: Arc<RwLock<Option<Instant>>>,
}

impl VoltConnection {
    /// Create a new connection manager
    pub fn new(config: &VoltClientConfig) -> Self {
        Self {
            state: Arc::new(RwLock::new(ConnectionState::Disconnected)),
            connection_id: Arc::new(RwLock::new(String::new())),
            _config: config.clone(),
            auto_retry: config.auto_reconnect,
            ping_interval: Duration::from_millis(config.ping_interval),
            reconnect_interval: Duration::from_millis(config.reconnect_interval),
            timeout_interval: Duration::from_millis(config.timeout_interval),
            dying: Arc::new(RwLock::new(false)),
            event_tx: None,
            request_tx: Arc::new(Mutex::new(None)),
            last_ping: Arc::new(RwLock::new(None)),
        }
    }

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

    /// Get the connection ID
    pub async fn connection_id(&self) -> String {
        self.connection_id.read().await.clone()
    }

    /// Check if connected
    pub async fn is_connected(&self) -> bool {
        *self.state.read().await == ConnectionState::Connected
    }

    /// Connect to the Volt server
    ///
    /// Returns a receiver for connection events
    pub async fn connect(
        &mut self,
        hello_payload: Option<serde_json::Value>,
    ) -> Result<EventReceiver> {
        // Create event channel
        let (event_tx, event_rx) = mpsc::channel(100);
        self.event_tx = Some(event_tx);

        // Set state to connecting
        *self.state.write().await = ConnectionState::Connecting;

        // Start connection task
        self.start_connection(hello_payload).await?;

        Ok(event_rx)
    }

    /// Start the connection process
    async fn start_connection(&self, _hello_payload: Option<serde_json::Value>) -> Result<()> {
        let state = self.state.clone();
        let connection_id = self.connection_id.clone();
        let event_tx = self.event_tx.clone();
        let dying = self.dying.clone();
        let auto_retry = self.auto_retry;
        let reconnect_interval = self.reconnect_interval;
        let ping_interval = self.ping_interval;
        let timeout_interval = self.timeout_interval;
        let _request_tx = self.request_tx.clone();
        let last_ping = self.last_ping.clone();

        // Build hello message
        let hello = ConnectHello {
            ping_interval: ping_interval.as_millis() as u64,
            timestamp: std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_millis() as u64,
        };

        // Spawn connection task
        tokio::spawn(async move {
            loop {
                // Check if we're shutting down
                if *dying.read().await {
                    tracing::debug!("Connection dying, exiting loop");
                    break;
                }

                // TODO: Establish actual gRPC connection here
                // This is a placeholder for the connection logic
                // In a real implementation, this would:
                // 1. Create a gRPC channel
                // 2. Call the Connect RPC
                // 3. Handle the bidirectional stream

                tracing::debug!("Would connect with hello: {:?}", hello);

                // Simulate connection established
                *state.write().await = ConnectionState::Connected;
                *connection_id.write().await = uuid::Uuid::new_v4().to_string();

                if let Some(ref tx) = event_tx {
                    let conn_id = connection_id.read().await.clone();
                    let _ = tx.send(ConnectionEvent::Connected(conn_id)).await;
                }

                // Start ping loop
                let mut ping_timer = interval(ping_interval);
                loop {
                    tokio::select! {
                        _ = ping_timer.tick() => {
                            // Send ping
                            let now = std::time::SystemTime::now()
                                .duration_since(std::time::UNIX_EPOCH)
                                .unwrap()
                                .as_millis() as u64;

                            *last_ping.write().await = Some(Instant::now());
                            tracing::trace!("Sending ping at {}", now);

                            // TODO: Actually send ping on the stream
                        }

                        // Check for timeout
                        _ = tokio::time::sleep(timeout_interval) => {
                            let last = last_ping.read().await;
                            if let Some(last_ping_time) = *last {
                                if last_ping_time.elapsed() > timeout_interval {
                                    tracing::warn!("Connection timed out");
                                    break;
                                }
                            }
                        }
                    }

                    // Check if dying
                    if *dying.read().await {
                        break;
                    }
                }

                // Connection lost
                *state.write().await = ConnectionState::Disconnected;
                *connection_id.write().await = String::new();

                if let Some(ref tx) = event_tx {
                    let _ = tx.send(ConnectionEvent::Disconnected).await;
                }

                // Should we reconnect?
                if !auto_retry || *dying.read().await {
                    break;
                }

                // Wait before reconnecting
                *state.write().await = ConnectionState::Reconnecting;
                tokio::time::sleep(reconnect_interval).await;
            }
        });

        Ok(())
    }

    /// Disconnect from the server
    pub async fn disconnect(&mut self) {
        *self.dying.write().await = true;

        // Close the request channel
        *self.request_tx.lock().await = None;

        *self.state.write().await = ConnectionState::Disconnected;
        *self.connection_id.write().await = String::new();
    }

    /// Send a message on the connection
    pub async fn send(&self, request: ConnectRequest) -> Result<()> {
        let tx = self.request_tx.lock().await;

        if let Some(ref sender) = *tx {
            sender
                .send(request)
                .await
                .map_err(|_| VoltError::NotConnected)?;
            Ok(())
        } else {
            Err(VoltError::NotConnected)
        }
    }
}

impl Drop for VoltConnection {
    fn drop(&mut self) {
        // Signal shutdown
        // Note: This is sync context, so we can't use async
        // The actual cleanup happens in the spawned task
    }
}

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

    #[tokio::test]
    async fn test_connection_state() {
        let config = VoltClientConfig {
            client_name: "test".to_string(),
            volt: VoltConfig::default(),
            ..Default::default()
        };

        let conn = VoltConnection::new(&config);
        assert_eq!(conn.state().await, ConnectionState::Disconnected);
    }
}