baobab/core/worker/cmd/system.rs

//! 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<String>,

    /// Override Redis URL from config
    #[arg(long, help = "Override Redis URL from configuration file")]
    redis_url: Option<String>,

    /// Override database path from config
    #[arg(long, help = "Override database path from configuration file")]
    db_path: Option<String>,

    /// Override default timeout in seconds
    #[arg(long, help = "Override default job timeout in seconds")]
    timeout: Option<u64>,

    /// 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<dyn std::error::Error + Send + Sync>> {
    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<dyn std::error::Error + Send + Sync>> {
    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)));
    }
}