eshanized_polaris_core/

cluster.rs

//! Cluster management and orchestration.
//!
//! This module provides the main cluster API for submitting and managing tasks.

use crate::config::{ClusterConfig, NodeConfig};
use crate::errors::{PolarisError, PolarisResult};
use crate::executor::Executor;
use crate::node::{Node, NodeRegistry, NodeStatus};
use crate::observability::MetricsCollector;
use crate::scheduler::{RoundRobinScheduler, Scheduler};
use crate::task::{Task, TaskHandle, TaskId, TaskResult, TaskStatus};
use std::sync::Arc;
use std::time::Duration;

/// Main cluster handle
#[derive(Clone)]
pub struct Cluster {
    config: Arc<ClusterConfig>,
    registry: NodeRegistry,
    scheduler: Arc<dyn Scheduler>,
    executor: Executor,
    metrics: MetricsCollector,
    tasks: Arc<dashmap::DashMap<TaskId, Task>>,
    local_mode: bool,
}

impl Cluster {
    /// Create a new cluster builder
    pub fn builder() -> ClusterBuilder {
        ClusterBuilder::default()
    }

    /// Connect to an existing cluster
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use polaris_core::prelude::*;
    /// # async fn example() -> PolarisResult<()> {
    /// let cluster = Cluster::connect(["10.0.0.1:7001", "10.0.0.2:7001"]).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn connect<I, S>(seeds: I) -> PolarisResult<Self>
    where
        I: IntoIterator<Item = S>,
        S: AsRef<str>,
    {
        let seed_list: Vec<String> = seeds.into_iter().map(|s| s.as_ref().to_string()).collect();

        if seed_list.is_empty() {
            return Err(PolarisError::InvalidConfig("No seed nodes provided".into()));
        }

        let config = ClusterConfig {
            seed_nodes: seed_list.clone(),
            ..Default::default()
        };

        let builder = ClusterBuilder::new().with_config(config);

        // Connect to seed nodes
        for seed in &seed_list {
            tracing::info!(seed = %seed, "Connecting to seed node");
            // In a real implementation, this would establish network connections
        }

        builder.build().await
    }

    /// Submit a task to the cluster
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use polaris_core::prelude::*;
    /// # use bytes::Bytes;
    /// # async fn example(cluster: Cluster) -> PolarisResult<()> {
    /// let task = Task::new("my_task", Bytes::from("payload"));
    /// let handle = cluster.submit(task).await?;
    /// let result = handle.result().await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn submit(&self, mut task: Task) -> PolarisResult<TaskHandle> {
        tracing::info!(task_id = %task.id, task_name = %task.name, "Submitting task");
        
        // Record metrics
        self.metrics.record_task_submitted();

        // Get available nodes
        let nodes = self.registry.available();
        if nodes.is_empty() {
            return Err(PolarisError::scheduling_failed("No available nodes"));
        }

        // Schedule the task
        let node_id = self.scheduler.schedule(&task, &nodes).await?;

        // Update task status
        task.transition_to(TaskStatus::Scheduled)?;
        task.metadata.node_id = Some(node_id.to_string());
        task.metadata.scheduled_at = Some(chrono::Utc::now());

        // Store task
        self.tasks.insert(task.id, task.clone());

        // Create handle
        let handle = TaskHandle::new(task.id);
        handle.update_status(TaskStatus::Scheduled);

        // In local mode, execute immediately
        if self.local_mode {
            self.execute_task_local(task, handle.clone()).await?;
        }

