rusty_bin/monitor/collector/

manager.rs

//! Collection manager for coordinating data collection from multiple exchanges
//!
//! This module manages the collection of market data from various exchanges and symbols
//! using the ExchangeClient infrastructure.

use crate::monitor::collector::{
    CollectionStatus, CollectionTask, DataType, MarketDataEvent, Result,
    exchange_client::ExchangeClient,
};
use crate::monitor::config::MonitorConfig;
use flume::Sender;
use futures::future;
use quanta::Clock;
use rusty_common::collections::FxHashMap;
use smartstring::alias::String as SmartString;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::{RwLock, oneshot};
use tokio::task::JoinHandle;

/// Retry state for failed exchange connections
#[derive(Debug, Clone)]
struct RetryState {
    failure_time: Instant,
    retry_count: u32,
    next_retry_time: Instant,
    last_error: String,
}

impl RetryState {
    fn new(error: String) -> Self {
        let now = Instant::now();
        Self {
            failure_time: now,
            retry_count: 0,
            next_retry_time: now + Duration::from_secs(5), // Initial retry after 5 seconds
            last_error: error,
        }
    }

    fn calculate_next_retry(&mut self) {
        self.retry_count += 1;
        // Exponential backoff with jitter: 5s, 10s, 20s, 40s, 80s, 160s, max 300s (5 min)
        let base_delay =
            Duration::from_secs(5).min(Duration::from_secs(5 * (1 << self.retry_count)));
        let max_delay = Duration::from_secs(300);
        let delay = base_delay.min(max_delay);

        // Add jitter (±20%) to avoid thundering herd
        // Add jitter factor between 0.8 and 1.2
        let jitter_factor = 0.8 + (rand::random::<f64>() * 0.4);
        let final_delay = delay.mul_f64(jitter_factor);

        self.next_retry_time = Instant::now() + final_delay;
    }

    fn should_retry(&self) -> bool {
        // Give up after 10 retries
        self.retry_count < 10 && Instant::now() >= self.next_retry_time
    }
}

/// Collection manager that coordinates data collection from multiple exchanges
pub struct CollectionManager {
    _config: MonitorConfig,
    exchange_clients: Arc<RwLock<FxHashMap<SmartString, ExchangeClient>>>,
    active_tasks: Arc<RwLock<FxHashMap<String, CollectionTask>>>,
    collection_handles: Arc<RwLock<Vec<JoinHandle<()>>>>,
    shutdown_tx: Arc<RwLock<Option<oneshot::Sender<()>>>>,
    failed_exchanges: Arc<RwLock<FxHashMap<SmartString, RetryState>>>,
}

impl CollectionManager {
    /// Create a new collection manager
    pub async fn new(config: MonitorConfig) -> Result<Self> {
        Ok(Self {
            _config: config,
            exchange_clients: Arc::new(RwLock::new(FxHashMap::default())),
            active_tasks: Arc::new(RwLock::new(FxHashMap::default())),
            collection_handles: Arc::new(RwLock::new(Vec::new())),
            shutdown_tx: Arc::new(RwLock::new(None)),
            failed_exchanges: Arc::new(RwLock::new(FxHashMap::default())),
        })
    }

