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> { // 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 = 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(()) }