mecha10_core/

rate_limit.rs

//! Rate Limiting
//!
//! Per-topic and per-node rate limiting to prevent message flooding and ensure QoS.
//!
//! # Features
//!
//! - Token bucket algorithm for smooth rate limiting
//! - Per-topic limits
//! - Per-node limits
//! - Burst capacity
//! - Configurable policies
//! - Metrics integration
//! - Low overhead with async
//!
//! # Example
//!
//! ```rust
//! use mecha10::prelude::*;
//! use mecha10::rate_limit::{RateLimiter, RateLimit};
//!
//! # async fn example() -> Result<()> {
//! // Create rate limiter: 10 messages/sec with burst of 20
//! let limiter = RateLimiter::new(10.0, 20);
//!
//! // Check if message is allowed
//! if limiter.check_and_consume().await {
//!     // Allowed - publish message
//!     println!("Message allowed");
//! } else {
//!     // Rate limited - drop or queue
//!     println!("Rate limited!");
//! }
//! # Ok(())
//! # }
//! ```

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;

// ============================================================================
// Token Bucket Rate Limiter
// ============================================================================

/// Token bucket rate limiter
///
/// Implements the token bucket algorithm for smooth rate limiting with burst support.
///
/// # Algorithm
///
/// - Tokens are added at a constant rate (refill_rate)
/// - Each message consumes one token
/// - Bucket has maximum capacity (burst capacity)
/// - If no tokens available, request is denied
///
/// # Example
///
/// ```rust
/// use mecha10::rate_limit::RateLimiter;
///
/// # async fn example() {
/// // 10 msg/sec, burst of 20
/// let limiter = RateLimiter::new(10.0, 20);
///
/// // Check if allowed
/// if limiter.check_and_consume().await {
///     println!("Allowed");
/// }
/// # }
/// ```
pub struct RateLimiter {
    /// Current number of tokens
    tokens: Arc<RwLock<f64>>,
    /// Maximum tokens (burst capacity)
    capacity: f64,
    /// Tokens added per second
    refill_rate: f64,
    /// Last refill time
    last_refill: Arc<RwLock<Instant>>,
}

impl RateLimiter {
    /// Create a new rate limiter
    ///
    /// # Arguments
    ///
    /// * `rate` - Messages per second
    /// * `burst` - Maximum burst size
    ///
    /// # Example
    ///
    /// ```rust
    /// use mecha10::rate_limit::RateLimiter;
    ///
    /// // 100 msg/sec with burst of 200
    /// let limiter = RateLimiter::new(100.0, 200);
    /// ```
    pub fn new(rate: f64, burst: usize) -> Self {
        Self {
            tokens: Arc::new(RwLock::new(burst as f64)),
            capacity: burst as f64,
            refill_rate: rate,
            last_refill: Arc::new(RwLock::new(Instant::now())),
        }
    }

    /// Check if request is allowed (does not consume token)
    pub async fn check(&self) -> bool {
        self.refill().await;
        let tokens = self.tokens.read().await;
        *tokens >= 1.0
    }

    /// Check and consume a token if available
    pub async fn check_and_consume(&self) -> bool {
        self.refill().await;

        let mut tokens = self.tokens.write().await;
        if *tokens >= 1.0 {
            *tokens -= 1.0;
            true
        } else {
            false
        }
    }

    /// Try to consume multiple tokens
    pub async fn try_consume(&self, count: f64) -> bool {
        self.refill().await;

        let mut tokens = self.tokens.write().await;
        if *tokens >= count {
            *tokens -= count;
            true
        } else {
            false
        }
    }

    /// Wait until a token is available (with timeout)
    ///
    /// Returns true if token was acquired, false if timed out
    pub async fn wait_for_token(&self, timeout: Duration) -> bool {
        let start = Instant::now();

        while start.elapsed() < timeout {
            if self.check_and_consume().await {
                return true;
            }

            // Sleep for a short time before retrying
            tokio::time::sleep(Duration::from_millis(10)).await;
        }

        false
    }

    /// Get current token count
    pub async fn available_tokens(&self) -> f64 {
        self.refill().await;
        *self.tokens.read().await
    }

    /// Refill tokens based on elapsed time
    async fn refill(&self) {
        let now = Instant::now();
        let mut last_refill = self.last_refill.write().await;
        let elapsed = now.duration_since(*last_refill).as_secs_f64();

        if elapsed > 0.0 {
            let tokens_to_add = elapsed * self.refill_rate;
            let mut tokens = self.tokens.write().await;
            *tokens = (*tokens + tokens_to_add).min(self.capacity);
            *last_refill = now;
        }
    }

    /// Reset the rate limiter (fill bucket)
    pub async fn reset(&self) {
        let mut tokens = self.tokens.write().await;
        *tokens = self.capacity;
        let mut last_refill = self.last_refill.write().await;
        *last_refill = Instant::now();
    }
}

// ============================================================================
// Rate Limit Policy
// ============================================================================

/// Rate limit policy configuration
#[derive(Debug, Clone)]
pub struct RateLimit {
    /// Messages per second
    pub rate: f64,
    /// Burst capacity
    pub burst: usize,
}

impl RateLimit {
    /// Create a new rate limit policy
    pub fn new(rate: f64, burst: usize) -> Self {
        Self { rate, burst }
    }

    /// Create rate limit from messages per second (burst = rate * 2)
    pub fn per_second(rate: f64) -> Self {
        Self {
            rate,
            burst: (rate * 2.0) as usize,
        }
    }

    /// Create rate limit from messages per minute
    pub fn per_minute(rate: f64) -> Self {
        let per_sec = rate / 60.0;
        Self {
            rate: per_sec,
            burst: (per_sec * 2.0) as usize,
        }
    }

