mecha10_core/priority_queue/

mod.rs

//! Priority queue for message processing
//!
//! Enables quality-of-service (QoS) by prioritizing critical messages over
//! low-priority ones. Useful for safety-critical systems where some messages
//! (e.g., emergency stop) must be processed before others (e.g., telemetry).
//!
//! # Features
//!
//! - **Priority Levels**: Define message importance (Critical, High, Normal, Low)
//! - **Fair Processing**: Round-robin within same priority level
//! - **Starvation Prevention**: Ensures low-priority messages eventually process
//! - **Backpressure Support**: Integrates with bounded channels
//! - **Zero-Copy**: Messages moved, not copied
//!
//! # Example
//!
//! ```rust
//! use mecha10::prelude::*;
//! use mecha10::priority_queue::{PriorityReceiver, Priority};
//!
//! # async fn example(ctx: &Context) -> Result<()> {
//! // Create priority receiver
//! let mut prio_rx = PriorityReceiver::new();
//!
//! // Subscribe to multiple topics with different priorities
//! let emergency = ctx.subscribe("/safety/emergency").await?;
//! let commands = ctx.subscribe("/motor/command").await?;
//! let telemetry = ctx.subscribe("/sensor/telemetry").await?;
//!
//! // Add with priorities
//! prio_rx.add_source(emergency, Priority::Critical);
//! prio_rx.add_source(commands, Priority::High);
//! prio_rx.add_source(telemetry, Priority::Low);
//!
//! // Process messages in priority order
//! while let Some((msg, priority)) = prio_rx.recv().await {
//!     match priority {
//!         Priority::Critical => handle_emergency(msg).await?,
//!         Priority::High => handle_command(msg).await?,
//!         _ => handle_telemetry(msg).await?,
//!     }
//! }
//! # Ok(())
//! # }
//! ```

use crate::context::Receiver;
use std::collections::VecDeque;

/// Message priority levels
///
/// Higher priority messages are processed before lower priority ones.
/// Within the same priority level, messages are processed FIFO.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
pub enum Priority {
    /// Lowest priority - background tasks, non-critical telemetry
    Low = 0,
    /// Normal priority - standard operations
    #[default]
    Normal = 1,
    /// High priority - time-sensitive commands
    High = 2,
    /// Highest priority - safety-critical, emergency stop
    Critical = 3,
}

/// Configuration for priority queue behavior
#[derive(Debug, Clone)]
pub struct PriorityQueueConfig {
    /// Maximum messages to buffer per priority level
    /// None = unbounded (risk of memory growth)
    pub max_per_priority: Option<usize>,

    /// Enable starvation prevention
    /// If true, occasionally process lower priority messages even when
    /// higher priority messages are available
    pub prevent_starvation: bool,

    /// Starvation threshold: after processing N high-priority messages,
    /// process one lower-priority message (if available)
    pub starvation_threshold: usize,
}

impl Default for PriorityQueueConfig {
    fn default() -> Self {
        Self {
            max_per_priority: Some(1000),
            prevent_starvation: true,
            starvation_threshold: 100,
        }
    }
}

/// Priority receiver that processes messages by priority
///
/// Combines multiple receivers and delivers messages in priority order.
/// Messages with higher priority are processed before lower priority ones.
pub struct PriorityReceiver<T> {
    sources: Vec<PrioritySource<T>>,
    config: PriorityQueueConfig,
    high_priority_count: usize,
}

struct PrioritySource<T> {
    receiver: Receiver<T>,
    priority: Priority,
    buffer: VecDeque<T>,
}

