rusty_monitor/

main.rs

//! Rusty Monitor - Real-time Market Data Collection System
//!
//! A high-performance market data collection system that gathers real-time
//! trade and orderbook data from multiple cryptocurrency exchanges and stores
//! them in `FlatBuffers` format with compression.
use mimalloc::MiMalloc;

#[global_allocator]
static GLOBAL: MiMalloc = MiMalloc;

mod cli;

use anyhow::{Context, Result};
use clap::Parser;
use cli::{Cli, Commands, ConfigArgs, ConfigCommands, StartArgs};
use flume::{Receiver, Sender};
use rusty_bin::monitor::system_monitor::SystemMonitor;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::signal;
use tokio::task::JoinHandle;

use rusty_bin::monitor::collector::{
    CollectionManager, DataPipeline, MarketDataEvent, PipelineConfig,
};
use rusty_bin::monitor::config::MonitorConfig;
use rusty_bin::monitor::config::storage::StorageConfig;
use rusty_bin::monitor::storage::{CompressionMode, manager::StorageManager};

struct MonitorServices {
    collection_manager: Arc<CollectionManager>,
    data_pipeline: DataPipeline,
    system_monitor: Arc<SystemMonitor>,
    event_sender: Sender<MarketDataEvent>,
    event_receiver: Receiver<MarketDataEvent>,
    config: MonitorConfig,
}

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize logging early
    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info"))
        .format_timestamp_nanos()
        .init();

    // Parse command line arguments using clap's derive API
    let cli = Cli::parse();

    // Handle different commands
    // If no subcommand is specified, default to the `start` command.
    match cli
        .command
        .as_ref()
        .unwrap_or(&Commands::Start(StartArgs {}))
    {
        Commands::Start(_) => start_monitor(&cli).await,
        Commands::Config(config_args) => handle_config_command(config_args, &cli.config).await,
        Commands::Status(_) => handle_status_command(&cli).await,
    }
}

/// Start the market data monitor
async fn start_monitor(cli: &Cli) -> Result<()> {
    log::info!("🚀 Starting Rusty Market Data Monitor");

    let mut services = initialize_services(cli).await?;
    run_and_shutdown_services(&mut services, cli).await?;

    log::info!("✅ Rusty Market Data Monitor shut down gracefully.");
    Ok(())
}

async fn initialize_services(cli: &Cli) -> Result<MonitorServices> {
    // 1. Load configuration
    let config = match load_config(&cli.config) {
        Ok(config) => config,
        Err(e) => {
            log::error!("Failed to load configuration: {e}");
            log::error!(
                "If you are running for the first time, please create default configuration file with `config generate`"
            );
            return Err(e);
        }
    };

    // 2. Validate and setup
    config
        .validate()
        .context("Configuration validation failed")?;
    config
        .create_directories()
        .context("Failed to create directories")?;

    // 3. Create component configurations
    let storage_config = StorageConfig {
        market_data_path: config.storage.market_data_path.clone(),
        rotation_check_interval_s: 1,
        max_file_size_mb: 1,
        max_records_per_file: 1000,
        ..Default::default()
    };
    let compression_mode = if config.storage.enable_compression {
        CompressionMode::Realtime
    } else {
        CompressionMode::None
    };
    let pipeline_config = PipelineConfig {
        batch_size: 100,
        batch_timeout_ms: 100,
        worker_count: 4,
        compression_mode,
        rotation_check_interval_seconds: 5,
    };

    // 4. Create communication channel
    let (event_sender, event_receiver) = flume::unbounded();

    // 5. Initialize managers and services
    log::info!("📊 Initializing components...");
    let collection_manager = Arc::new(CollectionManager::new(config.clone()).await?);
    let storage_manager = Arc::new(StorageManager::new(storage_config).await?);
    let data_pipeline = DataPipeline::new(pipeline_config, storage_manager);
    let system_monitor = Arc::new(SystemMonitor::new(config.clone()));
    log::info!("✅ Components initialized.");

    Ok(MonitorServices {
        collection_manager,
        data_pipeline,
        system_monitor,
        event_sender,
        event_receiver,
        config,
    })
}