        Ok(handle)
    }

    /// Execute a task locally (for local mode or testing)
    async fn execute_task_local(&self, task: Task, handle: TaskHandle) -> PolarisResult<()> {
        self.executor.submit(task, handle)
    }

    /// Submit multiple tasks and collect results
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use polaris_core::prelude::*;
    /// # use bytes::Bytes;
    /// # async fn example(cluster: Cluster) -> PolarisResult<()> {
    /// let tasks = vec![
    ///     Task::new("task1", Bytes::from("input1")),
    ///     Task::new("task2", Bytes::from("input2")),
    /// ];
    /// let results = cluster.submit_batch(tasks).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn submit_batch(&self, tasks: Vec<Task>) -> PolarisResult<Vec<TaskHandle>> {
        let mut handles = Vec::with_capacity(tasks.len());

        for task in tasks {
            let handle = self.submit(task).await?;
            handles.push(handle);
        }

        Ok(handles)
    }

    /// Map over a list of tasks
    /// 
    /// Submit multiple tasks and collect their results in parallel.
    /// 
    /// # Example
    /// 
    /// ```no_run
    /// # use polaris_core::prelude::*;
    /// # use bytes::Bytes;
    /// # async fn example(cluster: Cluster) -> PolarisResult<()> {
    /// let tasks = vec![
    ///     Task::new("task1", Bytes::from("input1")),
    ///     Task::new("task2", Bytes::from("input2")),
    /// ];
    /// 
    /// let results = cluster.map(tasks)
    ///     .with_retry(5)
    ///     .await_impl()
    ///     .await?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn map<T, I: IntoIterator<Item = Task>>(&self, tasks: I) -> MapBuilder<T> {
        MapBuilder::new(self.clone(), tasks.into_iter().collect())
    }

    /// Get task status
    pub fn task_status(&self, task_id: TaskId) -> Option<TaskStatus> {
        self.tasks.get(&task_id).map(|t| t.status)
    }

    /// Get cluster statistics
    pub fn stats(&self) -> ClusterStats {
        let nodes = self.registry.all();
        let total_nodes = nodes.len();
        let available_nodes = nodes.iter().filter(|n| n.is_available()).count();

        let total_tasks = self.tasks.len();
        let running_tasks = self
            .tasks
            .iter()
            .filter(|t| t.status == TaskStatus::Running)
            .count();

        ClusterStats {
            total_nodes,
            available_nodes,
            total_tasks,
            running_tasks,
        }
    }

    /// Shutdown the cluster gracefully
    pub async fn shutdown(&self) -> PolarisResult<()> {
        tracing::info!("Shutting down cluster");

        // Mark all nodes as draining
        for node in self.registry.all() {
            node.set_status(NodeStatus::Draining);
        }

        // Wait for running tasks to complete (with timeout)
        let timeout = Duration::from_secs(30);
        let start = std::time::Instant::now();

        while start.elapsed() < timeout {
            let running = self
                .tasks
                .iter()
                .filter(|t| t.status == TaskStatus::Running)
                .count();

            if running == 0 {
                break;
            }

            tokio::time::sleep(Duration::from_millis(100)).await;
        }

        tracing::info!("Cluster shutdown complete");
        Ok(())
    }
}

impl std::fmt::Debug for Cluster {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Cluster")
            .field("config", &self.config)
            .field("node_count", &self.registry.count())
            .field("task_count", &self.tasks.len())
            .finish()
    }
}

/// Builder for creating a cluster
#[derive(Default)]
pub struct ClusterBuilder {
    config: Option<ClusterConfig>,
    scheduler: Option<Arc<dyn Scheduler>>,
    local_mode: bool,
}

impl ClusterBuilder {
    /// Create a new cluster builder
    pub fn new() -> Self {
        Self::default()
    }

    /// Set cluster configuration
    pub fn with_config(mut self, config: ClusterConfig) -> Self {
        self.config = Some(config);
        self
    }

    /// Set custom scheduler
    pub fn with_scheduler<S: Scheduler + 'static>(mut self, scheduler: S) -> Self {
        self.scheduler = Some(Arc::new(scheduler));
        self
    }

    /// Enable local mode (single node, in-process execution)
    pub fn with_local_node(mut self) -> Self {
        self.local_mode = true;
        self
    }

    /// Build the cluster
    pub async fn build(self) -> PolarisResult<Cluster> {
        let config = self.config.unwrap_or_default();
        config.validate()?;

        let registry = NodeRegistry::new();

        // Create local node if in local mode
        if self.local_mode {
            let node_config = NodeConfig::default();
            let node = Node::new("local-node", node_config.bind_addr);
            node.set_status(NodeStatus::Ready);
            registry.register(node)?;
            tracing::info!("Started cluster in local mode");
        }

        let scheduler = self
            .scheduler
            .unwrap_or_else(|| Arc::new(RoundRobinScheduler::new()));

        Ok(Cluster {
            config: Arc::new(config),
            registry,
            scheduler,
            executor: Executor::new(),
            metrics: MetricsCollector::new(),
            tasks: Arc::new(dashmap::DashMap::new()),
            local_mode: self.local_mode,
        })
    }
}

/// Cluster statistics
#[derive(Debug, Clone)]
pub struct ClusterStats {
    /// Total number of nodes
    pub total_nodes: usize,
    /// Number of available nodes
    pub available_nodes: usize,
    /// Total tasks submitted
    pub total_tasks: usize,
    /// Currently running tasks
    pub running_tasks: usize,
}

/// Builder for map operations
/// 
/// Allows batching multiple tasks with retry configuration.
pub struct MapBuilder<T> {
    cluster: Cluster,
    tasks: Vec<Task>,
    retry_count: u32,
    _phantom: std::marker::PhantomData<T>,
}

impl<T> MapBuilder<T> {
    fn new(cluster: Cluster, tasks: Vec<Task>) -> Self {
        Self {
            cluster,
            tasks,
            retry_count: 3,
            _phantom: std::marker::PhantomData,
        }
    }

