rusty_bin/monitor/

lockfree_exchange_manager.rs

//! Lock-free exchange manager for high-frequency trading monitoring
//!
//! This module provides a lock-free exchange connection manager that eliminates
//! contention in critical paths while maintaining thread safety through atomic operations.

use crate::monitor::collector::{
    CollectionTask, DataType, MarketDataEvent, Result, exchange_client::ExchangeClient,
};
use crossbeam_utils::CachePadded;
use flume::Sender;
use parking_lot::RwLock;
use quanta::Clock;
use rusty_common::collections::FxHashMap;
use smartstring::alias::String as SmartString;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::time::Instant;

/// Atomic exchange connection state
#[derive(Debug)]
struct AtomicExchangeState {
    /// Whether the exchange is currently connected
    connected: AtomicBool,
    /// Last successful message timestamp
    last_message_timestamp: AtomicU64,
    /// Number of reconnection attempts
    reconnection_attempts: AtomicUsize,
    /// Connection establishment timestamp
    connection_start_time: AtomicU64,
    /// Total messages received
    messages_received: AtomicU64,
    /// Total bytes received
    bytes_received: AtomicU64,
    /// Last error timestamp
    last_error_timestamp: AtomicU64,
    /// Whether health check is currently in progress
    health_check_in_progress: AtomicBool,
}

impl AtomicExchangeState {
    const fn new() -> Self {
        Self {
            connected: AtomicBool::new(false),
            last_message_timestamp: AtomicU64::new(0),
            reconnection_attempts: AtomicUsize::new(0),
            connection_start_time: AtomicU64::new(0),
            messages_received: AtomicU64::new(0),
            bytes_received: AtomicU64::new(0),
            last_error_timestamp: AtomicU64::new(0),
            health_check_in_progress: AtomicBool::new(false),
        }
    }

    /// Mark exchange as connected (lock-free)
    #[inline(always)]
    fn set_connected(&self, connected: bool) {
        self.connected.store(connected, Ordering::Release);
        if connected {
            self.connection_start_time
                .store(Clock::new().raw(), Ordering::Relaxed);
        }
    }

    /// Check if exchange is connected (lock-free)
    #[inline(always)]
    fn is_connected(&self) -> bool {
        self.connected.load(Ordering::Acquire)
    }

    /// Record a received message (lock-free)
    #[inline(always)]
    fn record_message(&self, bytes: u64) {
        self.messages_received.fetch_add(1, Ordering::Relaxed);
        self.bytes_received.fetch_add(bytes, Ordering::Relaxed);
        self.last_message_timestamp
            .store(Clock::new().raw(), Ordering::Relaxed);
    }

    /// Record an error (lock-free)
    #[inline(always)]
    fn record_error(&self) {
        self.last_error_timestamp
            .store(Clock::new().raw(), Ordering::Relaxed);
        self.set_connected(false);
    }

    /// Increment reconnection attempts (lock-free)
    #[inline(always)]
    fn increment_reconnection_attempts(&self) -> usize {
        self.reconnection_attempts.fetch_add(1, Ordering::Relaxed)
    }

    /// Reset reconnection attempts (lock-free)
    #[inline(always)]
    fn reset_reconnection_attempts(&self) {
        self.reconnection_attempts.store(0, Ordering::Relaxed);
    }

    /// Start health check if not already in progress (lock-free)
    #[inline(always)]
    fn try_start_health_check(&self) -> bool {
        self.health_check_in_progress
            .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
            .is_ok()
    }

    /// Finish health check (lock-free)
    #[inline(always)]
    fn finish_health_check(&self) {
        self.health_check_in_progress
            .store(false, Ordering::Release);
    }
}

/// Exchange connection metadata with atomic fields
struct ExchangeMetadata {
    /// Exchange client (wrapped in RwLock for rare mutations)
    client: RwLock<Option<ExchangeClient>>,
    /// Atomic state for frequent reads
    state: AtomicExchangeState,
    /// Active collection tasks
    active_tasks: RwLock<FxHashMap<String, CollectionTask>>,
    /// Last health check timestamp
    last_health_check: AtomicU64,
}

impl ExchangeMetadata {
    fn new() -> Self {
        Self {
            client: RwLock::new(None),
            state: AtomicExchangeState::new(),
            active_tasks: RwLock::new(FxHashMap::default()),
            last_health_check: AtomicU64::new(0),
        }
    }
}

/// Lock-free exchange manager optimized for high-frequency operations
pub struct LockFreeExchangeManager {
    /// Exchange metadata by name (uses RwLock only for structural changes)
    exchanges: RwLock<FxHashMap<SmartString, Arc<ExchangeMetadata>>>,

