//! System Worker Binary - Asynchronous worker for high-throughput concurrent processing use clap::Parser; use log::{error, info, warn}; use rhailib_worker::async_worker_impl::AsyncWorker; use rhailib_worker::config::{ConfigError, WorkerConfig}; use rhailib_worker::engine::create_heromodels_engine; use rhailib_worker::worker_trait::{spawn_worker, WorkerConfig as TraitWorkerConfig}; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; use tokio::signal; use tokio::sync::mpsc; #[derive(Parser, Debug)] #[command( name = "system", version = "0.1.0", about = "System Worker - Asynchronous Worker with Concurrent Job Processing", long_about = "An asynchronous worker for Hero framework that processes multiple jobs \ concurrently with timeout support. Ideal for high-throughput scenarios \ where jobs can be executed in parallel." )] struct Args { /// Path to TOML configuration file #[arg(short, long, help = "Path to TOML configuration file")] config: PathBuf, /// Override worker ID from config #[arg(long, help = "Override worker ID from configuration file")] worker_id: Option, /// Override Redis URL from config #[arg(long, help = "Override Redis URL from configuration file")] redis_url: Option, /// Override database path from config #[arg(long, help = "Override database path from configuration file")] db_path: Option, /// Override default timeout in seconds #[arg(long, help = "Override default job timeout in seconds")] timeout: Option, /// Enable verbose logging (debug level) #[arg(short, long, help = "Enable verbose logging")] verbose: bool, /// Disable timestamps in log output #[arg(long, help = "Remove timestamps from log output")] no_timestamp: bool, /// Show worker statistics periodically #[arg(long, help = "Show periodic worker statistics")] show_stats: bool, } #[tokio::main] async fn main() -> Result<(), Box> { let args = Args::parse(); // Load configuration from TOML file let mut config = match WorkerConfig::from_file(&args.config) { Ok(config) => config, Err(e) => { eprintln!("Failed to load configuration from {:?}: {}", args.config, e); std::process::exit(1); } }; // Validate that this is an async worker configuration if !config.is_async() { eprintln!("Error: System worker requires an async worker configuration"); eprintln!("Expected: [worker_type] type = \"async\""); eprintln!("Found: {:?}", config.worker_type); std::process::exit(1); } // Apply command line overrides if let Some(worker_id) = args.worker_id { config.worker_id = worker_id; } if let Some(redis_url) = args.redis_url { config.redis_url = redis_url; } if let Some(db_path) = args.db_path { config.db_path = db_path; } // Override timeout if specified if let Some(timeout_secs) = args.timeout { if let rhailib_worker::config::WorkerType::Async { ref mut default_timeout_seconds } = config.worker_type { *default_timeout_seconds = timeout_secs; } } // Configure logging setup_logging(&config, args.verbose, args.no_timestamp)?; info!("🚀 System Worker starting..."); info!("Worker ID: {}", config.worker_id); info!("Redis URL: {}", config.redis_url); info!("Database Path: {}", config.db_path); info!("Preserve Tasks: {}", config.preserve_tasks); if let Some(timeout) = config.get_default_timeout() { info!("Default Timeout: {:?}", timeout); } // Create Rhai engine let engine = create_heromodels_engine(); info!("✅ Rhai engine initialized"); // Create worker configuration for the trait-based interface let mut worker_config = TraitWorkerConfig::new( config.worker_id.clone(), config.db_path.clone(), config.redis_url.clone(), config.preserve_tasks, ); // Add timeout configuration for async worker if let Some(timeout) = config.get_default_timeout() { worker_config = worker_config.with_default_timeout(timeout); } // Create async worker instance let worker = Arc::new(AsyncWorker::default()); info!("✅ Async worker instance created"); // Setup shutdown signal handling let (shutdown_tx, shutdown_rx) = mpsc::channel(1); // Spawn shutdown signal handler let shutdown_tx_clone = shutdown_tx.clone(); tokio::spawn(async move { if let Err(e) = signal::ctrl_c().await { error!("Failed to listen for shutdown signal: {}", e); return; } info!("🛑 Shutdown signal received"); if let Err(e) = shutdown_tx_clone.send(()).await { error!("Failed to send shutdown signal: {}", e); } }); // Spawn statistics reporter if requested if args.show_stats { let worker_stats = Arc::clone(&worker); tokio::spawn(async move { let mut interval = tokio::time::interval(Duration::from_secs(30)); loop { interval.tick().await; let running_count = worker_stats.running_job_count().await; if running_count > 0 { info!("📊 Worker Stats: {} jobs currently running", running_count); } else { info!("📊 Worker Stats: No jobs currently running"); } } }); } // Spawn the worker info!("🔄 Starting worker loop..."); let worker_handle = spawn_worker(worker, engine, shutdown_rx); // Wait for the worker to complete match worker_handle.await { Ok(Ok(())) => { info!("✅ System Worker shut down gracefully"); } Ok(Err(e)) => { error!("❌ System Worker encountered an error: {}", e); std::process::exit(1); } Err(e) => { error!("❌ Failed to join worker task: {}", e); std::process::exit(1); } } Ok(()) } /// Setup logging based on configuration and command line arguments fn setup_logging( config: &WorkerConfig, verbose: bool, no_timestamp: bool, ) -> Result<(), Box> { let mut builder = env_logger::Builder::new(); // Determine log level let log_level = if verbose { "debug" } else { &config.logging.level }; // Set log level builder.filter_level(match log_level.to_lowercase().as_str() { "trace" => log::LevelFilter::Trace, "debug" => log::LevelFilter::Debug, "info" => log::LevelFilter::Info, "warn" => log::LevelFilter::Warn, "error" => log::LevelFilter::Error, _ => { warn!("Invalid log level: {}. Using 'info'", log_level); log::LevelFilter::Info } }); // Configure timestamps let show_timestamps = !no_timestamp && config.logging.timestamps; if !show_timestamps { builder.format_timestamp(None); } builder.init(); Ok(()) } #[cfg(test)] mod tests { use super::*; use std::io::Write; use tempfile::NamedTempFile; #[test] fn test_config_validation() { let config_toml = r#" worker_id = "test_system" redis_url = "redis://localhost:6379" db_path = "/tmp/test_db" [worker_type] type = "async" default_timeout_seconds = 600 [logging] level = "info" "#; let mut temp_file = NamedTempFile::new().unwrap(); temp_file.write_all(config_toml.as_bytes()).unwrap(); let config = WorkerConfig::from_file(temp_file.path()).unwrap(); assert!(!config.is_sync()); assert!(config.is_async()); assert_eq!(config.worker_id, "test_system"); assert_eq!(config.get_default_timeout(), Some(Duration::from_secs(600))); } #[test] fn test_sync_config_rejection() { let config_toml = r#" worker_id = "test_system" redis_url = "redis://localhost:6379" db_path = "/tmp/test_db" [worker_type] type = "sync" [logging] level = "info" "#; let mut temp_file = NamedTempFile::new().unwrap(); temp_file.write_all(config_toml.as_bytes()).unwrap(); let config = WorkerConfig::from_file(temp_file.path()).unwrap(); assert!(config.is_sync()); assert!(!config.is_async()); // This would be rejected in main() function } #[test] fn test_timeout_override() { let config_toml = r#" worker_id = "test_system" redis_url = "redis://localhost:6379" db_path = "/tmp/test_db" [worker_type] type = "async" default_timeout_seconds = 300 "#; let mut temp_file = NamedTempFile::new().unwrap(); temp_file.write_all(config_toml.as_bytes()).unwrap(); let mut config = WorkerConfig::from_file(temp_file.path()).unwrap(); assert_eq!(config.get_default_timeout(), Some(Duration::from_secs(300))); // Test timeout override if let rhailib_worker::config::WorkerType::Async { ref mut default_timeout_seconds } = config.worker_type { *default_timeout_seconds = 600; } assert_eq!(config.get_default_timeout(), Some(Duration::from_secs(600))); } }