sz-orm-pool 1.2.2

//! # 限流器（Rate Limiter）
//!
//! P1-4 修复：将限流器抽象从扩展层（sz-orm-limit）提升到核心层，
//! 消除 core → limit 的反向依赖。
//!
//! ## 设计
//!
//! - [`RateLimiter`] trait：限流器抽象接口，核心层定义
//! - [`RateLimitResult`]：限流判定结果
//! - [`RateLimitError`]：限流器内部错误
//!
//! 扩展包 sz-orm-limit 提供多种具体实现（令牌桶/滑动窗口/固定窗口/分布式），
//! 但必须实现核心 trait，保证核心层不依赖扩展层。
//!
//! 连接池 `Pool` 仅依赖核心 trait，调用方通过 `set_rate_limiter`
//! 注入任意实现（sz-orm-limit 提供的或自定义的）。

/// 限流判定结果
///
/// P1-4：从 sz-orm-limit 迁移到核心层，作为限流器抽象的一部分。
#[derive(Debug, Clone)]
pub struct RateLimitResult {
    /// 是否放行
    pub allowed: bool,
    /// 当前窗口剩余配额
    pub remaining: u64,
    /// 窗口重置时间戳（毫秒）
    pub reset_at: i64,
}

impl RateLimitResult {
    /// 构造“放行”结果
    ///
    /// - `remaining`：剩余配额
    /// - `reset_at`：窗口重置时间戳（毫秒）
    pub fn allowed(remaining: u64, reset_at: i64) -> Self {
        Self {
            allowed: true,
            remaining,
            reset_at,
        }
    }

    /// 构造“拒绝”结果
    ///
    /// - `remaining`：剩余配额（通常为 0）
    /// - `reset_at`：窗口重置时间戳（毫秒）
    pub fn rejected(remaining: u64, reset_at: i64) -> Self {
        Self {
            allowed: false,
            remaining,
            reset_at,
        }
    }
}

/// 限流器错误
///
/// P1-4：从 sz-orm-limit 迁移到核心层。
#[derive(Debug)]
pub enum RateLimitError {
    /// 指定 key 未找到
    KeyNotFound(String),
    /// 限流器内部错误（如锁中毒、后端不可用）
    Internal(String),
}

impl std::fmt::Display for RateLimitError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            RateLimitError::KeyNotFound(key) => write!(f, "Key not found: {}", key),
            RateLimitError::Internal(msg) => write!(f, "Internal error: {}", msg),
        }
    }
}

impl std::error::Error for RateLimitError {}

impl serde::Serialize for RateLimitError {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        serializer.serialize_str(&self.to_string())
    }
}

/// 限流器抽象 trait
///
/// P1-4：核心层定义的限流器接口，消除对 sz-orm-limit 的反向依赖。
///
/// 实现方需保证线程安全（`Send + Sync`），以便在连接池中作为
/// `Arc<RwLock<Option<Arc<dyn RateLimiter>>>>` 使用。
///
/// 常见实现（位于 sz-orm-limit 扩展包）：
/// - `SlidingWindowRateLimiter`：滑动窗口
/// - `TokenBucketRateLimiter`：令牌桶
/// - `FixedWindowRateLimiter`：固定窗口
/// - `DistributedRateLimiter`：分布式限流（基于 Redis 等共享存储）
pub trait RateLimiter: Send + Sync {
    /// 获取一个令牌（阻塞语义：可能等待）
    ///
    /// 返回 [`RateLimitResult`] 表示放行或拒绝，或 [`RateLimitError`] 表示内部错误。
    fn acquire(&self, key: &str) -> Result<RateLimitResult, RateLimitError>;

    /// 尝试获取一个令牌（非阻塞语义：立即返回）
    ///
    /// 默认实现委托给 `acquire`；具体实现可覆盖以提供真正的非阻塞路径。
    fn try_acquire(&self, key: &str) -> Result<RateLimitResult, RateLimitError> {
        self.acquire(key)
    }

    /// 重置指定 key 的限流状态
    fn reset(&self, key: &str) -> Result<(), RateLimitError>;
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicU64, Ordering};
    use std::sync::Arc;

    #[test]
    fn test_rate_limit_result_allowed() {
        let result = RateLimitResult::allowed(5, 1000);
        assert!(result.allowed);
        assert_eq!(result.remaining, 5);
        assert_eq!(result.reset_at, 1000);
    }

    #[test]
    fn test_rate_limit_result_rejected() {
        let result = RateLimitResult::rejected(0, 2000);
        assert!(!result.allowed);
        assert_eq!(result.remaining, 0);
    }

    #[test]
    fn test_rate_limit_error_display() {
        let e = RateLimitError::KeyNotFound("user-1".to_string());
        assert!(format!("{}", e).contains("user-1"));
        let e = RateLimitError::Internal("lock poisoned".to_string());
        assert!(format!("{}", e).contains("lock poisoned"));
    }

    /// 简单计数器限流器（仅用于测试 trait 抽象可用性）
    struct CounterLimiter {
        max: u64,
        count: AtomicU64,
    }

    impl CounterLimiter {
        fn new(max: u64) -> Self {
            Self {
                max,
                count: AtomicU64::new(0),
            }
        }
    }

    impl RateLimiter for CounterLimiter {
        fn acquire(&self, _key: &str) -> Result<RateLimitResult, RateLimitError> {
            let prev = self.count.fetch_add(1, Ordering::SeqCst);
            if prev < self.max {
                Ok(RateLimitResult::allowed(self.max - prev - 1, 0))
            } else {
                Ok(RateLimitResult::rejected(0, 0))
            }
        }

        fn reset(&self, _key: &str) -> Result<(), RateLimitError> {
            self.count.store(0, Ordering::SeqCst);
            Ok(())
        }
    }

    #[test]
    fn test_counter_limiter_via_trait() {
        let limiter: Arc<dyn RateLimiter> = Arc::new(CounterLimiter::new(2));
        let r1 = limiter.acquire("k").unwrap();
        assert!(r1.allowed);
        let r2 = limiter.acquire("k").unwrap();
        assert!(r2.allowed);
        let r3 = limiter.acquire("k").unwrap();
        assert!(!r3.allowed);
    }

    #[test]
    fn test_counter_limiter_reset() {
        let limiter: Arc<dyn RateLimiter> = Arc::new(CounterLimiter::new(1));
        assert!(limiter.acquire("k").unwrap().allowed);
        assert!(!limiter.acquire("k").unwrap().allowed);
        limiter.reset("k").unwrap();
        assert!(limiter.acquire("k").unwrap().allowed);
    }

    #[test]
    fn test_default_try_acquire_delegates_to_acquire() {
        let limiter: Arc<dyn RateLimiter> = Arc::new(CounterLimiter::new(1));
        // 默认 try_acquire 应委托给 acquire
        let r1 = limiter.try_acquire("k").unwrap();
        assert!(r1.allowed);
        let r2 = limiter.try_acquire("k").unwrap();
        assert!(!r2.allowed);
    }

    #[test]
    fn test_rate_limiter_is_send_sync() {
        fn assert_send_sync<T: Send + Sync>() {}
        assert_send_sync::<Arc<dyn RateLimiter>>();
    }
}