    /// Global atomic counters for fast statistics
    total_exchanges: CachePadded<AtomicUsize>,
    total_active_tasks: CachePadded<AtomicUsize>,
    total_messages_received: CachePadded<AtomicU64>,
    total_bytes_received: CachePadded<AtomicU64>,
    total_errors: CachePadded<AtomicU64>,

    /// Manager state
    manager_start_time: Instant,
    shutdown_requested: AtomicBool,
}

impl LockFreeExchangeManager {
    /// Create a new lock-free exchange manager
    pub fn new() -> Self {
        Self {
            exchanges: RwLock::new(FxHashMap::default()),
            total_exchanges: CachePadded::new(AtomicUsize::new(0)),
            total_active_tasks: CachePadded::new(AtomicUsize::new(0)),
            total_messages_received: CachePadded::new(AtomicU64::new(0)),
            total_bytes_received: CachePadded::new(AtomicU64::new(0)),
            total_errors: CachePadded::new(AtomicU64::new(0)),
            manager_start_time: Instant::now(),
            shutdown_requested: AtomicBool::new(false),
        }
    }

    /// Add an exchange to the manager (infrequent operation, uses lock)
    pub async fn add_exchange(
        &self,
        exchange_name: SmartString,
        mut client: ExchangeClient,
    ) -> Result<()> {
        // Connect the client first
        match client.connect().await {
            Ok(()) => {
                log::info!("Successfully connected to exchange: {exchange_name}");

                // Create metadata
                let metadata = Arc::new(ExchangeMetadata::new());
                metadata.state.set_connected(true);
                *metadata.client.write() = Some(client);

                // Add to exchanges map
                {
                    let mut exchanges = self.exchanges.write();
                    exchanges.insert(exchange_name.clone(), metadata);
                    self.total_exchanges.fetch_add(1, Ordering::Relaxed);
                }

                Ok(())
            }
            Err(e) => {
                log::error!("Failed to connect to exchange {exchange_name}: {e}");

                // Still add to map but mark as disconnected for retry
                let metadata = Arc::new(ExchangeMetadata::new());
                metadata.state.record_error();
                *metadata.client.write() = Some(client);

                {
                    let mut exchanges = self.exchanges.write();
                    exchanges.insert(exchange_name, metadata);
                    self.total_exchanges.fetch_add(1, Ordering::Relaxed);
                }

                self.total_errors.fetch_add(1, Ordering::Relaxed);
                Err(e)
            }
        }
    }

    /// Remove an exchange from the manager (infrequent operation)
    pub async fn remove_exchange(&self, exchange_name: &str) -> Result<()> {
        let exchange_key: SmartString = exchange_name.into();

        let metadata = {
            let exchanges = self.exchanges.read();
            exchanges.get(&exchange_key).cloned()
        };

        if let Some(metadata) = metadata {
            // Disconnect client if present
            let client_opt = {
                let mut client_lock = metadata.client.write();
                client_lock.take()
            }; // Lock is dropped here

            if let Some(mut client) = client_opt
                && let Err(e) = client.disconnect().await
            {
                log::error!("Error disconnecting exchange {exchange_name}: {e}");
            }

            // Mark as disconnected
            metadata.state.set_connected(false);

            // Remove from map
            {
                let mut exchanges = self.exchanges.write();
                exchanges.remove(&exchange_key);
                self.total_exchanges.fetch_sub(1, Ordering::Relaxed);
            }

            log::info!("Successfully removed exchange: {exchange_name}");
        }

        Ok(())
    }

    /// Start collection for an exchange (mostly lock-free)
    pub async fn start_collection(
        &self,
        exchange_name: &str,
        symbols: Vec<SmartString>,
        data_types: Vec<DataType>,
        event_sender: Sender<MarketDataEvent>,
    ) -> Result<()> {
        let exchange_key: SmartString = exchange_name.into();

        let metadata = {
            let exchanges = self.exchanges.read();
            exchanges.get(&exchange_key).cloned()
        };

        let metadata = metadata.ok_or_else(|| {
            Box::new(std::io::Error::new(
                std::io::ErrorKind::NotFound,
                format!("Exchange {exchange_name} not found"),
            )) as Box<dyn std::error::Error + Send + Sync>
        })?;

        // Check if connected (lock-free)
        if !metadata.state.is_connected() {
            return Err(Box::new(std::io::Error::new(
                std::io::ErrorKind::NotConnected,
                format!("Exchange {exchange_name} is not connected"),
            )) as Box<dyn std::error::Error + Send + Sync>)?;
        }

        // Start collection on client with panic safety
        {
            let has_client = metadata.client.read().is_some();
            if !has_client {
                return Err(Box::new(std::io::Error::new(
                    std::io::ErrorKind::NotFound,
                    format!("No client for exchange {exchange_name}"),
                ))
                    as Box<dyn std::error::Error + Send + Sync>)?;
            }
        }

        // PANIC SAFETY: Use a guard to ensure client is always put back
        struct ClientGuard<'a> {
            client: Option<ExchangeClient>,
            storage: &'a RwLock<Option<ExchangeClient>>,
        }