    /// Start data collection for specified exchanges and symbols
    pub async fn start_collection(
        &self,
        exchange_filter: Option<Vec<SmartString>>,
        symbol_filter: Option<Vec<SmartString>>,
        event_sender: Sender<MarketDataEvent>,
    ) -> Result<()> {
        log::info!(
            "Starting data collection with filters: exchanges={exchange_filter:?}, symbols={symbol_filter:?}"
        );

        // Determine which exchanges to use
        let exchanges = match exchange_filter {
            Some(exchanges) => exchanges,
            None => vec![
                "binance".into(),
                // "upbit".into(),
                // "bybit".into(),
                // "coinbase".into(),
                // "bithumb".into(),
            ],
        };

        // Determine which symbols to collect (default popular symbols if none specified)
        let symbols = match symbol_filter {
            Some(symbols) => symbols,
            None => vec![
                "BTCUSDT".into(),  // Binance Futures perpetual
                "ETHUSDT".into(),  // Binance Futures perpetual
                "ADAUSDT".into(),  // Binance Futures perpetual
                "SOLUSDT".into(),  // Binance Futures perpetual
                "AVAXUSDT".into(), // Binance Futures perpetual
                "DOTUSDT".into(),  // Binance Futures perpetual
                "KRW-BTC".into(),  // Upbit format
                "BTC-USD".into(),  // Coinbase format
                "BTC_KRW".into(),  // Bithumb format
            ],
        };

        // Initialize exchange clients
        for exchange_name in exchanges {
            let mut client = ExchangeClient::new(exchange_name.clone());

            match client.connect().await {
                Ok(()) => {
                    log::info!("Successfully connected to {exchange_name}");
                    // Remove from failed exchanges if it was there
                    self.failed_exchanges.write().await.remove(&exchange_name);
                }
                Err(e) => {
                    log::error!("Failed to connect to exchange {exchange_name}: {e}");
                    // Add to failed exchanges for retry
                    self.failed_exchanges
                        .write()
                        .await
                        .insert(exchange_name.clone(), RetryState::new(e.to_string()));
                    continue; // Skip this exchange for now
                }
            }

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

            // Filter symbols appropriate for this exchange
            let exchange_symbols = self.filter_symbols_for_exchange(&exchange_name, &symbols);

            if exchange_symbols.is_empty() {
                log::warn!("No symbols to collect for exchange: {exchange_name}");
                let _ = client.disconnect().await;
                continue;
            }

            // Start collection for both trades and orderbook
            let data_types = vec![DataType::Trades, DataType::OrderBook];

            client
                .start_collection(
                    exchange_symbols.clone(),
                    data_types.clone(),
                    event_sender.clone(),
                )
                .await?;
            log::info!(
                "Started collection for exchange {exchange_name} with symbols: {exchange_symbols:?}"
            );

            // Add tasks to tracking
            let mut active_tasks = self.active_tasks.write().await;

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

            // Store the client
            self.exchange_clients
                .write()
                .await
                .insert(exchange_name, client);
        }

        log::info!("All exchange connections initialized. Starting management loop...");

        // Management loop - keep the collection manager running
        let (tx, mut rx) = oneshot::channel();
        *self.shutdown_tx.write().await = Some(tx);
        let mut health_check_interval = tokio::time::interval(std::time::Duration::from_secs(30));
        let mut retry_interval = tokio::time::interval(std::time::Duration::from_secs(60));

        loop {
            tokio::select! {
                _ = &mut rx => {
                    log::info!("CollectionManager received shutdown signal. Exiting management loop.");
                    break;
                }
                _ = health_check_interval.tick() => {
                    // Perform health checks on all connections
                    let health_status = self.health_check().await;

                    health_status
                        .into_iter()
                        .filter(|(_, healthy)| !healthy)
                        .for_each(|(exchange, _)| {
                            log::warn!("Exchange {exchange} connection unhealthy");
                            // Mark unhealthy exchanges for retry
                            let failed_exchanges = self.failed_exchanges.clone();
                            let exchange = exchange.clone();
                            tokio::spawn(async move {
                                let mut failed = failed_exchanges.write().await;
                                failed.entry(exchange.clone())
                                    .or_insert_with(|| RetryState::new("Health check failed".to_string()));
                            });
                        });
                }
                _ = retry_interval.tick() => {
                    // Retry failed connections
                    self.retry_failed_exchanges(&event_sender).await;
                }
            }
        }
        log::info!("Collection management loop has stopped.");
        Ok(())
    }

    /// Stop all collection tasks
    pub async fn stop_all(&self) -> Result<()> {
        log::info!("Stopping all collection tasks");

        if let Some(tx) = self.shutdown_tx.write().await.take()
            && tx.send(()).is_err()
        {
            log::warn!(
                "Failed to send shutdown signal to management loop; it might have already stopped."
            );
        }

        // Stop all collection handles
        let mut handles = self.collection_handles.write().await;
        for handle in handles.drain(..) {
            handle.abort();
        }

        // Disconnect all exchange clients
        let mut clients = self.exchange_clients.write().await;
        for (exchange_name, mut client) in clients.drain() {
            if let Err(e) = client.disconnect().await {
                log::error!("Failed to disconnect from exchange {exchange_name}: {e}");
            }
        }

        // Clear active tasks
        self.active_tasks.write().await.clear();

        log::info!("All collection tasks stopped");

        Ok(())
    }

    /// Get active tasks
    pub async fn get_active_tasks(&self) -> FxHashMap<String, CollectionTask> {
        self.active_tasks.read().await.clone()
    }

    /// Get collection status for all exchanges
    pub async fn get_collection_status(&self) -> FxHashMap<SmartString, Vec<CollectionStatus>> {
        let mut status_map = FxHashMap::default();
        let clients = self.exchange_clients.read().await;

        for (exchange_name, client) in clients.iter() {
            if let Some(status) = client.get_status().await {
                status_map
                    .entry(exchange_name.clone())
                    .or_insert_with(Vec::new)
                    .push(status);
            }
        }

        status_map
    }

    /// Stop collection for a specific exchange
    pub async fn stop_exchange(&self, exchange_name: &str) -> Result<()> {
        log::info!("Stopping collection for exchange: {exchange_name}");

        let exchange_name_smart: SmartString = exchange_name.into();
        let mut clients = self.exchange_clients.write().await;

        if let Some(mut client) = clients.remove(&exchange_name_smart) {
            client.disconnect().await?;

            // Remove related tasks
            let mut active_tasks = self.active_tasks.write().await;
            active_tasks.retain(|task_id, _| !task_id.starts_with(exchange_name));

            log::info!("Successfully stopped collection for exchange: {exchange_name}");
        } else {
            log::warn!("Exchange {exchange_name} was not active");
        }

        Ok(())
    }

