sz_orm_limit/
config_builder.rs

1//! 限流配置构建器（Rate Limit Config Builder）
2//!
3//! 提供链式 API 构建限流器配置，支持多种算法和策略组合。
4//! 适用于从配置文件或环境变量构建限流器。
5
6use std::time::Duration;
7
8use crate::{
9    FixedWindowRateLimiter, LeakyBucketLimiter, SlidingWindowLogLimiter, SlidingWindowRateLimiter,
10    TokenBucketRateLimiter,
11};
12
13/// 限流算法类型
14#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
15pub enum LimitAlgorithm {
16    /// 固定窗口
17    FixedWindow,
18    /// 滑动窗口（计数器）
19    SlidingWindow,
20    /// 滑动窗口（日志）
21    SlidingWindowLog,
22    /// 令牌桶
23    TokenBucket,
24    /// 漏桶
25    LeakyBucket,
26}
27
28/// 限流配置
29///
30/// 描述一个限流器的完整配置，可序列化/反序列化。
31#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
32pub struct RateLimitConfig {
33    /// 算法类型
34    pub algorithm: LimitAlgorithm,
35    /// 容量/最大请求数
36    pub capacity: u64,
37    /// 补充速率（令牌桶）/ 漏出速率（漏桶），单位：请求/秒
38    pub rate: f64,
39    /// 窗口大小（毫秒），用于固定窗口和滑动窗口
40    pub window_ms: u64,
41    /// 最大 key 数量
42    pub max_keys: usize,
43}
44
45impl Default for RateLimitConfig {
46    fn default() -> Self {
47        Self {
48            algorithm: LimitAlgorithm::TokenBucket,
49            capacity: 100,
50            rate: 10.0,
51            window_ms: 1000,
52            max_keys: crate::DEFAULT_MAX_KEYS,
53        }
54    }
55}
56
57impl RateLimitConfig {
58    /// 创建默认配置
59    pub fn new() -> Self {
60        Self::default()
61    }
62
63    /// 从 JSON 字符串解析配置
64    pub fn from_json_str(s: &str) -> Result<Self, ConfigError> {
65        serde_json::from_str(s).map_err(|e| ConfigError::Parse(e.to_string()))
66    }
67
68    /// 序列化为 JSON 字符串
69    pub fn to_json_string(&self) -> Result<String, ConfigError> {
70        serde_json::to_string(self).map_err(|e| ConfigError::Serialize(e.to_string()))
71    }
72
73    /// 校验配置合理性
74    pub fn validate(&self) -> Result<(), ConfigError> {
75        if self.capacity == 0 {
76            return Err(ConfigError::InvalidCapacity);
77        }
78        match self.algorithm {
79            LimitAlgorithm::TokenBucket | LimitAlgorithm::LeakyBucket => {
80                if self.rate < 0.0 {
81                    return Err(ConfigError::InvalidRate);
82                }
83            }
84            _ => {}
85        }
86        match self.algorithm {
87            LimitAlgorithm::FixedWindow
88            | LimitAlgorithm::SlidingWindow
89            | LimitAlgorithm::SlidingWindowLog => {
90                if self.window_ms == 0 {
91                    return Err(ConfigError::InvalidWindow);
92                }
93            }
94            _ => {}
95        }
96        if self.max_keys == 0 {
97            return Err(ConfigError::InvalidMaxKeys);
98        }
99        Ok(())
100    }
101}
102
103/// 配置错误
104#[derive(Debug, Clone, PartialEq, Eq)]
105pub enum ConfigError {
106    Parse(String),
107    Serialize(String),
108    InvalidCapacity,
109    InvalidRate,
110    InvalidWindow,
111    InvalidMaxKeys,
112}
113
114impl std::fmt::Display for ConfigError {
115    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
116        match self {
117            ConfigError::Parse(msg) => write!(f, "config parse error: {}", msg),
118            ConfigError::Serialize(msg) => write!(f, "config serialize error: {}", msg),
119            ConfigError::InvalidCapacity => write!(f, "capacity must be positive"),
120            ConfigError::InvalidRate => write!(f, "rate must be non-negative"),
121            ConfigError::InvalidWindow => write!(f, "window size must be positive"),
122            ConfigError::InvalidMaxKeys => write!(f, "max_keys must be positive"),
123        }
124    }
125}
126
127impl std::error::Error for ConfigError {}
128
129/// 限流配置构建器
130///
131/// 链式 API 构建限流配置。
132///
133/// # 示例
134///
135/// ```rust
136/// use sz_orm_limit::config_builder::{RateLimitConfigBuilder, LimitAlgorithm};
137///
138/// let config = RateLimitConfigBuilder::new()
139///     .algorithm(LimitAlgorithm::TokenBucket)
140///     .capacity(100)
141///     .rate(10.0)
142///     .max_keys(5000)
143///     .build();
144/// ```
145pub struct RateLimitConfigBuilder {
146    config: RateLimitConfig,
147}
148
149impl Default for RateLimitConfigBuilder {
150    fn default() -> Self {
151        Self::new()
152    }
153}
154
155impl RateLimitConfigBuilder {
156    /// 创建构建器
157    pub fn new() -> Self {
158        Self {
159            config: RateLimitConfig::default(),
160        }
161    }
162
163    /// 设置算法类型
164    pub fn algorithm(mut self, algo: LimitAlgorithm) -> Self {
165        self.config.algorithm = algo;
166        self
167    }
168
169    /// 设置容量
170    pub fn capacity(mut self, cap: u64) -> Self {
171        self.config.capacity = cap;
172        self
173    }
174
175    /// 设置速率
176    pub fn rate(mut self, rate: f64) -> Self {
177        self.config.rate = rate;
178        self
179    }
180
181    /// 设置窗口大小（毫秒）
182    pub fn window_ms(mut self, ms: u64) -> Self {
183        self.config.window_ms = ms;
184        self
185    }
186
187    /// 设置窗口大小（秒）
188    pub fn window_secs(mut self, secs: u64) -> Self {
189        self.config.window_ms = secs * 1000;
190        self
191    }
192
193    /// 设置最大 key 数量
194    pub fn max_keys(mut self, max: usize) -> Self {
195        self.config.max_keys = max;
196        self
197    }
198
199    /// 构建配置（不校验）
200    pub fn build(self) -> RateLimitConfig {
201        self.config
202    }
203
204    /// 构建配置并校验
205    pub fn build_checked(self) -> Result<RateLimitConfig, ConfigError> {
206        self.config.validate()?;
207        Ok(self.config)
208    }
209
210    /// 从配置构建令牌桶限流器
211    pub fn build_token_bucket(&self) -> TokenBucketRateLimiter {
212        TokenBucketRateLimiter::new(self.config.capacity, self.config.rate)
213            .with_max_keys(self.config.max_keys)
214    }
215
216    /// 从配置构建滑动窗口限流器
217    pub fn build_sliding_window(&self) -> SlidingWindowRateLimiter {
218        SlidingWindowRateLimiter::new(
219            self.config.capacity,
220            Duration::from_millis(self.config.window_ms),
221        )
222        .with_max_keys(self.config.max_keys)
223    }
224
225    /// 从配置构建固定窗口限流器
226    pub fn build_fixed_window(&self) -> FixedWindowRateLimiter {
227        FixedWindowRateLimiter::new(
228            self.config.capacity,
229            Duration::from_millis(self.config.window_ms),
230        )
231        .with_max_keys(self.config.max_keys)
232    }
233
234    /// 从配置构建漏桶限流器
235    pub fn build_leaky_bucket(&self) -> LeakyBucketLimiter {
236        LeakyBucketLimiter::new(self.config.capacity, self.config.rate)
237            .with_max_keys(self.config.max_keys)
238    }
239
240    /// 从配置构建滑动窗口日志限流器
241    pub fn build_sliding_window_log(&self) -> SlidingWindowLogLimiter {
242        SlidingWindowLogLimiter::new(
243            self.config.capacity,
244            Duration::from_millis(self.config.window_ms),
245        )
246        .with_max_keys(self.config.max_keys)
247    }
248}
249
250/// 多层级限流配置
251///
252/// 为不同维度（IP、用户、API）配置不同的限流策略。
253#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
254pub struct TieredRateLimitConfig {
255    /// IP 级别配置
256    pub ip: RateLimitConfig,
257    /// 用户级别配置
258    pub user: RateLimitConfig,
259    /// API 级别配置
260    pub api: RateLimitConfig,
261    /// 全局配置
262    pub global: RateLimitConfig,
263}
264
265impl Default for TieredRateLimitConfig {
266    fn default() -> Self {
267        Self {
268            ip: RateLimitConfig {
269                algorithm: LimitAlgorithm::SlidingWindow,
270                capacity: 1000,
271                rate: 0.0,
272                window_ms: 60_000,
273                max_keys: 10_000,
274            },
275            user: RateLimitConfig {
276                algorithm: LimitAlgorithm::TokenBucket,
277                capacity: 100,
278                rate: 10.0,
279                window_ms: 1000,
280                max_keys: 10_000,
281            },
282            api: RateLimitConfig {
283                algorithm: LimitAlgorithm::FixedWindow,
284                capacity: 500,
285                rate: 0.0,
286                window_ms: 60_000,
287                max_keys: 1000,
288            },
289            global: RateLimitConfig {
290                algorithm: LimitAlgorithm::TokenBucket,
291                capacity: 10_000,
292                rate: 100.0,
293                window_ms: 1000,
294                max_keys: 1,
295            },
296        }
297    }
298}
299
300impl TieredRateLimitConfig {
301    /// 创建默认多层级配置
302    pub fn new() -> Self {
303        Self::default()
304    }
305
306    /// 校验所有层级配置
307    pub fn validate(&self) -> Result<(), ConfigError> {
308        self.ip.validate()?;
309        self.user.validate()?;
310        self.api.validate()?;
311        self.global.validate()?;
312        Ok(())
313    }
314
315    /// 从 JSON 字符串解析
316    pub fn from_json_str(s: &str) -> Result<Self, ConfigError> {
317        serde_json::from_str(s).map_err(|e| ConfigError::Parse(e.to_string()))
318    }
319
320    /// 序列化为 JSON 字符串
321    pub fn to_json_string(&self) -> Result<String, ConfigError> {
322        serde_json::to_string(self).map_err(|e| ConfigError::Serialize(e.to_string()))
323    }
324}
325
326/// 多层级配置构建器
327pub struct TieredConfigBuilder {
328    config: TieredRateLimitConfig,
329}
330
331impl Default for TieredConfigBuilder {
332    fn default() -> Self {
333        Self::new()
334    }
335}
336
337impl TieredConfigBuilder {
338    /// 创建构建器
339    pub fn new() -> Self {
340        Self {
341            config: TieredRateLimitConfig::default(),
342        }
343    }
344
345    /// 设置 IP 级别配置
346    pub fn ip(mut self, config: RateLimitConfig) -> Self {
347        self.config.ip = config;
348        self
349    }
350
351    /// 设置用户级别配置
352    pub fn user(mut self, config: RateLimitConfig) -> Self {
353        self.config.user = config;
354        self
355    }
356
357    /// 设置 API 级别配置
358    pub fn api(mut self, config: RateLimitConfig) -> Self {
359        self.config.api = config;
360        self
361    }
362
363    /// 设置全局配置
364    pub fn global(mut self, config: RateLimitConfig) -> Self {
365        self.config.global = config;
366        self
367    }
368
369    /// 构建
370    pub fn build(self) -> TieredRateLimitConfig {
371        self.config
372    }
373
374    /// 构建并校验
375    pub fn build_checked(self) -> Result<TieredRateLimitConfig, ConfigError> {
376        self.config.validate()?;
377        Ok(self.config)
378    }
379}
380
381#[cfg(test)]
382mod tests {
383    use super::*;
384    use crate::RateLimiter;
385
386    #[test]
387    fn test_limit_algorithm_serde() {
388        let algo = LimitAlgorithm::TokenBucket;
389        let json = serde_json::to_string(&algo).unwrap();
390        let back: LimitAlgorithm = serde_json::from_str(&json).unwrap();
391        assert_eq!(algo, back);
392    }
393
394    #[test]
395    fn test_rate_limit_config_default() {
396        let config = RateLimitConfig::default();
397        assert_eq!(config.algorithm, LimitAlgorithm::TokenBucket);
398        assert_eq!(config.capacity, 100);
399        assert_eq!(config.rate, 10.0);
400    }
401
402    #[test]
403    fn test_rate_limit_config_validate_ok() {
404        let config = RateLimitConfig {
405            algorithm: LimitAlgorithm::TokenBucket,
406            capacity: 100,
407            rate: 10.0,
408            window_ms: 1000,
409            max_keys: 1000,
410        };
411        assert!(config.validate().is_ok());
412    }
413
414    #[test]
415    fn test_rate_limit_config_validate_zero_capacity() {
416        let config = RateLimitConfig {
417            capacity: 0,
418            ..Default::default()
419        };
420        assert_eq!(config.validate(), Err(ConfigError::InvalidCapacity));
421    }
422
423    #[test]
424    fn test_rate_limit_config_validate_negative_rate() {
425        let config = RateLimitConfig {
426            algorithm: LimitAlgorithm::TokenBucket,
427            rate: -1.0,
428            ..Default::default()
429        };
430        assert_eq!(config.validate(), Err(ConfigError::InvalidRate));
431    }
432
433    #[test]
434    fn test_rate_limit_config_validate_zero_window() {
435        let config = RateLimitConfig {
436            algorithm: LimitAlgorithm::FixedWindow,
437            window_ms: 0,
438            ..Default::default()
439        };
440        assert_eq!(config.validate(), Err(ConfigError::InvalidWindow));
441    }
442
443    #[test]
444    fn test_rate_limit_config_validate_zero_max_keys() {
445        let config = RateLimitConfig {
446            max_keys: 0,
447            ..Default::default()
448        };
449        assert_eq!(config.validate(), Err(ConfigError::InvalidMaxKeys));
450    }
451
452    #[test]
453    fn test_rate_limit_config_json_roundtrip() {
454        let config = RateLimitConfig::default();
455        let json = config.to_json_string().unwrap();
456        let back = RateLimitConfig::from_json_str(&json).unwrap();
457        assert_eq!(config.algorithm, back.algorithm);
458        assert_eq!(config.capacity, back.capacity);
459    }
460
461    #[test]
462    fn test_config_builder_basic() {
463        let config = RateLimitConfigBuilder::new()
464            .algorithm(LimitAlgorithm::SlidingWindow)
465            .capacity(200)
466            .window_secs(60)
467            .max_keys(5000)
468            .build();
469        assert_eq!(config.algorithm, LimitAlgorithm::SlidingWindow);
470        assert_eq!(config.capacity, 200);
471        assert_eq!(config.window_ms, 60_000);
472        assert_eq!(config.max_keys, 5000);
473    }
474
475    #[test]
476    fn test_config_builder_checked_ok() {
477        let config = RateLimitConfigBuilder::new()
478            .algorithm(LimitAlgorithm::TokenBucket)
479            .capacity(100)
480            .rate(10.0)
481            .build_checked();
482        assert!(config.is_ok());
483    }
484
485    #[test]
486    fn test_config_builder_checked_fail() {
487        let config = RateLimitConfigBuilder::new().capacity(0).build_checked();
488        assert!(config.is_err());
489    }
490
491    #[test]
492    fn test_config_builder_build_token_bucket() {
493        let builder = RateLimitConfigBuilder::new()
494            .algorithm(LimitAlgorithm::TokenBucket)
495            .capacity(10)
496            .rate(1.0)
497            .max_keys(100);
498        let limiter = builder.build_token_bucket();
499        assert_eq!(limiter.capacity(), 10);
500        let r = limiter.acquire("k").unwrap();
501        assert!(r.allowed);
502    }
503
504    #[test]
505    fn test_config_builder_build_sliding_window() {
506        let builder = RateLimitConfigBuilder::new()
507            .algorithm(LimitAlgorithm::SlidingWindow)
508            .capacity(10)
509            .window_secs(60);
510        let limiter = builder.build_sliding_window();
511        assert_eq!(limiter.max_requests(), 10);
512        let r = limiter.acquire("k").unwrap();
513        assert!(r.allowed);
514    }
515
516    #[test]
517    fn test_config_builder_build_fixed_window() {
518        let builder = RateLimitConfigBuilder::new()
519            .algorithm(LimitAlgorithm::FixedWindow)
520            .capacity(10)
521            .window_secs(60);
522        let limiter = builder.build_fixed_window();
523        assert_eq!(limiter.max_requests(), 10);
524        let r = limiter.acquire("k").unwrap();
525        assert!(r.allowed);
526    }
527
528    #[test]
529    fn test_config_builder_build_leaky_bucket() {
530        let builder = RateLimitConfigBuilder::new()
531            .algorithm(LimitAlgorithm::LeakyBucket)
532            .capacity(10)
533            .rate(1.0);
534        let limiter = builder.build_leaky_bucket();
535        assert_eq!(limiter.capacity(), 10);
536        let r = limiter.acquire("k").unwrap();
537        assert!(r.allowed);
538    }
539
540    #[test]
541    fn test_config_builder_build_sliding_window_log() {
542        let builder = RateLimitConfigBuilder::new()
543            .algorithm(LimitAlgorithm::SlidingWindowLog)
544            .capacity(10)
545            .window_secs(60);
546        let limiter = builder.build_sliding_window_log();
547        assert_eq!(limiter.max_requests(), 10);
548        let r = limiter.acquire("k").unwrap();
549        assert!(r.allowed);
550    }
551
552    #[test]
553    fn test_tiered_config_default() {
554        let config = TieredRateLimitConfig::default();
555        assert_eq!(config.ip.capacity, 1000);
556        assert_eq!(config.user.capacity, 100);
557        assert_eq!(config.api.capacity, 500);
558        assert_eq!(config.global.capacity, 10_000);
559    }
560
561    #[test]
562    fn test_tiered_config_validate() {
563        let config = TieredRateLimitConfig::default();
564        assert!(config.validate().is_ok());
565    }
566
567    #[test]
568    fn test_tiered_config_json_roundtrip() {
569        let config = TieredRateLimitConfig::default();
570        let json = config.to_json_string().unwrap();
571        let back = TieredRateLimitConfig::from_json_str(&json).unwrap();
572        assert_eq!(back.ip.capacity, config.ip.capacity);
573    }
574
575    #[test]
576    fn test_tiered_config_builder() {
577        let config = TieredConfigBuilder::new()
578            .ip(RateLimitConfig {
579                capacity: 2000,
580                ..Default::default()
581            })
582            .user(RateLimitConfig {
583                capacity: 200,
584                ..Default::default()
585            })
586            .build();
587        assert_eq!(config.ip.capacity, 2000);
588        assert_eq!(config.user.capacity, 200);
589    }
590
591    #[test]
592    fn test_tiered_config_builder_checked() {
593        let config = TieredConfigBuilder::new().build_checked();
594        assert!(config.is_ok());
595    }
596
597    #[test]
598    fn test_config_builder_window_ms() {
599        let config = RateLimitConfigBuilder::new().window_ms(500).build();
600        assert_eq!(config.window_ms, 500);
601    }
602
603    #[test]
604    fn test_config_builder_rate() {
605        let config = RateLimitConfigBuilder::new().rate(5.0).build();
606        assert_eq!(config.rate, 5.0);
607    }
608}
sz_orm_limit/config_builder.rs

sz_orm_limit/
config_builder.rs