        impl<'a> Drop for ClientGuard<'a> {
            fn drop(&mut self) {
                if let Some(client) = self.client.take() {
                    *self.storage.write() = Some(client);
                }
            }
        }

        // Take client out with guard protection
        let mut guard = ClientGuard {
            client: metadata.client.write().take(),
            storage: &metadata.client,
        };

        if let Some(client) = &mut guard.client {
            client
                .start_collection(symbols.clone(), data_types.clone(), event_sender)
                .await?;
        }

        // Add tasks to tracking
        {
            let mut active_tasks = metadata.active_tasks.write();
            let mut task_count = 0;

            for symbol in &symbols {
                for data_type in &data_types {
                    let task_id = format!("{exchange_name}:{symbol}:{data_type:?}");
                    active_tasks.insert(
                        task_id,
                        CollectionTask {
                            exchange: exchange_key.clone(),
                            symbol: symbol.clone(),
                            data_types: vec![*data_type],
                            enabled: true,
                            retry_count: 0,
                            max_retries: 3,
                            backoff_ms: 1000,
                        },
                    );
                    task_count += 1;
                }
            }

            self.total_active_tasks
                .fetch_add(task_count, Ordering::Relaxed);
        }

        log::info!(
            "Started collection for exchange {} with {} symbols and {} data types",
            exchange_name,
            symbols.len(),
            data_types.len()
        );