    /// No limit
    pub fn unlimited() -> Self {
        Self {
            rate: f64::MAX,
            burst: usize::MAX,
        }
    }
}

// ============================================================================
// Multi-Key Rate Limiter
// ============================================================================

/// Rate limiter that tracks limits per key (topic, node, etc.)
///
/// # Example
///
/// ```rust
/// use mecha10::rate_limit::{MultiKeyRateLimiter, RateLimit};
///
/// # async fn example() {
/// let limiter = MultiKeyRateLimiter::new();
///
/// // Set limits per topic
/// limiter.set_limit("/camera/rgb", RateLimit::per_second(30.0)).await;
/// limiter.set_limit("/camera/depth", RateLimit::per_second(30.0)).await;
///
/// // Check limits
/// if limiter.check_and_consume("/camera/rgb").await {
///     println!("Allowed");
/// }
/// # }
/// ```
pub struct MultiKeyRateLimiter {
    limiters: Arc<RwLock<HashMap<String, RateLimiter>>>,
    default_limit: RateLimit,
}

impl MultiKeyRateLimiter {
    /// Create a new multi-key rate limiter with default unlimited
    pub fn new() -> Self {
        Self {
            limiters: Arc::new(RwLock::new(HashMap::new())),
            default_limit: RateLimit::unlimited(),
        }
    }

    /// Create with a default limit for all keys
    pub fn with_default(default: RateLimit) -> Self {
        Self {
            limiters: Arc::new(RwLock::new(HashMap::new())),
            default_limit: default,
        }
    }

    /// Set rate limit for a specific key
    pub async fn set_limit(&self, key: impl Into<String>, limit: RateLimit) {
        let key = key.into();
        let limiter = RateLimiter::new(limit.rate, limit.burst);

        let mut limiters = self.limiters.write().await;
        limiters.insert(key, limiter);
    }

    /// Remove rate limit for a key
    pub async fn remove_limit(&self, key: &str) {
        let mut limiters = self.limiters.write().await;
        limiters.remove(key);
    }

    /// Check if request for key is allowed
    pub async fn check(&self, key: &str) -> bool {
        let limiters = self.limiters.read().await;

        if let Some(limiter) = limiters.get(key) {
            limiter.check().await
        } else if self.default_limit.rate == f64::MAX {
            true
        } else {
            // Create limiter for this key on first use
            drop(limiters);
            self.set_limit(key, self.default_limit.clone()).await;
            true
        }
    }

    /// Check and consume token for key
    pub async fn check_and_consume(&self, key: &str) -> bool {
        // Get or create limiter
        let should_create = {
            let limiters = self.limiters.read().await;
            !limiters.contains_key(key) && self.default_limit.rate != f64::MAX
        };

        if should_create {
            self.set_limit(key, self.default_limit.clone()).await;
        }

        let limiters = self.limiters.read().await;

        if let Some(limiter) = limiters.get(key) {
            limiter.check_and_consume().await
        } else {
            // Unlimited
            true
        }
    }

    /// Get available tokens for a key
    pub async fn available_tokens(&self, key: &str) -> Option<f64> {
        let limiters = self.limiters.read().await;
        if let Some(limiter) = limiters.get(key) {
            Some(limiter.available_tokens().await)
        } else {
            None
        }
    }

    /// Get all configured limits
    pub async fn get_limits(&self) -> HashMap<String, RateLimit> {
        let limiters = self.limiters.read().await;
        let mut limits = HashMap::new();

        for (key, limiter) in limiters.iter() {
            limits.insert(
                key.clone(),
                RateLimit {
                    rate: limiter.refill_rate,
                    burst: limiter.capacity as usize,
                },
            );
        }

        limits
    }

    /// Clear all limiters
    pub async fn clear(&self) {
        let mut limiters = self.limiters.write().await;
        limiters.clear();
    }
}

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

// ============================================================================
// Rate Limiter for Context Integration
// ============================================================================

/// Global rate limiter manager for the system
pub struct RateLimiterManager {
    /// Per-topic rate limiters
    topic_limiters: Arc<MultiKeyRateLimiter>,
    /// Per-node rate limiters
    node_limiters: Arc<MultiKeyRateLimiter>,
}

impl RateLimiterManager {
    /// Create a new rate limiter manager
    pub fn new() -> Self {
        Self {
            topic_limiters: Arc::new(MultiKeyRateLimiter::new()),
            node_limiters: Arc::new(MultiKeyRateLimiter::new()),
        }
    }

    /// Set rate limit for a topic
    pub async fn set_topic_limit(&self, topic: impl Into<String>, limit: RateLimit) {
        self.topic_limiters.set_limit(topic, limit).await;
    }

    /// Set rate limit for a node
    pub async fn set_node_limit(&self, node_id: impl Into<String>, limit: RateLimit) {
        self.node_limiters.set_limit(node_id, limit).await;
    }

    /// Check if publishing to topic is allowed
    pub async fn check_topic(&self, topic: &str) -> bool {
        self.topic_limiters.check_and_consume(topic).await
    }

    /// Check if node can publish
    pub async fn check_node(&self, node_id: &str) -> bool {
        self.node_limiters.check_and_consume(node_id).await
    }

    /// Check both topic and node limits
    pub async fn check_publish(&self, node_id: &str, topic: &str) -> bool {
        self.check_node(node_id).await && self.check_topic(topic).await
    }

    /// Get topic limits
    pub async fn topic_limits(&self) -> HashMap<String, RateLimit> {
        self.topic_limiters.get_limits().await
    }

    /// Get node limits
    pub async fn node_limits(&self) -> HashMap<String, RateLimit> {
        self.node_limiters.get_limits().await
    }
}

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