nagisa_core/
ratelimit.rs

1//! 内置的、按需启用的限流器。一个令牌桶 `Middleware`：当某 peer——或全局桶——
2//! 超过配置速率时返回 `Flow::Stop`（静默丢弃）。纯内存，无外部依赖。
3//!
4//! 经 `App::layer(RateLimit::per_peer(max, per))` 启用。
5use crate::ctx::Ctx;
6use crate::middleware::{Flow, Middleware, Next};
7use async_trait::async_trait;
8use nagisa_types::id::Peer;
9use std::collections::HashMap;
10use std::sync::{Arc, Mutex};
11use std::time::{Duration, Instant};
12
13/// `RateLimit` 以什么为桶的键。
14#[derive(Clone, Copy, Debug, PartialEq, Eq)]
15pub enum RateLimitScope {
16    /// 每个可寻址 peer（群/好友/临时）一个桶。无 peer 事件放行。
17    PerPeer,
18    /// 所有事件共享一个桶。
19    Global,
20}
21
22/// 可当 `Middleware` 用的令牌桶限流器。
23#[derive(Clone)]
24pub struct RateLimit {
25    scope: RateLimitScope,
26    /// 最大令牌数（突发容量）。
27    max: f64,
28    /// 补满 `max` 个令牌的窗口。
29    per: Duration,
30    /// 按 peer 的桶（PerPeer）或单个以 `None` 为键的桶（Global）。
31    buckets: Arc<Mutex<HashMap<Option<Peer>, Bucket>>>,
32}
33
34#[derive(Clone, Copy)]
35struct Bucket {
36    tokens: f64,
37    last: Instant,
38}
39
40impl RateLimit {
41    /// 每个 peer 一个桶：每 `per` 窗口、每 peer 最多 `max` 个事件。
42    pub fn per_peer(max: u32, per: Duration) -> Self {
43        RateLimit {
44            scope: RateLimitScope::PerPeer,
45            max: max as f64,
46            per,
47            buckets: Arc::new(Mutex::new(HashMap::new())),
48        }
49    }
50    /// 单个共享桶：每 `per` 窗口全局最多 `max` 个事件。
51    pub fn global(max: u32, per: Duration) -> Self {
52        RateLimit { scope: RateLimitScope::Global, max: max as f64, per, buckets: Arc::new(Mutex::new(HashMap::new())) }
53    }
54
55    /// 为 `peer` 消耗一个令牌。放行返回 `true`，被限流返回 `false`。
56    /// `PerPeer` 下无 peer 的事件恒放行（没有可作键的东西）。
57    pub fn check(&self, peer: Option<Peer>) -> bool {
58        let key = match self.scope {
59            RateLimitScope::PerPeer => match peer {
60                Some(p) => Some(p),
61                None => return true, // 无可寻址 peer → 不限流
62            },
63            RateLimitScope::Global => None,
64        };
65        let refill_per_sec = self.max / self.per.as_secs_f64();
66        let now = Instant::now();
67        let mut map = self.buckets.lock().expect("ratelimit buckets poisoned");
68        let bucket = map.entry(key).or_insert(Bucket { tokens: self.max, last: now });
69        // 按流逝时间补充令牌，封顶到 `max`。
70        let elapsed = now.saturating_duration_since(bucket.last).as_secs_f64();
71        bucket.tokens = (bucket.tokens + elapsed * refill_per_sec).min(self.max);
72        bucket.last = now;
73        if bucket.tokens >= 1.0 {
74            bucket.tokens -= 1.0;
75            true
76        } else {
77            false
78        }
79    }
80}
81
82#[async_trait]
83impl Middleware for RateLimit {
84    async fn handle(&self, ctx: Arc<Ctx>, next: Next<'_>) -> Flow {
85        let peer = ctx.event().peer();
86        if self.check(peer) {
87            next.run(ctx).await
88        } else {
89            tracing::debug!(?peer, "rate limited; dropping event");
90            Flow::Stop
91        }
92    }
93}
nagisa_core/ratelimit.rs

nagisa_core/
ratelimit.rs
