refactor wip
This commit is contained in:
Timur Gordon
233
core/worker/cmd/osis.rs
Normal file
233
core/worker/cmd/osis.rs
Normal file
@@ -0,0 +1,233 @@
|
||||
//! OSIS Worker Binary - Synchronous worker for system-level operations
|
||||
|
||||
use clap::Parser;
|
||||
use log::{error, info};
|
||||
use rhailib_worker::config::{ConfigError, WorkerConfig};
|
||||
use rhailib_worker::engine::create_heromodels_engine;
|
||||
use rhailib_worker::sync_worker::SyncWorker;
|
||||
use rhailib_worker::worker_trait::{spawn_worker, WorkerConfig as TraitWorkerConfig};
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use tokio::signal;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(
|
||||
name = "osis",
|
||||
version = "0.1.0",
|
||||
about = "OSIS (Operating System Integration Service) - Synchronous Worker",
|
||||
long_about = "A synchronous worker for Hero framework that processes jobs sequentially. \
|
||||
Ideal for system-level operations that require careful resource management."
|
||||
)]
|
||||
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>,
|
||||
|
||||
/// 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,
|
||||
}
|
||||
|
||||
#[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 a sync worker configuration
|
||||
if !config.is_sync() {
|
||||
eprintln!("Error: OSIS worker requires a sync worker configuration");
|
||||
eprintln!("Expected: [worker_type] type = \"sync\"");
|
||||
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;
|
||||
}
|
||||
|
||||
// Configure logging
|
||||
setup_logging(&config, args.verbose, args.no_timestamp)?;
|
||||
|
||||
info!("🚀 OSIS 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);
|
||||
|
||||
// Create Rhai engine
|
||||
let engine = create_heromodels_engine();
|
||||
info!("✅ Rhai engine initialized");
|
||||
|
||||
// Create worker configuration for the trait-based interface
|
||||
let worker_config = TraitWorkerConfig::new(
|
||||
config.worker_id.clone(),
|
||||
config.db_path.clone(),
|
||||
config.redis_url.clone(),
|
||||
config.preserve_tasks,
|
||||
);
|
||||
|
||||
// Create sync worker instance
|
||||
let worker = Arc::new(SyncWorker::default());
|
||||
info!("✅ Sync 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 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!("✅ OSIS Worker shut down gracefully");
|
||||
}
|
||||
Ok(Err(e)) => {
|
||||
error!("❌ OSIS 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,
|
||||
_ => {
|
||||
eprintln!("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_osis"
|
||||
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());
|
||||
assert_eq!(config.worker_id, "test_osis");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_async_config_rejection() {
|
||||
let config_toml = r#"
|
||||
worker_id = "test_osis"
|
||||
redis_url = "redis://localhost:6379"
|
||||
db_path = "/tmp/test_db"
|
||||
|
||||
[worker_type]
|
||||
type = "async"
|
||||
default_timeout_seconds = 300
|
||||
|
||||
[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
|
||||
}
|
||||
}
|
||||
302
core/worker/cmd/system.rs
Normal file
302
core/worker/cmd/system.rs
Normal file
@@ -0,0 +1,302 @@
|
||||
//! 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)));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user