        Ok(())
    }

    /// Record a message for an exchange (lock-free)
    #[inline(always)]
    pub fn record_message(&self, exchange_name: &str, bytes: u64) {
        let exchange_key: SmartString = exchange_name.into();

        // Always update global counters first
        self.total_messages_received.fetch_add(1, Ordering::Relaxed);
        self.total_bytes_received
            .fetch_add(bytes, Ordering::Relaxed);

        // Fast path: get metadata without lock
        let metadata = {
            let exchanges = self.exchanges.read();
            exchanges.get(&exchange_key).cloned()
        };

        if let Some(metadata) = metadata {
            metadata.state.record_message(bytes);
        }
    }

    /// Record an error for an exchange (lock-free)
    #[inline(always)]
    pub fn record_error(&self, exchange_name: &str) {
        let exchange_key: SmartString = exchange_name.into();

        // Always update global counter first
        self.total_errors.fetch_add(1, Ordering::Relaxed);

        let metadata = {
            let exchanges = self.exchanges.read();
            exchanges.get(&exchange_key).cloned()
        };

        if let Some(metadata) = metadata {
            metadata.state.record_error();
        }
    }

    /// Get connection status for all exchanges (mostly lock-free reads)
    pub fn get_connection_status(&self) -> FxHashMap<SmartString, bool> {
        let mut status_map = FxHashMap::default();

        let exchanges = self.exchanges.read();
        for (exchange_name, metadata) in exchanges.iter() {
            status_map.insert(exchange_name.clone(), metadata.state.is_connected());
        }

        status_map
    }

    /// Perform health check on all exchanges (lock-free coordination)
    pub async fn health_check(&self) -> FxHashMap<SmartString, bool> {
        let mut health_status = FxHashMap::default();

        let exchanges = {
            let exchanges = self.exchanges.read();
            exchanges.clone()
        };

        for (exchange_name, metadata) in exchanges {
            // Try to start health check (lock-free)
            if metadata.state.try_start_health_check() {
                let is_healthy = {
                    // For now, just use the atomic connection state to avoid
                    // complex async handling in this health check
                    // A more sophisticated health check can be implemented later
                    metadata.state.is_connected()
                };

                // Update connection state if health check failed
                if !is_healthy {
                    metadata.state.record_error();
                }

                metadata
                    .last_health_check
                    .store(Clock::new().raw(), Ordering::Relaxed);
                metadata.state.finish_health_check();

                health_status.insert(exchange_name, is_healthy);
            } else {
                // Health check already in progress, use cached connection status
                health_status.insert(exchange_name, metadata.state.is_connected());
            }
        }

        health_status
    }

    /// Get comprehensive statistics (mostly lock-free reads)
    pub fn get_statistics(&self) -> ExchangeManagerStats {
        let uptime_seconds = self.manager_start_time.elapsed().as_secs_f64();

        // Global atomic reads
        let total_messages = self.total_messages_received.load(Ordering::Relaxed);
        let total_bytes = self.total_bytes_received.load(Ordering::Relaxed);

        // Calculate rates
        let messages_per_second = if uptime_seconds > 0.0 {
            total_messages as f64 / uptime_seconds
        } else {
            0.0
        };

        let bytes_per_second = if uptime_seconds > 0.0 {
            total_bytes as f64 / uptime_seconds
        } else {
            0.0
        };

        // Per-exchange statistics
        let mut exchange_stats = FxHashMap::default();
        {
            let exchanges = self.exchanges.read();
            for (exchange_name, metadata) in exchanges.iter() {
                let exchange_messages = metadata.state.messages_received.load(Ordering::Relaxed);
                let exchange_bytes = metadata.state.bytes_received.load(Ordering::Relaxed);
                let exchange_rate = if uptime_seconds > 0.0 {
                    exchange_messages as f64 / uptime_seconds
                } else {
                    0.0
                };

                exchange_stats.insert(
                    exchange_name.clone(),
                    ExchangeStats {
                        connected: metadata.state.is_connected(),
                        messages_received: exchange_messages,
                        bytes_received: exchange_bytes,
                        messages_per_second: exchange_rate,
                        reconnection_attempts: metadata
                            .state
                            .reconnection_attempts
                            .load(Ordering::Relaxed),
                        last_message_timestamp: metadata
                            .state
                            .last_message_timestamp
                            .load(Ordering::Relaxed),
                        connection_start_time: metadata
                            .state
                            .connection_start_time
                            .load(Ordering::Relaxed),
                    },
                );
            }
        }

        ExchangeManagerStats {
            total_exchanges: self.total_exchanges.load(Ordering::Relaxed),
            total_active_tasks: self.total_active_tasks.load(Ordering::Relaxed),
            total_messages_received: total_messages,
            total_bytes_received: total_bytes,
            total_errors: self.total_errors.load(Ordering::Relaxed),
            messages_per_second,
            bytes_per_second,
            uptime_seconds,
            exchange_stats,
        }
    }

    /// Request shutdown (lock-free)
    pub fn request_shutdown(&self) {
        self.shutdown_requested.store(true, Ordering::Release);
    }

    /// Check if shutdown was requested (lock-free)
    pub fn is_shutdown_requested(&self) -> bool {
        self.shutdown_requested.load(Ordering::Acquire)
    }

    /// Shutdown all exchanges
    pub async fn shutdown(&self) -> Result<()> {
        self.request_shutdown();

        let exchanges = {
            let mut exchanges = self.exchanges.write();
            std::mem::take(&mut *exchanges)
        };

        for (exchange_name, metadata) in exchanges {
            // Take client out of lock to avoid holding across await
            let client_opt = {
                let mut client_lock = metadata.client.write();
                client_lock.take()
            };

            if let Some(mut client) = client_opt
                && let Err(e) = client.disconnect().await
            {
                log::error!("Error disconnecting exchange {exchange_name} during shutdown: {e}");
            }
            metadata.state.set_connected(false);
        }

        // Reset counters
        self.total_exchanges.store(0, Ordering::Relaxed);
        self.total_active_tasks.store(0, Ordering::Relaxed);

        log::info!("Exchange manager shutdown completed");
        Ok(())
    }
}

/// Statistics for the exchange manager
///
/// Provides comprehensive performance metrics and operational statistics
/// for the lock-free exchange manager, including global counters and
/// per-exchange breakdown.
#[derive(Debug, Clone)]
pub struct ExchangeManagerStats {
    /// Total number of exchanges currently registered
    pub total_exchanges: usize,
    /// Total number of active collection tasks across all exchanges
    pub total_active_tasks: usize,
    /// Total messages received from all exchanges
    pub total_messages_received: u64,
    /// Total bytes received from all exchanges
    pub total_bytes_received: u64,
    /// Total errors encountered across all exchanges
    pub total_errors: u64,
    /// Average messages per second across all exchanges
    pub messages_per_second: f64,
    /// Average bytes per second across all exchanges
    pub bytes_per_second: f64,
    /// Manager uptime in seconds since initialization
    pub uptime_seconds: f64,
    /// Per-exchange statistics breakdown
    pub exchange_stats: FxHashMap<SmartString, ExchangeStats>,
}

