taskflow-rs 0.1.1

pub mod config;
pub mod handlers;
pub mod task_executor;

use std::collections::HashMap;
use std::sync::Arc;
use tokio::time::Duration;
use tracing::{error, info, warn};

use crate::Task;

use crate::error::Result;
use crate::scheduler::Scheduler;
use crate::task::TaskHandler;
pub use config::ExecutorConfig;
pub use handlers::{HttpTaskHandler, ShellTaskHandler};
pub use task_executor::TaskExecutor;

#[derive(Clone)]
pub struct Executor {
    config: ExecutorConfig,
    scheduler: Arc<Scheduler>,
    handlers: HashMap<String, Arc<dyn TaskHandler>>,
    executor: TaskExecutor,
}

impl Executor {
    pub fn new(scheduler: Arc<Scheduler>, config: ExecutorConfig) -> Self {
        Self {
            config,
            scheduler,
            handlers: HashMap::new(),
            executor: TaskExecutor::new(),
        }
    }

    pub fn register_handler(&mut self, handler: Arc<dyn TaskHandler>) {
        let task_type = handler.task_type().to_string();
        self.handlers.insert(task_type.clone(), handler);
        info!("Registered task handler for type: {}", task_type);
    }

    pub async fn start(&self) -> Result<()> {
        info!("Starting executor: {}", self.config.worker_id);

        let shutdown_signal = self.executor.wait_for_shutdown();

        let executor_arc = Arc::new(self.clone());

        let heartbeat_task = {
            let executor = Arc::clone(&executor_arc);
            tokio::spawn(async move {
                executor.start_heartbeat_loop().await;
            })
        };

        let execution_task = {
            let executor = Arc::clone(&executor_arc);
            tokio::spawn(async move {
                executor.start_execution_loop().await;
            })
        };

        tokio::select! {
            _ = shutdown_signal.notified() => {
                info!("Shutdown signal received");
            }
            result = heartbeat_task => {
                if let Err(e) = result {
                    error!("Heartbeat task failed: {}", e);
                }
            }
            result = execution_task => {
                if let Err(e) = result {
                    error!("Execution task failed: {}", e);
                }
            }
        }

        self.executor.wait_for_running_tasks().await;
        info!("Executor stopped: {}", self.config.worker_id);
        Ok(())
    }

    pub async fn shutdown(&self) {
        info!("Initiating executor shutdown: {}", self.config.worker_id);
        self.executor.shutdown();
    }

    async fn start_execution_loop(&self) {
        let mut interval = tokio::time::interval(Duration::from_secs(1));

        loop {
            tokio::select! {
                _ = self.executor.wait_for_shutdown().notified() => {
                    break;
                }
                _ = interval.tick() => {
                    if let Err(e) = self.process_tasks().await {
                        error!("Error processing tasks: {}", e);
                    }
                }
            }
        }
    }

    async fn start_heartbeat_loop(&self) {
        let mut interval =
            tokio::time::interval(Duration::from_secs(self.config.heartbeat_interval_seconds));

        loop {
            tokio::select! {
                _ = self.executor.wait_for_shutdown().notified() => {
                    break;
                }
                _ = interval.tick() => {
                    self.send_heartbeat().await;
                }
            }
        }
    }

    async fn process_tasks(&self) -> Result<()> {
        let running_count = self.executor.running_count().await;
        if running_count >= self.config.max_concurrent_tasks {
            return Ok(());
        }

        let available_slots = self.config.max_concurrent_tasks - running_count;

        for _ in 0..available_slots {
            if let Some(task) = self.scheduler.get_next_task().await {
                if self.handlers.contains_key(&task.definition.task_type) {
                    self.execute_task(task).await;
                } else {
                    warn!(
                        "No handler registered for task type: {}",
                        task.definition.task_type
                    );
                    self.scheduler
                        .complete_task(
                            &task.definition.id,
                            false,
                            None,
                            Some(format!(
                                "No handler for task type: {}",
                                task.definition.task_type
                            )),
                        )
                        .await?;
                }
            } else {
                break;
            }
        }

        Ok(())
    }

    async fn execute_task(&self, mut task: Task) {
        let task_id = task.definition.id.clone();
        let task_type = task.definition.task_type.clone();

        self.executor.add_running_task(task_id.clone()).await;

        task.start_execution(&self.config.worker_id);

        let handler = self.handlers.get(&task_type).unwrap().clone();
        let scheduler = Arc::clone(&self.scheduler);
        let executor = self.executor.clone();
        let timeout_duration = Duration::from_secs(
            task.definition
                .timeout_seconds
                .max(self.config.task_timeout_seconds),
        );

        tokio::spawn(async move {
            let result = tokio::time::timeout(timeout_duration, handler.execute(&task)).await;

            let (success, output, error) = match result {
                Ok(Ok(task_result)) => (task_result.success, task_result.output, task_result.error),
                Ok(Err(e)) => (false, None, Some(e.to_string())),
                Err(_) => (false, None, Some("Task execution timeout".to_string())),
            };

            if let Err(e) = scheduler
                .complete_task(&task_id, success, output, error)
                .await
            {
                error!("Failed to complete task {}: {}", task_id, e);
            }

            executor.remove_running_task(&task_id).await;

            info!(
                "Task execution finished: {} (success: {})",
                task_id, success
            );
        });
    }

    async fn send_heartbeat(&self) {
        let running_count = self.executor.running_count().await;
        info!(
            "Executor heartbeat: {} (running tasks: {})",
            self.config.worker_id, running_count
        );
    }
}