pagesift_api/
limits.rs

1//! 按 API key 的限流和配额：每分钟请求数（令牌桶）、同时进行的任务数、每天（UTC）的页面数。
2//! 超限时返回 Firecrawl 兼容的 429 错误体。没配置任何限额时什么都不检查。
3//!
4//! 当天已用的页数写进存储（有的话），重启不会把配额清零。
5//!
6//! 多服务器网关上（[`Limiter::with_ledger`]）每分钟请求数和每日页数改记在 `pagesift-cluster` 的账本里：
7//! 用量是追加的冻结 / 结算记录，崩溃重放、温备接管后一致；令牌桶在账本里，接管时由最近一分钟的
8//! 请求记录重建。单主机不经过账本，行为不变。
9
10use std::collections::HashMap;
11use std::sync::{Arc, Mutex};
12use std::time::{Instant, SystemTime, UNIX_EPOCH};
13
14use pagesift_store::Store;
15
16use crate::error::{ApiError, ErrorCode};
17
18/// 一个 API key 的限额，各项不设就是不限。
19#[derive(Debug, Clone, Default, PartialEq, Eq)]
20pub struct KeyLimits {
21    /// 日志、监控和配额记录里用的名字；不设时用 key 的指纹，key 本身不会写进日志或磁盘。
22    pub name: Option<String>,
23    pub requests_per_minute: Option<u32>,
24    pub max_concurrent_jobs: Option<usize>,
25    pub pages_per_day: Option<u64>,
26}
27
28/// 发起请求的调用方：鉴权中间件放进请求扩展里。
29#[derive(Debug, Clone)]
30pub(crate) struct Caller {
31    pub account: Arc<str>,
32    pub limits: Arc<KeyLimits>,
33}
34
35impl Caller {
36    pub(crate) fn new(key: Option<&str>, limits: Arc<KeyLimits>) -> Self {
37        let account = match (&limits.name, key) {
38            (Some(n), _) => n.clone(),
39            (None, Some(k)) => fingerprint(k),
40            (None, None) => "anonymous".into(),
41        };
42        Self {
43            account: account.into(),
44            limits,
45        }
46    }
47}
48
49impl Caller {
50    /// 恢复任务时按保存的名字找回调用方。
51    pub(crate) fn named(account: &str, limits: Arc<KeyLimits>) -> Self {
52        Self {
53            account: account.into(),
54            limits,
55        }
56    }
57}
58
59/// key 的指纹（FNV-1a 64 位），稳定、不可逆到能用的程度，用来区分调用方。
60fn fingerprint(key: &str) -> String {
61    let mut h: u64 = 0xcbf2_9ce4_8422_2325;
62    for b in key.bytes() {
63        h ^= u64::from(b);
64        h = h.wrapping_mul(0x0100_0000_01b3);
65    }
66    format!("key-{h:016x}")
67}
68
69struct Bucket {
70    tokens: f64,
71    last: Instant,
72}
73
74#[derive(Clone, Copy)]
75struct Usage {
76    day: u64,
77    used: u64,
78}
79
80#[derive(Default)]
81pub(crate) struct Limiter {
82    buckets: Mutex<HashMap<Arc<str>, Bucket>>,
83    usage: Mutex<HashMap<Arc<str>, Usage>>,
84    store: Option<Store>,
85    ledger: Option<Arc<pagesift_cluster::Ledger>>,
86    page_reservation: tokio::sync::Mutex<()>,
87}
88
89fn rate_error(rpm: u32, wait: u64) -> ApiError {
90    ApiError::new(
91        ErrorCode::RateLimitExceeded,
92        format!(
93            "rate limit exceeded: {rpm} requests per minute for this API key; retry after {wait}s"
94        ),
95    )
96    .with_retry_after(wait)
97}
98
99fn ledger_error(e: impl std::fmt::Display) -> ApiError {
100    ApiError::new(ErrorCode::Internal, format!("quota ledger: {e}"))
101}
102
103fn today() -> u64 {
104    SystemTime::now()
105        .duration_since(UNIX_EPOCH)
106        .map_or(0, |d| d.as_secs() / 86_400)
107}
108
109fn secs_to_midnight() -> u64 {
110    let now = SystemTime::now()
111        .duration_since(UNIX_EPOCH)
112        .map_or(0, |d| d.as_secs());
113    86_400 - now % 86_400
114}
115
116fn usage_key(account: &str) -> String {
117    format!("quota/{account}")
118}
119
120impl Limiter {
121    pub(crate) fn new(store: Option<Store>) -> Self {
122        Self {
123            store,
124            ..Default::default()
125        }
126    }
127
128    /// 网关：限额记在账本里。`accounts` 是全部账户（名字 → 限额），事先登记。
129    pub(crate) fn with_ledger(
130        ledger: Arc<pagesift_cluster::Ledger>,
131        accounts: &HashMap<String, Arc<KeyLimits>>,
132    ) -> Self {
133        for (name, l) in accounts {
134            ledger.set_limits(
135                name,
136                pagesift_cluster::AccountLimits {
137                    pages_per_day: l.pages_per_day,
138                    requests_per_minute: l.requests_per_minute,
139                },
140            );
141        }
142        Self {
143            ledger: Some(ledger),
144            ..Default::default()
145        }
146    }
147
148    /// 每分钟请求数：容量等于每分钟请求数的令牌桶，匀速补充，允许短时突发。
149    pub(crate) fn check_rate(&self, c: &Caller) -> Result<(), ApiError> {
150        let Some(rpm) = c.limits.requests_per_minute else {
151            return Ok(());
152        };
153        if let Some(l) = &self.ledger {
154            use pagesift_cluster::Refusal;
155            return match l.check_rate(&c.account) {
156                Ok(()) => Ok(()),
157                Err(Refusal::Rate { retry_after, .. }) => Err(rate_error(
158                    rpm,
159                    retry_after.as_secs_f64().ceil().max(1.0) as u64,
160                )),
161                Err(Refusal::Quota { .. }) => Err(quota_error(c)),
162                Err(Refusal::Error(e)) => Err(ledger_error(e)),
163            };
164        }
165        let cap = f64::from(rpm.max(1));
166        let per_sec = cap / 60.0;
167        let now = Instant::now();
168        let mut buckets = self.buckets.lock().unwrap_or_else(|e| e.into_inner());
169        let b = buckets.entry(c.account.clone()).or_insert(Bucket {
170            tokens: cap,
171            last: now,
172        });
173        b.tokens = (b.tokens + now.duration_since(b.last).as_secs_f64() * per_sec).min(cap);
174        b.last = now;
175        if b.tokens >= 1.0 {
176            b.tokens -= 1.0;
177            return Ok(());
178        }
179        let wait = ((1.0 - b.tokens) / per_sec).ceil().max(1.0) as u64;
180        Err(rate_error(rpm, wait))
181    }
182
183    /// 同时进行的任务数是否还能再加一个。
184    pub(crate) fn check_jobs(&self, c: &Caller, active: usize) -> Result<(), ApiError> {
185        match c.limits.max_concurrent_jobs {
186            Some(max) if active >= max => Err(ApiError::new(
187                ErrorCode::ConcurrencyLimitExceeded,
188                format!("concurrency limit exceeded: at most {max} crawl or batch jobs may run at once for this API key"),
189            )
190            .with_retry_after(5)),
191            _ => Ok(()),
192        }
193    }
194
195    fn usage(&self, account: &Arc<str>) -> Usage {
196        let day = today();
197        let mut map = self.usage.lock().unwrap_or_else(|e| e.into_inner());
198        if let Some(u) = map.get(account).filter(|u| u.day == day) {
199            return *u;
200        }
201        // 第一次用到这个调用方（或者跨天了）：从存储里读当天的记录
202        let stored = self
203            .store
204            .as_ref()
205            .and_then(|s| s.tasks().kv_get(usage_key(account)).ok().flatten())
206            .and_then(|v| {
207                let text = String::from_utf8(v).ok()?;
208                let (d, n) = text.split_once(' ')?;
209                Some(Usage {
210                    day: d.parse().ok()?,
211                    used: n.parse().ok()?,
212                })
213            })
214            .filter(|u| u.day == day);
215        let u = stored.unwrap_or(Usage { day, used: 0 });
216        map.insert(account.clone(), u);
217        u
218    }
219
220    /// 当天还剩多少页；没设配额时是 `None`。
221    pub(crate) fn remaining_pages(&self, c: &Caller) -> Option<u64> {
222        let max = c.limits.pages_per_day?;
223        if let Some(l) = &self.ledger {
224            return l.balance(&c.account);
225        }
226        Some(max.saturating_sub(self.usage(&c.account).used))
227    }
228
229    pub(crate) fn check_quota(&self, c: &Caller) -> Result<(), ApiError> {
230        if self.remaining_pages(c) == Some(0) {
231            return Err(quota_error(c));
232        }
233        Ok(())
234    }
235
236    pub(crate) async fn reserve_page(&self, c: &Caller) -> Result<(), ApiError> {
237        let _guard = self.page_reservation.lock().await;
238        if c.limits.pages_per_day.is_some()
239            && let Some(ledger) = &self.ledger
240        {
241            let hold = ledger.reserve(&c.account, 1).map_err(|err| match err {
242                pagesift_cluster::Refusal::Error(err) => ledger_error(err),
243                _ => quota_error(c),
244            })?;
245            return ledger.settle(&hold, 1).map_err(ledger_error);
246        }
247        self.check_quota(c)?;
248        self.consume_pages(c, 1).await;
249        Ok(())
250    }
251
252    /// 记入当天用掉的页数。只有设了配额的调用方才计数和落盘。
253    pub(crate) async fn consume_pages(&self, c: &Caller, n: u64) {
254        if c.limits.pages_per_day.is_none() || n == 0 {
255            return;
256        }
257        if let Some(l) = &self.ledger {
258            // 冻结后立即按实际用量结算：页面已经抓完。余额不够冻结时（并发请求都通过了检查）
259            // 当天已经用满，不再记账，余额照样是 0
260            match l.reserve(&c.account, n) {
261                Ok(hold) => {
262                    if let Err(e) = l.settle(&hold, n) {
263                        tracing::warn!(error = %e, "settling page quota failed");
264                    }
265                }
266                Err(pagesift_cluster::Refusal::Error(e)) => {
267                    tracing::warn!(error = %e, "reserving page quota failed");
268                }
269                Err(_) => {}
270            }
271            return;
272        }
273        let u = self.usage(&c.account);
274        let u = {
275            let mut map = self.usage.lock().unwrap_or_else(|e| e.into_inner());
276            let e = map.entry(c.account.clone()).or_insert(u);
277            if e.day != u.day {
278                *e = u;
279            }
280            e.used += n;
281            *e
282        };
283        if let Some(store) = &self.store {
284            let key = usage_key(&c.account);
285            let value = format!("{} {}", u.day, u.used);
286            let r = store
287                .blocking(move |s| s.tasks().kv_put(key, value.as_bytes()))
288                .await;
289            if let Err(e) = r {
290                tracing::warn!(error = %e, "saving page quota usage failed");
291            }
292        }
293    }
294}
295
296pub(crate) fn quota_error(c: &Caller) -> ApiError {
297    ApiError::new(
298        ErrorCode::QuotaExceeded,
299        format!(
300            "daily page quota exceeded: {} pages per day for this API key",
301            c.limits.pages_per_day.unwrap_or(0)
302        ),
303    )
304    .with_retry_after(secs_to_midnight())
305}
306
307#[cfg(test)]
308mod tests {
309    use super::*;
310
311    fn caller(limits: KeyLimits) -> Caller {
312        Caller::new(Some("secret-key"), Arc::new(limits))
313    }
314
315    #[test]
316    fn token_bucket_allows_a_burst_then_429s() {
317        let l = Limiter::default();
318        let c = caller(KeyLimits {
319            requests_per_minute: Some(3),
320            ..Default::default()
321        });
322        for _ in 0..3 {
323            l.check_rate(&c).unwrap();
324        }
325        let e = l.check_rate(&c).unwrap_err();
326        assert_eq!(e.code, ErrorCode::RateLimitExceeded);
327        assert!(e.retry_after.unwrap() >= 1);
328        // 别的调用方有自己的桶
329        let other = Caller::new(Some("other"), c.limits.clone());
330        l.check_rate(&other).unwrap();
331        // 不限速的调用方不受影响
332        l.check_rate(&caller(KeyLimits::default())).unwrap();
333    }
334
335    #[test]
336    fn accounts_use_names_or_fingerprints() {
337        let named = Caller::new(
338            Some("k"),
339            Arc::new(KeyLimits {
340                name: Some("team-a".into()),
341                ..Default::default()
342            }),
343        );
344        assert_eq!(&*named.account, "team-a");
345        let anon = caller(KeyLimits::default());
346        assert!(anon.account.starts_with("key-") && !anon.account.contains("secret"));
347        assert_eq!(&*Caller::new(None, Arc::default()).account, "anonymous");
348    }
349
350    /// 网关：限额记在账本里，行为与单主机一致（桶先突发后 429，配额用完 QUOTA_EXCEEDED）。
351    #[tokio::test]
352    async fn ledger_backed_limits_match_the_local_ones() {
353        let ledger = Arc::new(
354            pagesift_cluster::Ledger::open(pagesift_cluster::LedgerConfig::new(
355                pagesift_cluster::MonoClock::new(),
356                false,
357            ))
358            .unwrap(),
359        );
360        let limits = Arc::new(KeyLimits {
361            requests_per_minute: Some(2),
362            pages_per_day: Some(2),
363            ..Default::default()
364        });
365        let c = Caller::new(Some("k"), limits.clone());
366        let accounts = HashMap::from([(c.account.to_string(), limits)]);
367        let l = Limiter::with_ledger(ledger, &accounts);
368        l.check_rate(&c).unwrap();
369        l.check_rate(&c).unwrap();
370        let e = l.check_rate(&c).unwrap_err();
371        assert_eq!(e.code, ErrorCode::RateLimitExceeded);
372        assert!(e.retry_after.unwrap() >= 1);
373        assert_eq!(l.remaining_pages(&c), Some(2));
374        l.consume_pages(&c, 1).await;
375        assert!(l.check_quota(&c).is_ok());
376        l.consume_pages(&c, 1).await;
377        assert_eq!(l.remaining_pages(&c), Some(0));
378        assert_eq!(
379            l.check_quota(&c).unwrap_err().code,
380            ErrorCode::QuotaExceeded
381        );
382        // 不限额的调用方不受影响
383        l.check_rate(&caller(KeyLimits::default())).unwrap();
384    }
385
386    #[tokio::test]
387    async fn concurrent_page_reservations_stop_at_quota() {
388        let c = caller(KeyLimits {
389            pages_per_day: Some(1),
390            ..Default::default()
391        });
392        let limiter = Limiter::new(None);
393        let (a, b) = tokio::join!(limiter.reserve_page(&c), limiter.reserve_page(&c));
394        assert_ne!(a.is_ok(), b.is_ok());
395        assert_eq!(limiter.remaining_pages(&c), Some(0));
396    }
397
398    #[tokio::test]
399    async fn page_quota_survives_a_restart() {
400        let dir = tempfile::tempdir().unwrap();
401        let c = caller(KeyLimits {
402            pages_per_day: Some(5),
403            ..Default::default()
404        });
405        {
406            let store = Store::open_dir(dir.path()).unwrap();
407            let l = Limiter::new(Some(store));
408            assert_eq!(l.remaining_pages(&c), Some(5));
409            l.consume_pages(&c, 4).await;
410            assert!(l.check_quota(&c).is_ok());
411            l.consume_pages(&c, 1).await;
412            assert_eq!(
413                l.check_quota(&c).unwrap_err().code,
414                ErrorCode::QuotaExceeded
415            );
416        }
417        let l = Limiter::new(Some(Store::open_dir(dir.path()).unwrap()));
418        assert_eq!(l.remaining_pages(&c), Some(0));
419        let jobs = caller(KeyLimits {
420            max_concurrent_jobs: Some(1),
421            ..Default::default()
422        });
423        assert!(l.check_jobs(&jobs, 0).is_ok());
424        assert_eq!(
425            l.check_jobs(&jobs, 1).unwrap_err().code,
426            ErrorCode::ConcurrencyLimitExceeded
427        );
428    }
429}
pagesift_api/limits.rs

pagesift_api/
limits.rs
