sz_orm_core/
rate_limiter.rs

1//! 限流器抽象（P1-4：抽象提升到核心层，供连接池执行路径集成）
2//!
3//! 连接池在 `rate-limit` feature 下通过 `set_rate_limiter()` 配置限流器，
4//! 获取连接前调用 [`RateLimiter::try_acquire`] 检查是否放行。
5//!
6//! 本模块为自包含抽象（不依赖 sz-orm-limit），消除核心层对上层 crate 的反向依赖；
7//! 具体实现（滑动窗口/令牌桶等）由 sz-orm-limit 提供并实现本 trait。
8
9use std::fmt;
10
11/// 限流判定结果。
12#[derive(Debug, Clone)]
13pub struct RateLimitResult {
14    /// 是否放行。
15    pub allowed: bool,
16    /// 放行后剩余配额。
17    pub remaining: u64,
18    /// 配额重置时间戳（Unix 毫秒）。
19    pub reset_at: i64,
20}
21
22impl RateLimitResult {
23    /// 放行结果。
24    pub fn allowed(remaining: u64, reset_at: i64) -> Self {
25        Self {
26            allowed: true,
27            remaining,
28            reset_at,
29        }
30    }
31
32    /// 拒绝结果。
33    pub fn rejected(remaining: u64, reset_at: i64) -> Self {
34        Self {
35            allowed: false,
36            remaining,
37            reset_at,
38        }
39    }
40}
41
42/// 限流器内部错误（如后端不可用）。
43///
44/// 连接池对内部错误采取保守放行策略（避免误杀正常请求）。
45#[derive(Debug, Clone, PartialEq, Eq)]
46pub struct RateLimitError(pub String);
47
48impl fmt::Display for RateLimitError {
49    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
50        write!(f, "rate limiter error: {}", self.0)
51    }
52}
53
54impl std::error::Error for RateLimitError {}
55
56/// 限流器抽象 trait。
57pub trait RateLimiter: Send + Sync {
58    /// 尝试获取 `key` 的配额。`Ok(result)` 中 `result.allowed` 为 `false` 时拒绝放行。
59    fn try_acquire(&self, key: &str) -> Result<RateLimitResult, RateLimitError>;
60}
61
62#[cfg(test)]
63mod tests {
64    use super::*;
65
66    #[test]
67    fn test_rate_limit_result_allowed() {
68        let r = RateLimitResult::allowed(9, 1_700_000_000_000);
69        assert!(r.allowed);
70        assert_eq!(r.remaining, 9);
71    }
72
73    #[test]
74    fn test_rate_limit_result_rejected() {
75        let r = RateLimitResult::rejected(0, 1_700_000_000_000);
76        assert!(!r.allowed);
77        assert_eq!(r.remaining, 0);
78    }
79
80    #[test]
81    fn test_rate_limit_error_display() {
82        let e = RateLimitError("backend down".to_string());
83        assert!(e.to_string().contains("backend down"));
84    }
85
86    /// 桩实现：验证 trait 可被外部实现并驱动 try_acquire。
87    struct AlwaysAllow;
88
89    impl RateLimiter for AlwaysAllow {
90        fn try_acquire(&self, _key: &str) -> Result<RateLimitResult, RateLimitError> {
91            Ok(RateLimitResult::allowed(1, 0))
92        }
93    }
94
95    #[test]
96    fn test_trait_dispatch() {
97        let limiter: &dyn RateLimiter = &AlwaysAllow;
98        let r = limiter.try_acquire("test-key").unwrap();
99        assert!(r.allowed);
100    }
101}
sz_orm_core/rate_limiter.rs

sz_orm_core/
rate_limiter.rs