impl<T> PriorityReceiver<T>
where
    T: Send + 'static,
{
    /// Create a new priority receiver with default configuration
    pub fn new() -> Self {
        Self::with_config(PriorityQueueConfig::default())
    }

    /// Create a new priority receiver with custom configuration
    pub fn with_config(config: PriorityQueueConfig) -> Self {
        Self {
            sources: Vec::new(),
            config,
            high_priority_count: 0,
        }
    }

    /// Add a message source with the specified priority
    ///
    /// # Example
    ///
    /// ```rust
    /// # use mecha10::prelude::*;
    /// # use mecha10::priority_queue::{PriorityReceiver, Priority};
    /// # async fn example(ctx: &Context) -> Result<()> {
    /// let mut prio_rx = PriorityReceiver::new();
    ///
    /// let emergency = ctx.subscribe("/safety/emergency").await?;
    /// prio_rx.add_source(emergency, Priority::Critical);
    ///
    /// let telemetry = ctx.subscribe("/sensor/telemetry").await?;
    /// prio_rx.add_source(telemetry, Priority::Low);
    /// # Ok(())
    /// # }
    /// ```
    pub fn add_source(&mut self, receiver: Receiver<T>, priority: Priority) {
        self.sources.push(PrioritySource {
            receiver,
            priority,
            buffer: VecDeque::new(),
        });
    }

    /// Receive the next message in priority order
    ///
    /// Returns `None` when all sources are closed.
    /// Returns `Some((message, priority))` with the highest priority message available.
    ///
    /// # Starvation Prevention
    ///
    /// If configured, this method will occasionally process lower-priority messages
    /// even when higher-priority messages are available, to prevent starvation.
    pub async fn recv(&mut self) -> Option<(T, Priority)> {
        if self.sources.is_empty() {
            return None;
        }

        loop {
            // Try to fill buffers from all sources (non-blocking)
            self.try_fill_buffers();

            // Check if we should process a lower-priority message for starvation prevention
            if self.config.prevent_starvation && self.high_priority_count >= self.config.starvation_threshold {
                // Try to get a lower-priority message
                if let Some((msg, prio)) = self.take_lower_priority_message() {
                    self.high_priority_count = 0;
                    return Some((msg, prio));
                }
            }

            // Get highest priority message from buffers
            if let Some((msg, prio)) = self.take_highest_priority_message() {
                if prio >= Priority::High {
                    self.high_priority_count += 1;
                }
                return Some((msg, prio));
            }

            // No messages in buffers, wait for any source to have data
            if !self.wait_for_message().await {
                // All sources closed
                return None;
            }
        }
    }

    /// Try to receive messages from all sources without blocking
    fn try_fill_buffers(&mut self) {
        for source in &mut self.sources {
            // Respect buffer size limits
            let max_size = self.config.max_per_priority.unwrap_or(usize::MAX);

            while source.buffer.len() < max_size {
                match source.receiver.try_recv() {
                    Ok(msg) => source.buffer.push_back(msg),
                    Err(_) => break, // Empty or closed
                }
            }
        }
    }

    /// Wait for at least one message from any source
    async fn wait_for_message(&mut self) -> bool {
        if self.sources.is_empty() {
            return false;
        }

        // Use tokio::select! to wait on all sources simultaneously
        // This is efficient and doesn't poll
        loop {
            for source in &mut self.sources {
                // Try each source
                if let Some(msg) = source.receiver.recv().await {
                    source.buffer.push_back(msg);
                    return true;
                }
            }

            // Remove closed sources
            self.sources.retain(|s| !s.buffer.is_empty() || !s.receiver.is_empty());

            if self.sources.is_empty() {
                return false;
            }
        }
    }

    /// Take the highest priority message from buffers
    fn take_highest_priority_message(&mut self) -> Option<(T, Priority)> {
        // Find source with highest priority that has messages
        let mut best_idx = None;
        let mut best_priority = Priority::Low;

        for (idx, source) in self.sources.iter().enumerate() {
            if !source.buffer.is_empty() && source.priority > best_priority {
                best_idx = Some(idx);
                best_priority = source.priority;
            }
        }

        best_idx.and_then(|idx| {
            self.sources[idx]
                .buffer
                .pop_front()
                .map(|msg| (msg, self.sources[idx].priority))
        })
    }

    /// Take a lower-priority message for starvation prevention
    fn take_lower_priority_message(&mut self) -> Option<(T, Priority)> {
        // Find source with lowest priority that has messages
        let mut best_idx = None;
        let mut best_priority = Priority::Critical;

        for (idx, source) in self.sources.iter().enumerate() {
            if !source.buffer.is_empty() && source.priority < best_priority {
                best_idx = Some(idx);
                best_priority = source.priority;
            }
        }

        best_idx.and_then(|idx| {
            self.sources[idx]
                .buffer
                .pop_front()
                .map(|msg| (msg, self.sources[idx].priority))
        })
    }

    /// Get statistics about the priority queue
    pub fn stats(&self) -> PriorityQueueStats {
        let mut counts_by_priority = [0; 4];
        let mut total_buffered = 0;

        for source in &self.sources {
            let idx = source.priority as usize;
            counts_by_priority[idx] += source.buffer.len();
            total_buffered += source.buffer.len();
        }

        PriorityQueueStats {
            total_buffered,
            critical_buffered: counts_by_priority[Priority::Critical as usize],
            high_buffered: counts_by_priority[Priority::High as usize],
            normal_buffered: counts_by_priority[Priority::Normal as usize],
            low_buffered: counts_by_priority[Priority::Low as usize],
            source_count: self.sources.len(),
        }
    }
}

impl<T> Default for PriorityReceiver<T>
where
    T: Send + 'static,
{
    fn default() -> Self {
        Self::new()
    }
}

/// Statistics about priority queue state
#[derive(Debug, Clone)]
pub struct PriorityQueueStats {
    /// Total messages buffered across all priorities
    pub total_buffered: usize,
    /// Messages buffered at Critical priority
    pub critical_buffered: usize,
    /// Messages buffered at High priority
    pub high_buffered: usize,
    /// Messages buffered at Normal priority
    pub normal_buffered: usize,
    /// Messages buffered at Low priority
    pub low_buffered: usize,
    /// Number of active sources
    pub source_count: usize,
}

impl PriorityQueueStats {
    /// Check if any priority level is approaching its limit
    pub fn is_near_limit(&self, config: &PriorityQueueConfig) -> bool {
        if let Some(max) = config.max_per_priority {
            let threshold = (max as f32 * 0.8) as usize;
            self.critical_buffered > threshold
                || self.high_buffered > threshold
                || self.normal_buffered > threshold
                || self.low_buffered > threshold
        } else {
            false
        }
    }
}