/// Statistics for individual exchanges
///
/// Provides detailed performance metrics and connection information
/// for a specific exchange, including message throughput, connection
/// status, and error tracking.
#[derive(Debug, Clone)]
pub struct ExchangeStats {
    /// Whether the exchange is currently connected
    pub connected: bool,
    /// Total messages received from this exchange
    pub messages_received: u64,
    /// Total bytes received from this exchange
    pub bytes_received: u64,
    /// Average messages per second for this exchange
    pub messages_per_second: f64,
    /// Number of reconnection attempts made
    pub reconnection_attempts: usize,
    /// Timestamp of the last message received (nanoseconds)
    pub last_message_timestamp: u64,
    /// Timestamp when the connection was established (nanoseconds)
    pub connection_start_time: u64,
}

impl Default for LockFreeExchangeManager {
    fn default() -> Self {
        Self::new()
    }
}

// Thread safety
unsafe impl Send for LockFreeExchangeManager {}
unsafe impl Sync for LockFreeExchangeManager {}

#[cfg(test)]
mod tests {
    use super::*;
    use std::thread;
    use std::time::Duration;

    #[tokio::test]
    async fn test_lockfree_exchange_manager_basic_operations() {
        let manager = LockFreeExchangeManager::new();

        // Test initial state
        let stats = manager.get_statistics();
        assert_eq!(stats.total_exchanges, 0);
        assert_eq!(stats.total_messages_received, 0);

        // Test shutdown
        assert!(!manager.is_shutdown_requested());
        manager.request_shutdown();
        assert!(manager.is_shutdown_requested());
    }

    #[tokio::test]
    async fn test_lockfree_exchange_manager_message_recording() {
        let manager = Arc::new(LockFreeExchangeManager::new());

        // Test concurrent message recording
        let mut handles = Vec::new();
        for i in 0..10 {
            let manager_clone = Arc::clone(&manager);
            let handle = thread::spawn(move || {
                for j in 0..100 {
                    manager_clone.record_message(&format!("exchange_{i}"), j as u64);
                }
            });
            handles.push(handle);
        }

        // Wait for all threads
        for handle in handles {
            handle.join().unwrap();
        }

        let stats = manager.get_statistics();
        assert_eq!(stats.total_messages_received, 1000); // 10 threads * 100 messages
        assert!(stats.total_bytes_received > 0);
    }

    #[tokio::test]
    async fn test_lockfree_exchange_manager_error_recording() {
        let manager = LockFreeExchangeManager::new();

        // Test error recording
        manager.record_error("test_exchange");
        manager.record_error("test_exchange");

        let stats = manager.get_statistics();
        assert_eq!(stats.total_errors, 2);
    }

    #[test]
    fn test_atomic_exchange_state() {
        let state = AtomicExchangeState::new();

        // Test initial state
        assert!(!state.is_connected());
        assert_eq!(state.reconnection_attempts.load(Ordering::Relaxed), 0);

        // Test connection state changes
        state.set_connected(true);
        assert!(state.is_connected());

        state.record_error();
        assert!(!state.is_connected());

        // Test message recording
        state.record_message(1024);
        assert_eq!(state.messages_received.load(Ordering::Relaxed), 1);
        assert_eq!(state.bytes_received.load(Ordering::Relaxed), 1024);

        // Test reconnection attempts
        state.increment_reconnection_attempts();
        state.increment_reconnection_attempts();
        assert_eq!(state.reconnection_attempts.load(Ordering::Relaxed), 2);

        state.reset_reconnection_attempts();
        assert_eq!(state.reconnection_attempts.load(Ordering::Relaxed), 0);

        // Test health check coordination
        assert!(state.try_start_health_check());
        assert!(!state.try_start_health_check()); // Should fail the second time

        state.finish_health_check();
        assert!(state.try_start_health_check()); // Should succeed again
    }

    #[tokio::test]
    async fn test_lockfree_exchange_manager_statistics() {
        let manager = LockFreeExchangeManager::new();

        // Record some activity
        manager.record_message("exchange1", 1000);
        manager.record_message("exchange2", 2000);
        manager.record_error("exchange1");

        // Wait a bit for time-based calculations
        thread::sleep(Duration::from_millis(100));

        let stats = manager.get_statistics();
        assert_eq!(stats.total_messages_received, 2);
        assert_eq!(stats.total_bytes_received, 3000);
        assert_eq!(stats.total_errors, 1);
        assert!(stats.messages_per_second > 0.0);
        assert!(stats.bytes_per_second > 0.0);
        assert!(stats.uptime_seconds > 0.0);
    }
}