    /// Set retry count for failed tasks
    /// 
    /// Each task will be retried up to this many times if it fails.
    pub fn with_retry(mut self, count: u32) -> Self {
        self.retry_count = count;
        self
    }

    /// Execute the map operation
    /// 
    /// Submits all tasks concurrently, waits for results, and retries failures.
    /// 
    /// # Returns
    /// 
    /// Returns task results in the order they were submitted.
    /// Note: Type parameter T is currently not used for deserialization.
    pub async fn await_impl(self) -> PolarisResult<Vec<TaskResult>> {
        use std::time::Duration;
        
        if self.tasks.is_empty() {
            return Ok(Vec::new());
        }

        tracing::info!(
            task_count = self.tasks.len(),
            retry_count = self.retry_count,
            "Starting map operation"
        );

        let mut results = Vec::with_capacity(self.tasks.len());
        let mut handles = Vec::with_capacity(self.tasks.len());

        // Submit all tasks
        for mut task in self.tasks {
            // Configure retry on the task itself
            task.max_retries = self.retry_count;
            
            let handle = self.cluster.submit(task).await?;
            handles.push(handle);
        }

        tracing::debug!("All tasks submitted, waiting for results");

        // Collect results with timeout
        for (idx, handle) in handles.into_iter().enumerate() {
            // Wait for task completion
            let result = handle.result().await?;
            
            if result.success {
                tracing::debug!(task_idx = idx, "Task succeeded");
            } else {
                tracing::warn!(
                    task_idx = idx,
                    error = ?result.error,
                    "Task failed after retries"
                );
            }
            
            results.push(result);
        }

        let success_count = results.iter().filter(|r| r.success).count();
        let failure_count = results.len() - success_count;

        tracing::info!(
            total = results.len(),
            succeeded = success_count,
            failed = failure_count,
            "Map operation completed"
        );

        Ok(results)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::scheduler::LoadAwareScheduler;
    use bytes::Bytes;

    #[tokio::test]
    async fn test_cluster_builder_local() {
        let cluster = Cluster::builder().with_local_node().build().await.unwrap();

        assert_eq!(cluster.registry.count(), 1);
        assert!(cluster.local_mode);
    }

    #[tokio::test]
    async fn test_cluster_submit_task() {
        let cluster = Cluster::builder().with_local_node().build().await.unwrap();

        let task = Task::new("test_task", Bytes::from("payload"));
        let handle = cluster.submit(task).await.unwrap();

        assert!(!handle.id.as_uuid().is_nil());
    }

    #[tokio::test]
    async fn test_cluster_map_operation() {
        let cluster = Cluster::builder().with_local_node().build().await.unwrap();

        // Create multiple tasks
        let tasks = vec![
            Task::new("map_task_1", Bytes::from("input1")),
            Task::new("map_task_2", Bytes::from("input2")),
            Task::new("map_task_3", Bytes::from("input3")),
        ];

        // Execute map operation with retry
        let results = cluster.map::<String, _>(tasks)
            .with_retry(2)
            .await_impl()
            .await
            .unwrap();

        // Verify results
        assert_eq!(results.len(), 3);
        assert!(results.iter().all(|r| r.success));
    }

    #[tokio::test]
    async fn test_cluster_submit_batch() {
        let cluster = Cluster::builder().with_local_node().build().await.unwrap();

        let tasks = vec![
            Task::new("task1", Bytes::new()),
            Task::new("task2", Bytes::new()),
            Task::new("task3", Bytes::new()),
        ];

        let handles = cluster.submit_batch(tasks).await.unwrap();
        assert_eq!(handles.len(), 3);
    }

    #[tokio::test]
    async fn test_cluster_stats() {
        let cluster = Cluster::builder().with_local_node().build().await.unwrap();

        let stats = cluster.stats();
        assert_eq!(stats.total_nodes, 1);
        assert_eq!(stats.available_nodes, 1);
    }

    #[tokio::test]
    async fn test_cluster_custom_scheduler() {
        let cluster = Cluster::builder()
            .with_local_node()
            .with_scheduler(LoadAwareScheduler::new())
            .build()
            .await
            .unwrap();

        let task = Task::new("test", Bytes::new());
        let result = cluster.submit(task).await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_cluster_no_nodes_error() {
        let cluster = Cluster::builder().build().await.unwrap();

        let task = Task::new("test", Bytes::new());
        let result = cluster.submit(task).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_cluster_connect_validation() {
        let result = Cluster::connect::<_, &str>([]).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_cluster_shutdown() {
        let cluster = Cluster::builder().with_local_node().build().await.unwrap();

        let result = cluster.shutdown().await;
        assert!(result.is_ok());
    }
}