baobab/cmd/main.rs

use std::env;
use std::sync::Arc;
use std::time::Duration;
use hero_supervisor::{SupervisorBuilder, SupervisorError};
use hero_websocket_server::ServerBuilder;
use tokio::signal;
use log::{info, error};
use env_logger::Builder;

/// The main entry point of the Hero Supervisor.
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Initialize logging
    env_logger::Builder::from_default_env()
        .filter_level(log::LevelFilter::Info)
        .init();
    
    info!("Hero Supervisor starting up...");
    
    // Get config path from command line arguments or use default
    let args: Vec<String> = env::args().collect();
    let config_path = if let Some(config_index) = args.iter().position(|arg| arg == "--config") {
        if config_index + 1 < args.len() {
            &args[config_index + 1]
        } else {
            "cmd/config.toml"
        }
    } else {
        "cmd/config.toml"
    };
    
    println!("Loading configuration from: {}", config_path);
    
    let supervisor = SupervisorBuilder::from_toml(config_path)?
        .build().await?;
    
    // Wrap supervisor in Arc for sharing across tasks
    let supervisor = Arc::new(supervisor);

    // Extract actor configurations from TOML config
    let actor_configs = supervisor.get_actor_configs()?;
    info!("Loaded {} actor configurations from TOML", actor_configs.len());
    
    // Spawn the background lifecycle manager with 5-minute health check interval
    let health_check_interval = Duration::from_secs(5 * 60); // 5 minutes
    let mut lifecycle_handle = supervisor.clone().spawn_lifecycle_manager(actor_configs, health_check_interval);
    
    info!("Hero Supervisor started successfully!");
    info!("Background lifecycle manager is running with 5-minute health checks.");
    info!("Actors are being monitored and will be automatically restarted if they fail.");
    
    // Start WebSocket server for job dispatching
    info!("Starting WebSocket server for job dispatching...");
    let ws_supervisor = supervisor.clone();
    
    // Get WebSocket server config from TOML or use defaults
    let ws_config = supervisor.get_websocket_config().unwrap_or_else(|_| {
        info!("Using default WebSocket server configuration");
        hero_supervisor::WebSocketServerConfig {
            host: "127.0.0.1".to_string(),
            port: 8443,
            redis_url: "redis://127.0.0.1/".to_string(),
            auth: false,
            tls: false,
            cert: None,
            key: None,
            tls_port: None,
            circles: std::collections::HashMap::new(),
        }
    });
    
    let mut websocket_handle = tokio::spawn(async move {
        info!("WebSocket server starting on {}:{}", ws_config.host, ws_config.port);
        
        // Create the WebSocket server with our supervisor
        let mut server_builder = ServerBuilder::new()
            .host(&ws_config.host)
            .port(ws_config.port)
            .redis_url(&ws_config.redis_url)
            .with_supervisor(ws_supervisor);
            
        // Configure auth if enabled
        if ws_config.auth {
            server_builder = server_builder.with_auth();
        }
        
        // Configure TLS if enabled
        if ws_config.tls {
            if let (Some(cert), Some(key)) = (&ws_config.cert, &ws_config.key) {
                server_builder = server_builder.with_tls(cert.clone(), key.clone());
                if let Some(tls_port) = ws_config.tls_port {
                    server_builder = server_builder.with_tls_port(tls_port);
                }
            }
        }
        
        // Configure circles
        if !ws_config.circles.is_empty() {
            server_builder = server_builder.circles(ws_config.circles.clone());
        }
        
        let server = match server_builder.build() {
                Ok(server) => server,
                Err(e) => {
                    error!("Failed to build WebSocket server: {}", e);
                    return;
                }
            };
        
        // Start the WebSocket server
        match server.spawn_circle_server() {
            Ok((server_handle, _)) => {
                info!("WebSocket server successfully started and ready to dispatch jobs");
                if let Err(e) = server_handle.await {
                    error!("WebSocket server error: {:?}", e);
                }
            }
            Err(e) => {
                error!("Failed to start WebSocket server: {}", e);
            }
        }
    });
    
    info!("WebSocket server started - ready to accept connections and dispatch jobs");
    
    // Set up graceful shutdown signal handlers
    let shutdown_signal = async {
        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 => {},
            _ = terminate => {},
        }
        
        info!("Shutdown signal received, initiating graceful shutdown...");
    };
    
    // Wait for shutdown signal or task completion
    tokio::select! {
        _ = shutdown_signal => {
            info!("Graceful shutdown initiated");
            
            // Cancel background tasks
            lifecycle_handle.abort();
            websocket_handle.abort();
            
            info!("Background tasks stopped");
        }
        result = &mut lifecycle_handle => {
            match result {
                Ok(Ok(())) => info!("Lifecycle manager completed successfully"),
                Ok(Err(e)) => error!("Lifecycle manager error: {}", e),
                Err(e) => error!("Lifecycle manager task panicked: {}", e),
            }
            // Also stop the websocket handle
            websocket_handle.abort();
        }
        result = &mut websocket_handle => {
            match result {
                Ok(()) => info!("WebSocket server completed successfully"),
                Err(e) => error!("WebSocket server task panicked: {}", e),
            }
            // Also stop the lifecycle handle
            lifecycle_handle.abort();
        }
    }
    
    info!("Hero Supervisor shutdown complete");
    
    Ok(())
}