    /// Filter symbols to only include those appropriate for the given exchange
    fn filter_symbols_for_exchange(
        &self,
        exchange_name: &str,
        symbols: &[SmartString],
    ) -> Vec<SmartString> {
        symbols
            .iter()
            .filter(|symbol| self.is_symbol_valid_for_exchange(exchange_name, symbol))
            .cloned()
            .collect()
    }

    /// Check if a symbol is valid for a specific exchange based on naming conventions
    fn is_symbol_valid_for_exchange(&self, exchange_name: &str, symbol: &str) -> bool {
        match exchange_name.to_lowercase().as_str() {
            "binance" => {
                // Binance uses format like BTCUSDT
                symbol.len() >= 6
                    && symbol.chars().all(|c| c.is_ascii_uppercase())
                    && !symbol.contains('-')
                    && !symbol.contains('_')
            }
            "upbit" => {
                // Upbit uses format like KRW-BTC
                symbol.contains('-') && symbol.starts_with("KRW-")
            }
            "bybit" => {
                // Bybit uses format like BTCUSDT (similar to Binance)
                symbol.len() >= 6
                    && symbol.chars().all(|c| c.is_ascii_uppercase())
                    && !symbol.contains('-')
                    && !symbol.contains('_')
            }
            "coinbase" => {
                // Coinbase uses format like BTC-USD
                symbol.contains('-') && !symbol.starts_with("KRW-")
            }
            "bithumb" => {
                // Bithumb uses format like BTC_KRW
                symbol.contains('_')
            }
            _ => false,
        }
    }

    /// Get statistics for all active collections
    pub async fn get_statistics(&self) -> FxHashMap<SmartString, u32> {
        let mut stats = FxHashMap::default();
        let clients = self.exchange_clients.read().await;

        for (exchange_name, _) in clients.iter() {
            let active_tasks = self.active_tasks.read().await;
            let count = active_tasks
                .values()
                .filter(|task| task.exchange == *exchange_name)
                .count() as u32;
            stats.insert(exchange_name.clone(), count);
        }

        stats
    }

    /// Health check for all exchange connections
    /// Retry failed exchange connections
    async fn retry_failed_exchanges(&self, event_sender: &Sender<MarketDataEvent>) {
        let failed_exchanges = self.failed_exchanges.read().await.clone();

        for (exchange_name, mut retry_state) in failed_exchanges {
            if !retry_state.should_retry() {
                if retry_state.retry_count >= 10 {
                    log::error!(
                        "Exchange {} has failed {} times. Giving up. Last error: {}",
                        exchange_name,
                        retry_state.retry_count,
                        retry_state.last_error
                    );
                    // Remove from retry list
                    self.failed_exchanges.write().await.remove(&exchange_name);
                }
                continue;
            }

            log::info!(
                "Retrying connection to {} (attempt {} after {}s)",
                exchange_name,
                retry_state.retry_count + 1,
                retry_state
                    .next_retry_time
                    .duration_since(retry_state.failure_time)
                    .as_secs()
            );

            // Create a new client and try to connect
            let mut client = ExchangeClient::new(exchange_name.clone());

            match client.connect().await {
                Ok(()) => {
                    log::info!("Successfully reconnected to {exchange_name}");

                    // Remove from failed list
                    self.failed_exchanges.write().await.remove(&exchange_name);

                    // Store the reconnected client
                    self.exchange_clients
                        .write()
                        .await
                        .insert(exchange_name.clone(), client);

                    // TODO: Re-subscribe to previously active symbols for this exchange
                    // For now, we'll just log that reconnection succeeded

                    // Send a connection status event
                    let _ = event_sender.send(MarketDataEvent::ConnectionStatus {
                        exchange: exchange_name.clone(),
                        symbol: "".into(), // No specific symbol for overall connection
                        connected: true,
                        timestamp: Clock::new().raw(),
                    });
                }
                Err(e) => {
                    log::error!(
                        "Failed to reconnect to {} (attempt {}): {}",
                        exchange_name,
                        retry_state.retry_count + 1,
                        e
                    );

                    // Update retry state
                    retry_state.last_error = e.to_string();
                    retry_state.calculate_next_retry();

                    // Update in the map
                    self.failed_exchanges
                        .write()
                        .await
                        .insert(exchange_name, retry_state);
                }
            }
        }
    }

    /// Perform health check on all registered exchanges
    ///
    /// Returns a map of exchange names to their health status (true if connected and healthy, false otherwise)
    pub async fn health_check(&self) -> FxHashMap<SmartString, bool> {
        let clients = self.exchange_clients.read().await;

        let health_futures = clients.iter().map(|(name, client)| async move {
            let is_healthy = client
                .get_status()
                .await
                .map(|status| status.connected)
                .unwrap_or(false);
            (name.clone(), is_healthy)
        });

        let health_results = future::join_all(health_futures).await;

        health_results.into_iter().collect()
    }
}