pagesift_api/
limits.rs

1//! 按 API key 的限流和配额：每分钟请求数（令牌桶）、同时进行的任务数、每天（UTC）的页面数。
2//! 超限时返回 Firecrawl 兼容的 429 错误体。没配置任何限额时什么都不检查。
3//!
4//! 当天已用的页数写进存储（有的话），重启不会把配额清零。
5
6use std::collections::HashMap;
7use std::sync::{Arc, Mutex};
8use std::time::{Instant, SystemTime, UNIX_EPOCH};
9
10use pagesift_store::Store;
11
12use crate::error::{ApiError, ErrorCode};
13
14/// 一个 API key 的限额，各项不设就是不限。
15#[derive(Debug, Clone, Default, PartialEq, Eq)]
16pub struct KeyLimits {
17    /// 日志、监控和配额记录里用的名字；不设时用 key 的指纹，key 本身不会写进日志或磁盘。
18    pub name: Option<String>,
19    pub requests_per_minute: Option<u32>,
20    pub max_concurrent_jobs: Option<usize>,
21    pub pages_per_day: Option<u64>,
22}
23
24/// 发起请求的调用方：鉴权中间件放进请求扩展里。
25#[derive(Debug, Clone)]
26pub(crate) struct Caller {
27    pub account: Arc<str>,
28    pub limits: Arc<KeyLimits>,
29}
30
31impl Caller {
32    pub(crate) fn new(key: Option<&str>, limits: Arc<KeyLimits>) -> Self {
33        let account = match (&limits.name, key) {
34            (Some(n), _) => n.clone(),
35            (None, Some(k)) => fingerprint(k),
36            (None, None) => "anonymous".into(),
37        };
38        Self {
39            account: account.into(),
40            limits,
41        }
42    }
43}
44
45impl Caller {
46    /// 恢复任务时按保存的名字找回调用方。
47    pub(crate) fn named(account: &str, limits: Arc<KeyLimits>) -> Self {
48        Self {
49            account: account.into(),
50            limits,
51        }
52    }
53}
54
55/// key 的指纹（FNV-1a 64 位），稳定、不可逆到能用的程度，用来区分调用方。
56fn fingerprint(key: &str) -> String {
57    let mut h: u64 = 0xcbf2_9ce4_8422_2325;
58    for b in key.bytes() {
59        h ^= u64::from(b);
60        h = h.wrapping_mul(0x0100_0000_01b3);
61    }
62    format!("key-{:08x}", h >> 32)
63}
64
65struct Bucket {
66    tokens: f64,
67    last: Instant,
68}
69
70#[derive(Clone, Copy)]
71struct Usage {
72    day: u64,
73    used: u64,
74}
75
76#[derive(Default)]
77pub(crate) struct Limiter {
78    buckets: Mutex<HashMap<Arc<str>, Bucket>>,
79    usage: Mutex<HashMap<Arc<str>, Usage>>,
80    store: Option<Store>,
81}
82
83fn today() -> u64 {
84    SystemTime::now()
85        .duration_since(UNIX_EPOCH)
86        .map_or(0, |d| d.as_secs() / 86_400)
87}
88
89fn secs_to_midnight() -> u64 {
90    let now = SystemTime::now()
91        .duration_since(UNIX_EPOCH)
92        .map_or(0, |d| d.as_secs());
93    86_400 - now % 86_400
94}
95
96fn usage_key(account: &str) -> String {
97    format!("quota/{account}")
98}
99
100impl Limiter {
101    pub(crate) fn new(store: Option<Store>) -> Self {
102        Self {
103            store,
104            ..Default::default()
105        }
106    }
107
108    /// 每分钟请求数：容量等于每分钟请求数的令牌桶，匀速补充，允许短时突发。
109    pub(crate) fn check_rate(&self, c: &Caller) -> Result<(), ApiError> {
110        let Some(rpm) = c.limits.requests_per_minute else {
111            return Ok(());
112        };
113        let cap = f64::from(rpm.max(1));
114        let per_sec = cap / 60.0;
115        let now = Instant::now();
116        let mut buckets = self.buckets.lock().unwrap_or_else(|e| e.into_inner());
117        let b = buckets.entry(c.account.clone()).or_insert(Bucket {
118            tokens: cap,
119            last: now,
120        });
121        b.tokens = (b.tokens + now.duration_since(b.last).as_secs_f64() * per_sec).min(cap);
122        b.last = now;
123        if b.tokens >= 1.0 {
124            b.tokens -= 1.0;
125            return Ok(());
126        }
127        let wait = ((1.0 - b.tokens) / per_sec).ceil().max(1.0) as u64;
128        Err(ApiError::new(
129            ErrorCode::RateLimitExceeded,
130            format!("rate limit exceeded: {rpm} requests per minute for this API key; retry after {wait}s"),
131        )
132        .with_retry_after(wait))
133    }
134
135    /// 同时进行的任务数是否还能再加一个。
136    pub(crate) fn check_jobs(&self, c: &Caller, active: usize) -> Result<(), ApiError> {
137        match c.limits.max_concurrent_jobs {
138            Some(max) if active >= max => Err(ApiError::new(
139                ErrorCode::ConcurrencyLimitExceeded,
140                format!("concurrency limit exceeded: at most {max} crawl or batch jobs may run at once for this API key"),
141            )
142            .with_retry_after(5)),
143            _ => Ok(()),
144        }
145    }
146
147    fn usage(&self, account: &Arc<str>) -> Usage {
148        let day = today();
149        let mut map = self.usage.lock().unwrap_or_else(|e| e.into_inner());
150        if let Some(u) = map.get(account).filter(|u| u.day == day) {
151            return *u;
152        }
153        // 第一次用到这个调用方（或者跨天了）：从存储里读当天的记录
154        let stored = self
155            .store
156            .as_ref()
157            .and_then(|s| s.tasks().kv_get(usage_key(account)).ok().flatten())
158            .and_then(|v| {
159                let text = String::from_utf8(v).ok()?;
160                let (d, n) = text.split_once(' ')?;
161                Some(Usage {
162                    day: d.parse().ok()?,
163                    used: n.parse().ok()?,
164                })
165            })
166            .filter(|u| u.day == day);
167        let u = stored.unwrap_or(Usage { day, used: 0 });
168        map.insert(account.clone(), u);
169        u
170    }
171
172    /// 当天还剩多少页；没设配额时是 `None`。
173    pub(crate) fn remaining_pages(&self, c: &Caller) -> Option<u64> {
174        let max = c.limits.pages_per_day?;
175        Some(max.saturating_sub(self.usage(&c.account).used))
176    }
177
178    pub(crate) fn check_quota(&self, c: &Caller) -> Result<(), ApiError> {
179        if self.remaining_pages(c) == Some(0) {
180            return Err(quota_error(c));
181        }
182        Ok(())
183    }
184
185    /// 记入当天用掉的页数。只有设了配额的调用方才计数和落盘。
186    pub(crate) async fn consume_pages(&self, c: &Caller, n: u64) {
187        if c.limits.pages_per_day.is_none() || n == 0 {
188            return;
189        }
190        let u = self.usage(&c.account);
191        let u = {
192            let mut map = self.usage.lock().unwrap_or_else(|e| e.into_inner());
193            let e = map.entry(c.account.clone()).or_insert(u);
194            if e.day != u.day {
195                *e = u;
196            }
197            e.used += n;
198            *e
199        };
200        if let Some(store) = &self.store {
201            let key = usage_key(&c.account);
202            let value = format!("{} {}", u.day, u.used);
203            let r = store
204                .blocking(move |s| s.tasks().kv_put(key, value.as_bytes()))
205                .await;
206            if let Err(e) = r {
207                tracing::warn!(error = %e, "saving page quota usage failed");
208            }
209        }
210    }
211}
212
213pub(crate) fn quota_error(c: &Caller) -> ApiError {
214    ApiError::new(
215        ErrorCode::QuotaExceeded,
216        format!(
217            "daily page quota exceeded: {} pages per day for this API key",
218            c.limits.pages_per_day.unwrap_or(0)
219        ),
220    )
221    .with_retry_after(secs_to_midnight())
222}
223
224#[cfg(test)]
225mod tests {
226    use super::*;
227
228    fn caller(limits: KeyLimits) -> Caller {
229        Caller::new(Some("secret-key"), Arc::new(limits))
230    }
231
232    #[test]
233    fn token_bucket_allows_a_burst_then_429s() {
234        let l = Limiter::default();
235        let c = caller(KeyLimits {
236            requests_per_minute: Some(3),
237            ..Default::default()
238        });
239        for _ in 0..3 {
240            l.check_rate(&c).unwrap();
241        }
242        let e = l.check_rate(&c).unwrap_err();
243        assert_eq!(e.code, ErrorCode::RateLimitExceeded);
244        assert!(e.retry_after.unwrap() >= 1);
245        // 别的调用方有自己的桶
246        let other = Caller::new(Some("other"), c.limits.clone());
247        l.check_rate(&other).unwrap();
248        // 不限速的调用方不受影响
249        l.check_rate(&caller(KeyLimits::default())).unwrap();
250    }
251
252    #[test]
253    fn accounts_use_names_or_fingerprints() {
254        let named = Caller::new(
255            Some("k"),
256            Arc::new(KeyLimits {
257                name: Some("team-a".into()),
258                ..Default::default()
259            }),
260        );
261        assert_eq!(&*named.account, "team-a");
262        let anon = caller(KeyLimits::default());
263        assert!(anon.account.starts_with("key-") && !anon.account.contains("secret"));
264        assert_eq!(&*Caller::new(None, Arc::default()).account, "anonymous");
265    }
266
267    #[tokio::test]
268    async fn page_quota_survives_a_restart() {
269        let dir = tempfile::tempdir().unwrap();
270        let c = caller(KeyLimits {
271            pages_per_day: Some(5),
272            ..Default::default()
273        });
274        {
275            let store = Store::open_dir(dir.path()).unwrap();
276            let l = Limiter::new(Some(store));
277            assert_eq!(l.remaining_pages(&c), Some(5));
278            l.consume_pages(&c, 4).await;
279            assert!(l.check_quota(&c).is_ok());
280            l.consume_pages(&c, 1).await;
281            assert_eq!(
282                l.check_quota(&c).unwrap_err().code,
283                ErrorCode::QuotaExceeded
284            );
285        }
286        let l = Limiter::new(Some(Store::open_dir(dir.path()).unwrap()));
287        assert_eq!(l.remaining_pages(&c), Some(0));
288        let jobs = caller(KeyLimits {
289            max_concurrent_jobs: Some(1),
290            ..Default::default()
291        });
292        assert!(l.check_jobs(&jobs, 0).is_ok());
293        assert_eq!(
294            l.check_jobs(&jobs, 1).unwrap_err().code,
295            ErrorCode::ConcurrencyLimitExceeded
296        );
297    }
298}
pagesift_api/limits.rs

pagesift_api/
limits.rs
