eshanized_polaris_core/

scheduler.rs

//! Task scheduling implementations.
//!
//! This module provides pluggable schedulers for distributing tasks across nodes.

use crate::errors::{PolarisError, PolarisResult};
use crate::node::{Node, NodeId};
use crate::task::Task;
use async_trait::async_trait;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

/// Trait for implementing custom schedulers
#[async_trait]
pub trait Scheduler: Send + Sync {
    /// Schedule a task on one of the available nodes
    ///
    /// # Arguments
    ///
    /// * `task` - The task to schedule
    /// * `nodes` - Available nodes for scheduling
    ///
    /// # Returns
    ///
    /// Returns the selected node ID, or an error if scheduling fails
    async fn schedule(&self, task: &Task, nodes: &[Node]) -> PolarisResult<NodeId>;

    /// Get scheduler name
    fn name(&self) -> &str;
}

/// Round-robin scheduler
///
/// Distributes tasks evenly across nodes in a circular fashion.
#[derive(Debug)]
pub struct RoundRobinScheduler {
    counter: Arc<AtomicUsize>,
}

impl RoundRobinScheduler {
    /// Create a new round-robin scheduler
    pub fn new() -> Self {
        Self {
            counter: Arc::new(AtomicUsize::new(0)),
        }
    }
}

impl Default for RoundRobinScheduler {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl Scheduler for RoundRobinScheduler {
    async fn schedule(&self, _task: &Task, nodes: &[Node]) -> PolarisResult<NodeId> {
        if nodes.is_empty() {
            return Err(PolarisError::scheduling_failed("No available nodes"));
        }

        let index = self.counter.fetch_add(1, Ordering::Relaxed) % nodes.len();
        let node = &nodes[index];

        tracing::debug!(
            task_id = %_task.id,
            node_id = %node.id(),
            scheduler = "round_robin",
            "Task scheduled"
        );

        Ok(node.id())
    }

    fn name(&self) -> &str {
        "round_robin"
    }
}

/// Load-aware scheduler
///
/// Schedules tasks on the least loaded node based on current resource usage.
#[derive(Debug)]
pub struct LoadAwareScheduler {
    /// Weight for CPU usage in load calculation
    cpu_weight: f64,
    /// Weight for memory usage in load calculation
    memory_weight: f64,
    /// Weight for task count in load calculation
    task_weight: f64,
}

impl LoadAwareScheduler {
    /// Create a new load-aware scheduler with default weights
    pub fn new() -> Self {
        Self {
            cpu_weight: 0.4,
            memory_weight: 0.3,
            task_weight: 0.3,
        }
    }

    /// Create a load-aware scheduler with custom weights
    pub fn with_weights(cpu_weight: f64, memory_weight: f64, task_weight: f64) -> Self {
        let total = cpu_weight + memory_weight + task_weight;
        Self {
            cpu_weight: cpu_weight / total,
            memory_weight: memory_weight / total,
            task_weight: task_weight / total,
        }
    }

    /// Calculate load score for a node (lower is better)
    fn calculate_load(&self, node: &Node) -> f64 {
        let info = node.info();
        let usage = &info.resource_usage;

        let cpu_load = usage.cpu_usage_percent / 100.0;
        let memory_load = if info.resource_limits.max_memory_bytes > 0 {
            usage.memory_usage_bytes as f64 / info.resource_limits.max_memory_bytes as f64
        } else {
            0.0
        };
        let task_load = node.load_factor();

        self.cpu_weight * cpu_load
            + self.memory_weight * memory_load
            + self.task_weight * task_load
    }
}

impl Default for LoadAwareScheduler {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl Scheduler for LoadAwareScheduler {
    async fn schedule(&self, task: &Task, nodes: &[Node]) -> PolarisResult<NodeId> {
        if nodes.is_empty() {
            return Err(PolarisError::scheduling_failed("No available nodes"));
        }

        // Find node with lowest load
        let (best_node, load) = nodes
            .iter()
            .map(|node| (node, self.calculate_load(node)))
            .min_by(|(_, load_a), (_, load_b)| {
                load_a.partial_cmp(load_b).unwrap_or(std::cmp::Ordering::Equal)
            })
            .unwrap();

        tracing::debug!(
            task_id = %task.id,
            node_id = %best_node.id(),
            load = %load,
            scheduler = "load_aware",
            "Task scheduled"
        );

        Ok(best_node.id())
    }

    fn name(&self) -> &str {
        "load_aware"
    }
}

/// Priority-based scheduler
///
/// Considers task priority when scheduling, favoring higher priority tasks.
#[derive(Debug)]
pub struct PriorityScheduler {
    base_scheduler: LoadAwareScheduler,
}

impl PriorityScheduler {
    /// Create a new priority scheduler
    pub fn new() -> Self {
        Self {
            base_scheduler: LoadAwareScheduler::new(),
        }
    }
}

impl Default for PriorityScheduler {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl Scheduler for PriorityScheduler {
    async fn schedule(&self, task: &Task, nodes: &[Node]) -> PolarisResult<NodeId> {
        // For now, delegate to load-aware scheduler
        // In a full implementation, this would maintain priority queues per node
        tracing::debug!(
            task_id = %task.id,
            priority = ?task.priority,
            scheduler = "priority",
            "Scheduling high-priority task"
        );
        self.base_scheduler.schedule(task, nodes).await
    }

