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

sz_orm_limit/
config_builder.rs
