mecha10-core 0.6.5

//! 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_core::prelude::*;
//! use mecha10_core::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::Instant;
use tokio::sync::RwLock;

/// 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_core::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_core::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
        }
    }

    /// 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 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,
        }
    }
}

/// Rate limiter that tracks limits per key (topic, node, etc.)
///
/// # Example
///
/// ```rust
/// use mecha10_core::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);
    }

    /// 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 {
            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 {
        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 {
            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
        }
    }

    // duplicate: false positive — see the note on `CircuitBreakerRegistry::reset_all` in
    // rpc_breaker.rs.
    /// 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()
    }
}

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

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

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

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