    fn name(&self) -> &str {
        "priority"
    }
}

/// Affinity-based scheduler
///
/// Schedules tasks with affinity constraints (e.g., locality, specific node labels).
#[derive(Debug)]
pub struct AffinityScheduler {
    base_scheduler: LoadAwareScheduler,
}

impl AffinityScheduler {
    /// Create a new affinity scheduler
    pub fn new() -> Self {
        Self {
            base_scheduler: LoadAwareScheduler::new(),
        }
    }

    /// Filter nodes by label selector
    fn filter_by_labels<'a>(
        &self,
        nodes: &'a [Node],
        required_labels: &std::collections::HashMap<String, String>,
    ) -> Vec<&'a Node> {
        if required_labels.is_empty() {
            return nodes.iter().collect();
        }

        nodes
            .iter()
            .filter(|node| {
                let node_labels = &node.info().metadata.labels;
                required_labels
                    .iter()
                    .all(|(k, v)| node_labels.get(k).map(|nv| nv == v).unwrap_or(false))
            })
            .collect()
    }
}

impl Default for AffinityScheduler {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl Scheduler for AffinityScheduler {
    async fn schedule(&self, task: &Task, nodes: &[Node]) -> PolarisResult<NodeId> {
        // Check for affinity constraints in task metadata
        let required_labels = &task.metadata.tags;

        let filtered_nodes: Vec<&Node> = self.filter_by_labels(nodes, required_labels);

        if filtered_nodes.is_empty() {
            return Err(PolarisError::scheduling_failed(
                "No nodes match affinity constraints",
            ));
        }

        // Convert to owned nodes for scheduling
        let owned_nodes: Vec<Node> = filtered_nodes.into_iter().cloned().collect();

        tracing::debug!(
            task_id = %task.id,
            constraints = ?required_labels,
            scheduler = "affinity",
            "Scheduling with affinity constraints"
        );

        self.base_scheduler.schedule(task, &owned_nodes).await
    }

    fn name(&self) -> &str {
        "affinity"
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::task::{TaskPriority, TaskStatus};
    use bytes::Bytes;

    fn create_test_nodes(count: usize) -> Vec<Node> {
        (0..count)
            .map(|i| {
                let addr = format!("127.0.0.1:700{}", i).parse().unwrap();
                let node = Node::new(format!("node{}", i), addr);
                node.set_status(crate::node::NodeStatus::Ready);
                node
            })
            .collect()
    }

    #[tokio::test]
    async fn test_round_robin_scheduler() {
        let scheduler = RoundRobinScheduler::new();
        let nodes = create_test_nodes(3);
        let task = Task::new("test", Bytes::new());

        // Schedule multiple tasks and verify round-robin behavior
        let mut node_counts = std::collections::HashMap::new();
        for _ in 0..9 {
            let node_id = scheduler.schedule(&task, &nodes).await.unwrap();
            *node_counts.entry(node_id).or_insert(0) += 1;
        }

        // Each node should get 3 tasks
        assert_eq!(node_counts.len(), 3);
        for count in node_counts.values() {
            assert_eq!(*count, 3);
        }
    }

    #[tokio::test]
    async fn test_round_robin_no_nodes() {
        let scheduler = RoundRobinScheduler::new();
        let nodes = vec![];
        let task = Task::new("test", Bytes::new());

        let result = scheduler.schedule(&task, &nodes).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_load_aware_scheduler() {
        let scheduler = LoadAwareScheduler::new();
        let nodes = create_test_nodes(3);

        // Make one node heavily loaded
        for _ in 0..10 {
            nodes[0].increment_active_tasks();
        }

        let task = Task::new("test", Bytes::new());
        let node_id = scheduler.schedule(&task, &nodes).await.unwrap();

        // Should not select the heavily loaded node
        assert_ne!(node_id, nodes[0].id());
    }

    #[tokio::test]
    async fn test_load_aware_custom_weights() {
        let scheduler = LoadAwareScheduler::with_weights(1.0, 0.0, 0.0);
        let nodes = create_test_nodes(2);
        let task = Task::new("test", Bytes::new());

        let result = scheduler.schedule(&task, &nodes).await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_priority_scheduler() {
        let scheduler = PriorityScheduler::new();
        let nodes = create_test_nodes(2);

        let task = Task::new("test", Bytes::new()).with_priority(TaskPriority::High);

        let result = scheduler.schedule(&task, &nodes).await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_affinity_scheduler_with_labels() {
        let scheduler = AffinityScheduler::new();
        let nodes = create_test_nodes(2);

        // Add label to first node
        {
            let mut info = nodes[0].info();
            info.metadata.labels.insert("zone".to_string(), "us-east".to_string());
        }

        let mut task = Task::new("test", Bytes::new());
        task.metadata.tags.insert("zone".to_string(), "us-east".to_string());

        let node_id = scheduler.schedule(&task, &nodes).await.unwrap();
        assert_eq!(node_id, nodes[0].id());
    }

    #[tokio::test]
    async fn test_affinity_scheduler_no_match() {
        let scheduler = AffinityScheduler::new();
        let nodes = create_test_nodes(2);

        let mut task = Task::new("test", Bytes::new());
        task.metadata.tags.insert("zone".to_string(), "us-west".to_string());

        let result = scheduler.schedule(&task, &nodes).await;
        assert!(result.is_err());
    }

    #[test]
    fn test_scheduler_names() {
        assert_eq!(RoundRobinScheduler::new().name(), "round_robin");
        assert_eq!(LoadAwareScheduler::new().name(), "load_aware");
        assert_eq!(PriorityScheduler::new().name(), "priority");
        assert_eq!(AffinityScheduler::new().name(), "affinity");
    }
}