async fn run_and_shutdown_services(services: &mut MonitorServices, cli: &Cli) -> Result<()> {
    // Extract command line filters
    let exchange_filter = cli.exchanges.clone().map(|exchanges| {
        exchanges
            .into_iter()
            .map(std::convert::Into::into)
            .collect()
    });
    let symbol_filter = cli
        .symbols
        .clone()
        .map(|symbols| symbols.into_iter().map(std::convert::Into::into).collect());

    // Start background tasks
    log::info!("🔄 Starting background services...");
    services
        .data_pipeline
        .start(services.event_receiver.clone())
        .await?;

    let collection_handle = {
        let manager = services.collection_manager.clone();
        let sender = services.event_sender.clone();
        tokio::spawn(async move {
            if let Err(e) = manager
                .start_collection(exchange_filter, symbol_filter, sender)
                .await
            {
                log::error!("Collection manager failed: {e}");
            } else {
                log::info!("Collection manager completed successfully");
            }
        })
    };

    log::info!("Starting health monitoring...");
    let health_handle = services.system_monitor.start_health_monitoring().await;

    let mut metrics_handle: Option<JoinHandle<()>> = None;
    if services.config.monitoring.enable_metrics {
        log::info!("Starting metrics monitoring...");
        metrics_handle = Some(services.system_monitor.start_monitoring().await);
        log::info!("✅ Monitoring services started.");
    }

    // Wait for shutdown signal (Ctrl+C)
    wait_for_shutdown_signal().await;

    // Graceful shutdown
    log::info!("🔌 Shutting down services...");
    log::info!("Stopping collection manager...");
    services.collection_manager.stop_all().await?;
    collection_handle.await?;
    log::info!("Collection manager stopped.");

    if services.config.monitoring.enable_metrics
        && let Some(handle) = metrics_handle
    {
        log::info!("Stopping metrics monitoring...");
        services.system_monitor.stop_monitoring().await;
        handle.await?;
        log::info!("Metrics monitoring stopped.");
    }

    log::info!("Stopping health monitoring...");
    health_handle.abort();
    log::info!("Health monitoring stopped.");

    log::info!("Stopping data pipeline...");
    services.data_pipeline.stop().await?;
    log::info!("Data pipeline stopped.");

    Ok(())
}

/// Load configuration from file
fn load_config(path: &PathBuf) -> Result<MonitorConfig> {
    log::info!("Loading configuration from: {path:?}");
    let config =
        MonitorConfig::from_file(path).context(format!("Failed to load config file: {path:?}"))?;
    Ok(config)
}

/// Handle the `config` subcommand
async fn handle_config_command(args: &ConfigArgs, config_path: &PathBuf) -> Result<()> {
    match &args.command {
        ConfigCommands::Validate => {
            log::info!("Validating configuration file: {config_path:?}");
            match load_config(config_path) {
                Ok(config) => match config.validate() {
                    Ok(()) => log::info!("✅ Configuration is valid."),
                    Err(e) => log::error!("❌ Configuration validation failed: {e}"),
                },
                Err(e) => log::error!("❌ Failed to load configuration: {e}"),
            }
        }
        ConfigCommands::Generate { output } => {
            log::info!("Generating default configuration file at: {output:?}");
            let config = MonitorConfig::default();
            let toml_content = toml::to_string_pretty(&config)?;

            // create directories if they don't exist
            let dir = Path::new(output)
                .parent()
                .ok_or(anyhow::anyhow!("Invalid output path"))?;
            if !dir.exists() {
                tokio::fs::create_dir_all(dir).await?;
            }

            tokio::fs::write(output, toml_content).await?;
            log::info!("✅ Default configuration file generated successfully.");
        }
    }
    Ok(())
}

/// Handle the `status` subcommand
async fn handle_status_command(cli: &Cli) -> Result<()> {
    log::info!("Showing system status...");

    let config = load_config(&cli.config)?;
    let collection_manager = Arc::new(CollectionManager::new(config.clone()).await?);
    let system_monitor = Arc::new(SystemMonitor::new(config.clone()));

    let status = collection_manager.get_collection_status().await;
    let stats = collection_manager.get_statistics().await;
    let system_metrics = system_monitor.get_system_metrics().await;
    let health_checks = system_monitor.get_health_checks().await;

    println!("--- System Status ---");
    display_system_status(&system_metrics, &health_checks);

    println!("\n--- Collection Status ---");
    display_collection_status(&status);

    println!("\n--- Collection Statistics ---");
    display_collection_statistics(&stats);

    Ok(())
}

