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_core::prelude::*;
//! use mecha10_core::priority_queue::{PriorityReceiver, Priority};
//! use mecha10_core::topics::Topic;
//!
//! #[derive(Debug, Deserialize)]
//! struct Event {
//!     value: f32,
//! }
//! impl Message for Event {}
//!
//! const EMERGENCY: Topic<Event> = Topic::new("/safety/emergency");
//! const COMMAND: Topic<Event> = Topic::new("/motor/command");
//! const TELEMETRY: Topic<Event> = Topic::new("/sensor/telemetry");
//!
//! # async fn handle_emergency(_msg: Event) -> Result<()> { Ok(()) }
//! # async fn handle_command(_msg: Event) -> Result<()> { Ok(()) }
//! # async fn handle_telemetry(_msg: Event) -> Result<()> { Ok(()) }
//!
//! # 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(EMERGENCY).await?;
//! let commands = ctx.subscribe(COMMAND).await?;
//! let telemetry = ctx.subscribe(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;

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
pub enum Priority {
    Low = 0,
    #[default]
    Normal = 1,
    High = 2,
    Critical = 3,
}

#[derive(Debug, Clone)]
pub struct PriorityQueueConfig {
    pub max_per_priority: Option<usize>,

    pub prevent_starvation: bool,

    pub starvation_threshold: usize,
}

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

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_core::prelude::*;
    /// # use mecha10_core::priority_queue::{PriorityReceiver, Priority};
    /// # use mecha10_core::topics::Topic;
    /// # #[derive(Debug, Deserialize)]
    /// # struct Event { value: f32 }
    /// # impl Message for Event {}
    /// # const EMERGENCY: Topic<Event> = Topic::new("/safety/emergency");
    /// # const TELEMETRY: Topic<Event> = Topic::new("/sensor/telemetry");
    /// # async fn example(ctx: &Context) -> Result<()> {
    /// let mut prio_rx = PriorityReceiver::new();
    ///
    /// let emergency = ctx.subscribe(EMERGENCY).await?;
    /// prio_rx.add_source(emergency, Priority::Critical);
    ///
    /// let telemetry = ctx.subscribe(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 {
            self.try_fill_buffers();

            if self.config.prevent_starvation && self.high_priority_count >= self.config.starvation_threshold {
                if let Some((msg, prio)) = self.take_lower_priority_message() {
                    self.high_priority_count = 0;
                    return Some((msg, prio));
                }
            }

            if let Some((msg, prio)) = self.take_highest_priority_message() {
                if prio >= Priority::High {
                    self.high_priority_count += 1;
                }
                return Some((msg, prio));
            }

            if !self.wait_for_message().await {
                return None;
            }
        }
    }

    /// Try to receive messages from all sources without blocking
    fn try_fill_buffers(&mut self) {
        for source in &mut self.sources {
            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,
                }
            }
        }
    }

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

        loop {
            for source in &mut self.sources {
                if let Some(msg) = source.receiver.recv().await {
                    source.buffer.push_back(msg);
                    return true;
                }
            }

            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)> {
        self.take_extreme_priority_message(|candidate, best| candidate > best)
    }

    /// Take a lower-priority message for starvation prevention
    fn take_lower_priority_message(&mut self) -> Option<(T, Priority)> {
        self.take_extreme_priority_message(|candidate, best| candidate < best)
    }

    /// Shared implementation behind `take_highest_priority_message`/`take_lower_priority_message`,
    /// which differ only in which direction `is_better` compares priorities.
    fn take_extreme_priority_message(
        &mut self,
        is_better: impl Fn(Priority, Priority) -> bool,
    ) -> Option<(T, Priority)> {
        let mut best: Option<(usize, Priority)> = None;

        for (idx, source) in self.sources.iter().enumerate() {
            if source.buffer.is_empty() {
                continue;
            }
            match best {
                None => best = Some((idx, source.priority)),
                Some((_, best_priority)) if is_better(source.priority, best_priority) => {
                    best = Some((idx, source.priority));
                }
                _ => {}
            }
        }

        best.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()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_source(priority: Priority, messages: Vec<u32>) -> PrioritySource<u32> {
        let (_tx, rx) = tokio::sync::mpsc::unbounded_channel();
        PrioritySource {
            receiver: Receiver::from(rx),
            priority,
            buffer: messages.into_iter().collect(),
        }
    }

    /// Regression test: when every buffered message is `Priority::Low` (the
    /// sentinel `take_highest_priority_message` used to seed `best_priority`
    /// with), the strict `>` comparison used to never select any of them,
    /// even though they were the only messages present. This starved those
    /// sources forever in `recv()`, since the underlying channel had already
    /// been drained into the buffer via `try_recv` and `wait_for_message`
    /// only awaits *new* channel messages, not buffered ones.
    #[test]
    fn take_highest_priority_message_selects_all_low_messages() {
        let mut prio_rx: PriorityReceiver<u32> = PriorityReceiver::new();
        prio_rx.sources.push(make_source(Priority::Low, vec![1, 2, 3]));

        let (msg, priority) = prio_rx
            .take_highest_priority_message()
            .expect("a Low-priority message should still be returned when it's the only one buffered");
        assert_eq!(msg, 1);
        assert_eq!(priority, Priority::Low);

        assert_eq!(prio_rx.take_highest_priority_message(), Some((2, Priority::Low)));
        assert_eq!(prio_rx.take_highest_priority_message(), Some((3, Priority::Low)));
        assert_eq!(prio_rx.take_highest_priority_message(), None);
    }

    /// Regression test: mirror of the above for `take_lower_priority_message`,
    /// whose sentinel was `Priority::Critical` compared with strict `<`.
    #[test]
    fn take_lower_priority_message_selects_all_critical_messages() {
        let mut prio_rx: PriorityReceiver<u32> = PriorityReceiver::new();
        prio_rx.sources.push(make_source(Priority::Critical, vec![10, 20]));

        let (msg, priority) = prio_rx
            .take_lower_priority_message()
            .expect("a Critical-priority message should still be returned when it's the only one buffered");
        assert_eq!(msg, 10);
        assert_eq!(priority, Priority::Critical);

        assert_eq!(prio_rx.take_lower_priority_message(), Some((20, Priority::Critical)));
        assert_eq!(prio_rx.take_lower_priority_message(), None);
    }

    /// End-to-end regression test through `recv()`: a single source producing
    /// only `Priority::Low` messages must have those messages delivered
    /// rather than hanging forever waiting on `wait_for_message`.
    #[tokio::test]
    async fn recv_delivers_messages_from_a_low_priority_only_source() {
        let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<u32>();
        let mut prio_rx: PriorityReceiver<u32> = PriorityReceiver::new();
        prio_rx.add_source(Receiver::from(rx), Priority::Low);

        tx.send(42).unwrap();
        drop(tx);

        let result = tokio::time::timeout(std::time::Duration::from_secs(5), prio_rx.recv())
            .await
            .expect("recv() should not hang starving a Low-only source");
        assert_eq!(result, Some((42, Priority::Low)));
    }

    /// take_highest_priority_message should still prefer higher priorities
    /// over lower ones when both are present.
    #[test]
    fn take_highest_priority_message_prefers_higher_priority() {
        let mut prio_rx: PriorityReceiver<u32> = PriorityReceiver::new();
        prio_rx.sources.push(make_source(Priority::Low, vec![1]));
        prio_rx.sources.push(make_source(Priority::Critical, vec![2]));

        assert_eq!(prio_rx.take_highest_priority_message(), Some((2, Priority::Critical)));
    }
}

#[derive(Debug, Clone)]
pub struct PriorityQueueStats {
    pub total_buffered: usize,
    pub critical_buffered: usize,
    pub high_buffered: usize,
    pub normal_buffered: usize,
    pub low_buffered: usize,
    pub source_count: usize,
}