wip

core/supervisor/src/lib.rs

@@ -0,0 +1,596 @@
+use log::{debug, error, info, warn};
+use redis::AsyncCommands;
+use std::collections::HashMap;
+use std::time::Duration;
+use hero_job::NAMESPACE_PREFIX;
+use zinit_client::ZinitClient;
+
+mod job;
+mod error;
+mod lifecycle;
+
+pub use crate::error::SupervisorError;
+pub use crate::job::JobBuilder;
+pub use crate::lifecycle::WorkerConfig;
+// Re-export types from hero_job for public API
+pub use hero_job::{Job, JobStatus, ScriptType};
+
+pub struct Supervisor {
+    redis_client: redis::Client,
+    zinit_client: ZinitClient,
+    builder_data: Option<SupervisorBuilderData>,
+}
+
+pub struct SupervisorBuilder {
+    redis_url: Option<String>,
+    zinit_socket_path: Option<String>,
+    osis_worker: Option<String>,
+    sal_worker: Option<String>,
+    v_worker: Option<String>,
+    python_worker: Option<String>,
+    worker_env_vars: HashMap<String, String>,
+}
+
+/// Helper struct to pass builder data to worker launch method
+struct SupervisorBuilderData {
+    osis_worker: Option<String>,
+    sal_worker: Option<String>,
+    v_worker: Option<String>,
+    python_worker: Option<String>,
+    worker_env_vars: HashMap<String, String>,
+}
+
+impl SupervisorBuilder {
+    pub fn new() -> Self {
+        Self {
+            redis_url: None,
+            zinit_socket_path: Some("/var/run/zinit.sock".to_string()),
+            osis_worker: None,
+            sal_worker: None,
+            v_worker: None,
+            python_worker: None,
+            worker_env_vars: HashMap::new(),
+        }
+    }
+
+    pub fn redis_url(mut self, url: &str) -> Self {
+        self.redis_url = Some(url.to_string());
+        self
+    }
+
+    pub fn zinit_socket_path(mut self, path: &str) -> Self {
+        self.zinit_socket_path = Some(path.to_string());
+        self
+    }
+
+    pub fn osis_worker(mut self, binary_path: &str) -> Self {
+        self.osis_worker = Some(binary_path.to_string());
+        self
+    }
+
+    pub fn sal_worker(mut self, binary_path: &str) -> Self {
+        self.sal_worker = Some(binary_path.to_string());
+        self
+    }
+
+    pub fn v_worker(mut self, binary_path: &str) -> Self {
+        self.v_worker = Some(binary_path.to_string());
+        self
+    }
+
+    pub fn python_worker(mut self, binary_path: &str) -> Self {
+        self.python_worker = Some(binary_path.to_string());
+        self
+    }
+
+    pub fn worker_env_var(mut self, key: &str, value: &str) -> Self {
+        self.worker_env_vars.insert(key.to_string(), value.to_string());
+        self
+    }
+
+    pub fn worker_env_vars(mut self, env_vars: HashMap<String, String>) -> Self {
+        self.worker_env_vars.extend(env_vars);
+        self
+    }
+
+    /// Builds the final `Supervisor` instance synchronously.
+    ///
+    /// This method validates the configuration and creates the Redis client.
+    /// Worker launching is deferred to the `start_workers()` method.
+    ///
+    /// # Returns
+    ///
+    /// * `Ok(Supervisor)` - Successfully configured client
+    /// * `Err(SupervisorError)` - Configuration or connection error
+    pub fn build(self) -> Result<Supervisor, SupervisorError> {
+        let url = self.redis_url
+            .unwrap_or_else(|| "redis://127.0.0.1/".to_string());
+        let client = redis::Client::open(url)?;
+        
+        let zinit_socket = self.zinit_socket_path
+            .unwrap_or_else(|| "/var/run/zinit.sock".to_string());
+        let zinit_client = ZinitClient::new(&zinit_socket);
+        
+        // Store builder data for later use in start_workers()
+        let builder_data = SupervisorBuilderData {
+            osis_worker: self.osis_worker,
+            sal_worker: self.sal_worker,
+            v_worker: self.v_worker,
+            python_worker: self.python_worker,
+            worker_env_vars: self.worker_env_vars,
+        };
+        
+        let supervisor = Supervisor {
+            redis_client: client,
+            zinit_client,
+            builder_data: Some(builder_data),
+        };
+        
+        Ok(supervisor)
+    }
+}
+
+impl Supervisor {
+    /// Start all configured workers asynchronously.
+    /// This method should be called after build() to launch the workers.
+    pub async fn start_workers(&self) -> Result<(), SupervisorError> {
+        // Clean up any existing worker services first
+        self.cleanup_existing_workers().await?;
+        
+        // Launch configured workers if builder data is available
+        if let Some(builder_data) = &self.builder_data {
+            self.launch_configured_workers(builder_data).await?;
+        }
+        
+        Ok(())
+    }
+    
+    /// Clean up all worker services from zinit on program exit
+    pub async fn cleanup_and_shutdown(&self) -> Result<(), SupervisorError> {
+        info!("Cleaning up worker services before shutdown...");
+        
+        let worker_names = vec![
+            "osis_worker_1",
+            "sal_worker_1", 
+            "v_worker_1",
+            "python_worker_1"
+        ];
+        
+        for worker_name in worker_names {
+            if let Err(e) = self.stop_and_delete_worker(worker_name).await {
+                warn!("Failed to cleanup worker {}: {}", worker_name, e);
+            }
+        }
+        
+        info!("Worker cleanup completed");
+        Ok(())
+    }
+    
+    /// Clean up any existing worker services on startup
+    async fn cleanup_existing_workers(&self) -> Result<(), SupervisorError> {
+        info!("Cleaning up any existing worker services...");
+        
+        let worker_names = vec![
+            "osis_worker_1",
+            "sal_worker_1", 
+            "v_worker_1",
+            "python_worker_1"
+        ];
+        
+        for worker_name in worker_names {
+            // Try to stop and delete, but don't fail if they don't exist
+            let _ = self.stop_and_delete_worker(worker_name).await;
+        }
+        
+        info!("Existing worker cleanup completed");
+        Ok(())
+    }
+    
+    /// Stop and delete a worker service from zinit
+    async fn stop_and_delete_worker(&self, worker_name: &str) -> Result<(), SupervisorError> {
+        // First try to stop the worker
+        if let Err(e) = self.zinit_client.stop(worker_name).await {
+            debug!("Worker {} was not running or failed to stop: {}", worker_name, e);
+        }
+        
+        // Then try to delete the service
+        if let Err(e) = self.zinit_client.delete(worker_name).await {
+            debug!("Worker {} service did not exist or failed to delete: {}", worker_name, e);
+        } else {
+            info!("Successfully deleted worker service: {}", worker_name);
+        }
+        
+        Ok(())
+    }
+
+    /// Get the hardcoded worker queue key for the script type
+    fn get_worker_queue_key(&self, script_type: &ScriptType) -> String {
+        format!("{}worker_queue:{}", NAMESPACE_PREFIX, script_type.worker_queue_suffix())
+    }
+
+    pub fn new_job(&self) -> JobBuilder {
+        JobBuilder::new(self)
+    }
+
+    // Internal helper to submit script details and push to work queue
+    async fn create_job_using_connection(
+        &self,
+        conn: &mut redis::aio::MultiplexedConnection,
+        job: &Job,
+    ) -> Result<(), SupervisorError> {
+        debug!(
+            "Submitting play request: {} for script type: {:?} with namespace prefix: {}",
+            job.id, job.script_type, NAMESPACE_PREFIX
+        );
+
+        // Use the shared Job struct's Redis storage method
+        job.store_in_redis(conn).await
+            .map_err(|e| SupervisorError::InvalidInput(format!("Failed to store job in Redis: {}", e)))?;
+            
+        Ok(())
+    }
+    
+    // Internal helper to submit script details and push to work queue
+    async fn start_job_using_connection(
+        &self,
+        conn: &mut redis::aio::MultiplexedConnection,
+        job_id: String,
+        script_type: &ScriptType
+    ) -> Result<(), SupervisorError> {
+        let worker_queue_key = self.get_worker_queue_key(script_type);
+
+        // lpush also infers its types, RV is typically i64 (length of list) or () depending on exact command variant
+        // For `redis::AsyncCommands::lpush`, it's `RedisResult<R>` where R: FromRedisValue
+        // Often this is the length of the list. Let's allow inference or specify if needed.
+        let _: redis::RedisResult<i64> =
+            conn.lpush(&worker_queue_key, job_id.clone()).await;
+
+        Ok(())
+    }
+
+    // Internal helper to await response from worker
+    async fn await_response_from_connection(
+        &self,
+        conn: &mut redis::aio::MultiplexedConnection,
+        job_key: &String,
+        reply_queue_key: &String,
+        timeout: Duration,
+    ) -> Result<String, SupervisorError> {
+        // BLPOP on the reply queue
+        // The timeout for BLPOP is in seconds (integer)
+        let blpop_timeout_secs = timeout.as_secs().max(1); // Ensure at least 1 second for BLPOP timeout
+
+        match conn
+            .blpop::<&String, Option<(String, String)>>(reply_queue_key, blpop_timeout_secs as f64)
+            .await
+        {
+            Ok(Some((_queue, result_message_str))) => {
+                Ok(result_message_str)
+            }
+            Ok(None) => {
+                // BLPOP timed out
+                warn!(
+                    "Timeout waiting for result on reply queue {} for job {}",
+                    reply_queue_key, job_key
+                );
+                // Optionally, delete the reply queue
+                let _: redis::RedisResult<i32> = conn.del(&reply_queue_key).await;
+                Err(SupervisorError::Timeout(job_key.clone()))
+            }
+            Err(e) => {
+                // Redis error
+                error!(
+                    "Redis error on BLPOP for reply queue {}: {}",
+                    reply_queue_key, e
+                );
+                // Optionally, delete the reply queue
+                let _: redis::RedisResult<i32> = conn.del(&reply_queue_key).await;
+                Err(SupervisorError::RedisError(e))
+            }
+        }
+    }
+
+    // New method using dedicated reply queue
+    pub async fn create_job(
+        &self,
+        job: &Job,
+    ) -> Result<(), SupervisorError> {
+        let mut conn = self.redis_client.get_multiplexed_async_connection().await?;
+
+        self.create_job_using_connection(
+            &mut conn,
+            &job, // Pass the job_id parameter
+        )
+        .await?;
+        Ok(())
+    }
+
+    // Method to start a previously created job
+    pub async fn start_job(
+        &self,
+        job_id: &str,
+    ) -> Result<(), SupervisorError> {
+        let mut conn = self.redis_client.get_multiplexed_async_connection().await?;
+        
+        // Load the job to get its script type
+        let job = Job::load_from_redis(&mut conn, job_id).await?;
+        
+        self.start_job_using_connection(&mut conn, job_id.to_string(), &job.script_type).await?;
+        Ok(())
+    }
+
+    // New method using dedicated reply queue with automatic worker selection
+    pub async fn run_job_and_await_result(
+        &self,
+        job: &Job
+    ) -> Result<String, SupervisorError> {
+        let mut conn = self.redis_client.get_multiplexed_async_connection().await?;
+
+        let reply_queue_key = format!("{}:reply:{}", NAMESPACE_PREFIX, job.id); // Derived from the passed job_id
+
+        self.create_job_using_connection(
+            &mut conn,
+            &job, // Pass the job_id parameter
+        )
+        .await?;
+
+        self.start_job_using_connection(&mut conn, job.id.clone(), &job.script_type).await?;
+
+        info!(
+            "Task {} submitted. Waiting for result on queue {} with timeout {:?}...",
+            job.id, // This is the UUID
+            reply_queue_key,
+            job.timeout
+        );
+
+        self.await_response_from_connection(
+            &mut conn,
+            &job.id,
+            &reply_queue_key,
+            job.timeout,
+        )
+        .await
+    }
+
+    // Method to get job status
+    pub async fn get_job_status(
+        &self,
+        job_id: &str,
+    ) -> Result<JobStatus, SupervisorError> {
+        let mut conn = self.redis_client.get_multiplexed_async_connection().await?;
+        let job_key = format!("{}{}", NAMESPACE_PREFIX, job_id);
+
+        let result_map: Option<std::collections::HashMap<String, String>> =
+            conn.hgetall(&job_key).await?;
+
+        match result_map {
+            Some(map) => {
+                let status_str = map.get("status").cloned().unwrap_or_else(|| {
+                    warn!("Task {}: 'status' field missing from Redis hash, defaulting to empty.", job_id);
+                    String::new()
+                });
+                
+                let status = match status_str.as_str() {
+                    "dispatched" => JobStatus::Dispatched,
+                    "started" => JobStatus::Started,
+                    "error" => JobStatus::Error,
+                    "finished" => JobStatus::Finished,
+                    _ => JobStatus::Dispatched, // default
+                };
+                
+                Ok(status)
+            }
+            None => {
+                warn!("Job {} not found in Redis", job_id);
+                Ok(JobStatus::Dispatched) // default for missing jobs
+            }
+        }
+    }
+
+    // Method to get job output
+    pub async fn get_job_output(
+        &self,
+        job_id: &str,
+    ) -> Result<Option<String>, SupervisorError> {
+        let mut conn = self.redis_client.get_multiplexed_async_connection().await?;
+        let job_key = format!("{}{}", NAMESPACE_PREFIX, job_id);
+
+        let result_map: Option<std::collections::HashMap<String, String>> =
+            conn.hgetall(&job_key).await?;
+
+        match result_map {
+            Some(map) => {
+                Ok(map.get("output").cloned())
+            }
+            None => {
+                warn!("Job {} not found in Redis", job_id);
+                Ok(None)
+            }
+        }
+    }
+
+    /// List all jobs in Redis
+    pub async fn list_jobs(&self) -> Result<Vec<String>, SupervisorError> {
+        let mut conn = self.redis_client.get_multiplexed_async_connection().await?;
+        
+        // Use the shared Job struct's list method
+        Job::list_all_job_ids(&mut conn).await
+            .map_err(|e| SupervisorError::InvalidInput(format!("Failed to list jobs: {}", e)))
+    }
+
+    /// Stop a job by pushing its ID to the stop queue
+    pub async fn stop_job(&self, job_id: &str) -> Result<(), SupervisorError> {
+        let mut conn = self.redis_client.get_multiplexed_async_connection().await?;
+        
+        // Get job details to determine script type and appropriate worker
+        let job_key = format!("{}job:{}", NAMESPACE_PREFIX, job_id);
+        let job_data: std::collections::HashMap<String, String> = conn.hgetall(&job_key).await?;
+        
+        if job_data.is_empty() {
+            return Err(SupervisorError::InvalidInput(format!("Job {} not found", job_id)));
+        }
+        
+        // Parse script type from job data
+        let script_type_str = job_data.get("script_type")
+            .ok_or_else(|| SupervisorError::InvalidInput("Job missing script_type field".to_string()))?;
+        
+        let script_type: ScriptType = serde_json::from_str(&format!("\"{}\"", script_type_str))
+            .map_err(|e| SupervisorError::InvalidInput(format!("Invalid script type: {}", e)))?;
+        
+        // Use hardcoded stop queue key for this script type
+        let stop_queue_key = format!("{}stop_queue:{}", NAMESPACE_PREFIX, script_type.worker_queue_suffix());
+        
+        // Push job ID to the stop queue
+        conn.lpush::<_, _, ()>(&stop_queue_key, job_id).await?;
+        
+        info!("Job {} added to stop queue {} for script type {:?}", job_id, stop_queue_key, script_type);
+        Ok(())
+    }
+
+    /// Get logs for a job by reading from its log file
+    pub async fn get_job_logs(&self, job_id: &str) -> Result<Option<String>, SupervisorError> {
+        let mut conn = self.redis_client.get_multiplexed_async_connection().await?;
+        let job_key = format!("{}job:{}", NAMESPACE_PREFIX, job_id);
+
+        // Get the job data to find the log path
+        let result_map: Option<std::collections::HashMap<String, String>> =
+            conn.hgetall(&job_key).await?;
+
+        match result_map {
+            Some(map) => {
+                if let Some(log_path) = map.get("log_path") {
+                    // Try to read the log file
+                    match std::fs::read_to_string(log_path) {
+                        Ok(contents) => Ok(Some(contents)),
+                        Err(e) => {
+                            warn!("Failed to read log file {}: {}", log_path, e);
+                            Ok(None)
+                        }
+                    }
+                } else {
+                    // No log path configured for this job
+                    Ok(None)
+                }
+            }
+            None => {
+                warn!("Job {} not found in Redis", job_id);
+                Ok(None)
+            }
+        }
+    }
+
+    /// Delete a specific job by ID
+    pub async fn delete_job(&self, job_id: &str) -> Result<(), SupervisorError> {
+        let mut conn = self.redis_client.get_multiplexed_async_connection().await?;
+        
+        // Use the shared Job struct's delete method
+        Job::delete_from_redis(&mut conn, job_id).await
+            .map_err(|e| SupervisorError::InvalidInput(format!("Failed to delete job: {}", e)))?;
+        
+        info!("Job {} deleted successfully", job_id);
+        Ok(())
+    }
+
+    /// Clear all jobs from Redis
+    pub async fn clear_all_jobs(&self) -> Result<usize, SupervisorError> {
+        let mut conn = self.redis_client.get_multiplexed_async_connection().await?;
+        
+        // Get all job IDs first
+        let job_ids = Job::list_all_job_ids(&mut conn).await
+            .map_err(|e| SupervisorError::InvalidInput(format!("Failed to list jobs: {}", e)))?;
+        
+        let count = job_ids.len();
+        
+        // Delete each job using the shared method
+        for job_id in job_ids {
+            Job::delete_from_redis(&mut conn, &job_id).await
+                .map_err(|e| SupervisorError::InvalidInput(format!("Failed to delete job {}: {}", job_id, e)))?;
+        }
+        
+        Ok(count)
+    }
+
+    /// Check if all prerequisites for a job are completed
+    pub async fn check_prerequisites_completed(&self, job_id: &str) -> Result<bool, SupervisorError> {
+        let mut conn = self.redis_client.get_multiplexed_async_connection().await?;
+        
+        // Load the job using the shared Job struct
+        let job = Job::load_from_redis(&mut conn, job_id).await
+            .map_err(|e| SupervisorError::InvalidInput(format!("Failed to load job: {}", e)))?;
+        
+        // Check each prerequisite job status
+        for prereq_id in &job.prerequisites {
+            let status = Job::get_status(&mut conn, prereq_id).await
+                .map_err(|e| SupervisorError::InvalidInput(format!("Failed to get prerequisite status: {}", e)))?;
+            
+            if status != JobStatus::Finished {
+                return Ok(false); // Prerequisite not completed
+            }
+        }
+        
+        Ok(true) // All prerequisites completed (or no prerequisites)
+    }
+
+    /// Update job status and check dependent jobs for readiness
+    pub async fn update_job_status_and_check_dependents(&self, job_id: &str, new_status: JobStatus) -> Result<Vec<String>, SupervisorError> {
+        let mut conn = self.redis_client.get_multiplexed_async_connection().await?;
+        
+        // Update job status using shared Job method
+        Job::update_status(&mut conn, job_id, new_status.clone()).await
+            .map_err(|e| SupervisorError::InvalidInput(format!("Failed to update job status: {}", e)))?;
+        
+        let mut ready_jobs = Vec::new();
+        
+        // If job finished, check dependent jobs
+        if new_status == JobStatus::Finished {
+            // Load the job to get its dependents
+            let job = Job::load_from_redis(&mut conn, job_id).await
+                .map_err(|e| SupervisorError::InvalidInput(format!("Failed to load job: {}", e)))?;
+            
+            // Check each dependent job
+            for dependent_id in &job.dependents {
+                let dependent_status = Job::get_status(&mut conn, dependent_id).await
+                    .map_err(|e| SupervisorError::InvalidInput(format!("Failed to get dependent status: {}", e)))?;
+                
+                // Only check jobs that are waiting for prerequisites
+                if dependent_status == JobStatus::WaitingForPrerequisites {
+                    // Check if all prerequisites are now completed
+                    if self.check_prerequisites_completed(dependent_id).await? {
+                        // Update status to dispatched and add to ready jobs
+                        Job::update_status(&mut conn, dependent_id, JobStatus::Dispatched).await
+                            .map_err(|e| SupervisorError::InvalidInput(format!("Failed to update dependent status: {}", e)))?;
+                        ready_jobs.push(dependent_id.clone());
+                    }
+                }
+            }
+        }
+        
+        Ok(ready_jobs)
+    }
+
+    /// Dispatch jobs that are ready (have all prerequisites completed)
+    pub async fn dispatch_ready_jobs(&self, ready_job_ids: Vec<String>) -> Result<(), SupervisorError> {
+        for job_id in ready_job_ids {
+            // Get job data to determine script type and select worker
+            let mut conn = self.redis_client.get_multiplexed_async_connection().await?;
+            let job_key = format!("{}job:{}", NAMESPACE_PREFIX, job_id);
+            let job_data: std::collections::HashMap<String, String> = conn.hgetall(&job_key).await?;
+            
+            if let Some(script_type_str) = job_data.get("script_type") {
+                // Parse script type (stored as Debug format, e.g., "OSIS")
+                let script_type = match script_type_str.as_str() {
+                    "OSIS" => ScriptType::OSIS,
+                    "SAL" => ScriptType::SAL,
+                    "V" => ScriptType::V,
+                    "Python" => ScriptType::Python,
+                    _ => return Err(SupervisorError::InvalidInput(format!("Unknown script type: {}", script_type_str))),
+                };
+                
+                // Dispatch job using hardcoded queue
+                self.start_job_using_connection(&mut conn, job_id, &script_type).await?;
+            }
+        }
+        Ok(())
+    }
+}
+