/// Display system status in a user-friendly format
fn display_system_status(
    metrics: &rusty_bin::monitor::SystemMetrics,
    health_checks: &[rusty_bin::monitor::HealthCheck],
) {
    // System uptime
    let uptime_seconds = metrics.uptime;
    let uptime_hours = uptime_seconds / 3600;
    let uptime_minutes = (uptime_seconds % 3600) / 60;
    let uptime_secs = uptime_seconds % 60;

    println!("  Uptime: {uptime_hours}h {uptime_minutes}m {uptime_secs}s");

    // CPU usage
    println!("  CPU Usage: {:.1}%", metrics.cpu_usage);

    // Memory usage
    let memory_gb = metrics.memory_total as f64 / (1024.0 * 1024.0 * 1024.0);
    let memory_used_gb = metrics.memory_used as f64 / (1024.0 * 1024.0 * 1024.0);
    println!(
        "  Memory: {:.2}GB / {:.2}GB ({:.1}%)",
        memory_used_gb, memory_gb, metrics.memory_usage
    );

    // Disk usage
    if !metrics.disk_usage.is_empty() {
        println!("  Disk Usage:");
        for (mount, disk) in &metrics.disk_usage {
            let total_gb = disk.total_space as f64 / (1024.0 * 1024.0 * 1024.0);
            let used_gb = disk.used_space as f64 / (1024.0 * 1024.0 * 1024.0);
            println!(
                "    {}: {:.2}GB / {:.2}GB ({:.1}%)",
                mount, used_gb, total_gb, disk.usage_percent
            );
        }
    }

    // Network stats
    let net_rx_mb = metrics.network_stats.bytes_received as f64 / (1024.0 * 1024.0);
    let net_tx_mb = metrics.network_stats.bytes_transmitted as f64 / (1024.0 * 1024.0);
    println!("  Network: RX {net_rx_mb:.2}MB, TX {net_tx_mb:.2}MB");

    // Health status
    println!("  Health Status:");
    for check in health_checks {
        let status_icon = match check.status {
            rusty_bin::monitor::HealthStatus::Healthy => "✅",
            rusty_bin::monitor::HealthStatus::Warning => "⚠️",
            rusty_bin::monitor::HealthStatus::Critical => "❌",
            rusty_bin::monitor::HealthStatus::Unknown => "❓",
        };
        println!("    {} {}: {}", status_icon, check.name, check.message);
    }
}

/// Display collection status in a user-friendly format
fn display_collection_status(
    status: &rusty_common::collections::FxHashMap<
        smartstring::alias::String,
        Vec<rusty_bin::monitor::collector::CollectionStatus>,
    >,
) {
    if status.is_empty() {
        println!("  No active collections");
        return;
    }

    for (exchange, statuses) in status {
        println!("  Exchange: {exchange}");
        for status in statuses {
            let status_icon = if status.connected { "✅" } else { "❌" };
            let data_types: Vec<String> = status
                .data_types
                .iter()
                .map(std::string::ToString::to_string)
                .collect();
            println!(
                "    {} {} [{}] - {}",
                status_icon,
                status.symbol,
                data_types.join(", "),
                if status.connected {
                    "connected"
                } else {
                    "disconnected"
                }
            );
        }
    }
}

/// Display collection statistics in a user-friendly format
fn display_collection_statistics(
    stats: &rusty_common::collections::FxHashMap<smartstring::alias::String, u32>,
) {
    if stats.is_empty() {
        println!("  No statistics available");
        return;
    }

    let total_tasks: u32 = stats.values().sum();
    println!("  Total active tasks: {total_tasks}");

    for (exchange, count) in stats {
        println!("  {exchange}: {count} tasks");
    }
}

/// Wait for a shutdown signal (e.g., Ctrl+C)
async fn wait_for_shutdown_signal() {
    let ctrl_c = async {
        signal::ctrl_c()
            .await
            .expect("Failed to install Ctrl+C handler");
    };

    #[cfg(unix)]
    let terminate = async {
        signal::unix::signal(signal::unix::SignalKind::terminate())
            .expect("Failed to install signal handler")
            .recv()
            .await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        () = ctrl_c => {log::info!("Received Ctrl+C, initiating shutdown...")},
        () = terminate => {log::info!("Received terminate signal, initiating shutdown...")},
    }
}