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}
87
88fn rate_error(rpm: u32, wait: u64) -> ApiError {
89    ApiError::new(
90        ErrorCode::RateLimitExceeded,
91        format!(
92            "rate limit exceeded: {rpm} requests per minute for this API key; retry after {wait}s"
93        ),
94    )
95    .with_retry_after(wait)
96}
97
98fn ledger_error(e: impl std::fmt::Display) -> ApiError {
99    ApiError::new(ErrorCode::Internal, format!("quota ledger: {e}"))
100}
101
102fn today() -> u64 {
103    SystemTime::now()
104        .duration_since(UNIX_EPOCH)
105        .map_or(0, |d| d.as_secs() / 86_400)
106}
107
108fn secs_to_midnight() -> u64 {
109    let now = SystemTime::now()
110        .duration_since(UNIX_EPOCH)
111        .map_or(0, |d| d.as_secs());
112    86_400 - now % 86_400
113}
114
115fn usage_key(account: &str) -> String {
116    format!("quota/{account}")
117}
118
119impl Limiter {
120    pub(crate) fn new(store: Option<Store>) -> Self {
121        Self {
122            store,
123            ..Default::default()
124        }
125    }
126
127    /// 网关：限额记在账本里。`accounts` 是全部账户（名字 → 限额），事先登记。
128    pub(crate) fn with_ledger(
129        ledger: Arc<pagesift_cluster::Ledger>,
130        accounts: &HashMap<String, Arc<KeyLimits>>,
131    ) -> Self {
132        for (name, l) in accounts {
133            ledger.set_limits(
134                name,
135                pagesift_cluster::AccountLimits {
136                    pages_per_day: l.pages_per_day,
137                    requests_per_minute: l.requests_per_minute,
138                },
139            );
140        }
141        Self {
142            ledger: Some(ledger),
143            ..Default::default()
144        }
145    }
146
147    /// 每分钟请求数：容量等于每分钟请求数的令牌桶，匀速补充，允许短时突发。
148    pub(crate) fn check_rate(&self, c: &Caller) -> Result<(), ApiError> {
149        let Some(rpm) = c.limits.requests_per_minute else {
150            return Ok(());
151        };
152        if let Some(l) = &self.ledger {
153            use pagesift_cluster::Refusal;
154            return match l.check_rate(&c.account) {
155                Ok(()) => Ok(()),
156                Err(Refusal::Rate { retry_after, .. }) => Err(rate_error(
157                    rpm,
158                    retry_after.as_secs_f64().ceil().max(1.0) as u64,
159                )),
160                Err(Refusal::Quota { .. }) => Err(quota_error(c)),
161                Err(Refusal::Error(e)) => Err(ledger_error(e)),
162            };
163        }
164        let cap = f64::from(rpm.max(1));
165        let per_sec = cap / 60.0;
166        let now = Instant::now();
167        let mut buckets = self.buckets.lock().unwrap_or_else(|e| e.into_inner());
168        let b = buckets.entry(c.account.clone()).or_insert(Bucket {
169            tokens: cap,
170            last: now,
171        });
172        b.tokens = (b.tokens + now.duration_since(b.last).as_secs_f64() * per_sec).min(cap);
173        b.last = now;
174        if b.tokens >= 1.0 {
175            b.tokens -= 1.0;
176            return Ok(());
177        }
178        let wait = ((1.0 - b.tokens) / per_sec).ceil().max(1.0) as u64;
179        Err(rate_error(rpm, wait))
180    }
181
182    /// 同时进行的任务数是否还能再加一个。
183    pub(crate) fn check_jobs(&self, c: &Caller, active: usize) -> Result<(), ApiError> {
184        match c.limits.max_concurrent_jobs {
185            Some(max) if active >= max => Err(ApiError::new(
186                ErrorCode::ConcurrencyLimitExceeded,
187                format!("concurrency limit exceeded: at most {max} crawl or batch jobs may run at once for this API key"),
188            )
189            .with_retry_after(5)),
190            _ => Ok(()),
191        }
192    }
193
194    fn load_usage(&self, account: &Arc<str>, day: u64) -> Usage {
195        // 第一次用到这个调用方（或者跨天了）：从存储里读当天的记录
196        let stored = self
197            .store
198            .as_ref()
199            .and_then(|s| s.tasks().kv_get(usage_key(account)).ok().flatten())
200            .and_then(|v| {
201                let text = String::from_utf8(v).ok()?;
202                let (d, n) = text.split_once(' ')?;
203                Some(Usage {
204                    day: d.parse().ok()?,
205                    used: n.parse().ok()?,
206                })
207            })
208            .filter(|u| u.day == day);
209        stored.unwrap_or(Usage { day, used: 0 })
210    }
211
212    fn usage_entry<'a>(
213        &self,
214        map: &'a mut HashMap<Arc<str>, Usage>,
215        account: &Arc<str>,
216        day: u64,
217    ) -> &'a mut Usage {
218        if map.get(account).is_none_or(|u| u.day != day) {
219            let u = self.load_usage(account, day);
220            map.insert(account.clone(), u);
221        }
222        map.get_mut(account).expect("inserted above")
223    }
224
225    fn usage(&self, account: &Arc<str>) -> Usage {
226        let day = today();
227        let mut map = self.usage.lock().unwrap_or_else(|e| e.into_inner());
228        *self.usage_entry(&mut map, account, day)
229    }
230
231    /// 当天还剩多少页；没设配额时是 `None`。已预留但还没结算的页也算在用量里。
232    pub(crate) fn remaining_pages(&self, c: &Caller) -> Option<u64> {
233        let max = c.limits.pages_per_day?;
234        if let Some(l) = &self.ledger {
235            return l.balance(&c.account);
236        }
237        Some(max.saturating_sub(self.usage(&c.account).used))
238    }
239
240    pub(crate) fn check_quota(&self, c: &Caller) -> Result<(), ApiError> {
241        if self.remaining_pages(c) == Some(0) {
242            return Err(quota_error(c));
243        }
244        Ok(())
245    }
246
247    /// 可计费操作发出之前原子地预留一页。返回的 [`PageHold`] 在成功时 `commit`，
248    /// 失败/取消时直接丢弃即自动退还（[`Drop`] 里完成），不会出现 check-then-act 的并发窗口。
249    ///
250    /// 账本模式下，进程在预留尚未结算时崩溃，冻结会保留到当天结束（UTC）：最多是崩溃瞬间
251    /// 在途的页数，不会双重扣费；单机模式的内存预留则随进程消失。
252    pub(crate) async fn reserve_page(self: &Arc<Self>, c: &Caller) -> Result<PageHold, ApiError> {
253        let day = today();
254        let amount = 1;
255        let Some(_) = c.limits.pages_per_day else {
256            return Ok(PageHold {
257                limiter: self.clone(),
258                account: c.account.clone(),
259                day,
260                amount,
261                ledger: None,
262                tracked: false,
263                settled: false,
264            });
265        };
266        if let Some(ledger) = &self.ledger {
267            let hold = ledger
268                .reserve(&c.account, amount)
269                .map_err(|err| match err {
270                    pagesift_cluster::Refusal::Error(err) => ledger_error(err),
271                    _ => quota_error(c),
272                })?;
273            return Ok(PageHold {
274                limiter: self.clone(),
275                account: c.account.clone(),
276                day,
277                amount,
278                ledger: Some(hold),
279                tracked: true,
280                settled: false,
281            });
282        }
283        let max = c.limits.pages_per_day.unwrap_or(0);
284        {
285            let mut map = self.usage.lock().unwrap_or_else(|e| e.into_inner());
286            let u = self.usage_entry(&mut map, &c.account, day);
287            if u.used.saturating_add(amount) > max {
288                return Err(quota_error(c));
289            }
290            u.used += amount;
291        }
292        Ok(PageHold {
293            limiter: self.clone(),
294            account: c.account.clone(),
295            day,
296            amount,
297            ledger: None,
298            tracked: true,
299            settled: false,
300        })
301    }
302
303    /// 退还预留（没有结算就丢弃）。只退当天的记录，跨天时旧记录已经作废。
304    fn refund(&self, account: &Arc<str>, day: u64, amount: u64) {
305        let mut map = self.usage.lock().unwrap_or_else(|e| e.into_inner());
306        if let Some(u) = map.get_mut(account)
307            && u.day == day
308        {
309            u.used = u.used.saturating_sub(amount);
310        }
311    }
312
313    async fn persist_usage(&self, account: &Arc<str>) {
314        let Some(store) = &self.store else {
315            return;
316        };
317        let u = {
318            let map = self.usage.lock().unwrap_or_else(|e| e.into_inner());
319            match map.get(account) {
320                Some(u) => *u,
321                None => return,
322            }
323        };
324        let key = usage_key(account);
325        let value = format!("{} {}", u.day, u.used);
326        let r = store
327            .blocking(move |s| s.tasks().kv_put(key, value.as_bytes()))
328            .await;
329        if let Err(e) = r {
330            tracing::warn!(error = %e, "saving page quota usage failed");
331        }
332    }
333}
334
335/// 一页的预留。`commit` 结算；未结算就丢弃时自动退还（失败、取消、panic 都走同一条路）。
336pub(crate) struct PageHold {
337    limiter: Arc<Limiter>,
338    account: Arc<str>,
339    day: u64,
340    amount: u64,
341    ledger: Option<pagesift_cluster::Hold>,
342    tracked: bool,
343    settled: bool,
344}
345
346impl PageHold {
347    pub(crate) async fn commit(mut self) {
348        self.settled = true;
349        if !self.tracked {
350            return;
351        }
352        if let Some(ledger) = &self.limiter.ledger {
353            if let Some(hold) = self.ledger.take()
354                && let Err(e) = ledger.settle(&hold, self.amount)
355            {
356                tracing::warn!(error = %e, "settling page quota failed");
357            }
358            return;
359        }
360        self.limiter.persist_usage(&self.account).await;
361    }
362}
363
364impl Drop for PageHold {
365    fn drop(&mut self) {
366        if self.settled || !self.tracked {
367            return;
368        }
369        if let Some(ledger) = &self.limiter.ledger {
370            if let Some(hold) = self.ledger.take()
371                && let Err(e) = ledger.release(&hold)
372            {
373                tracing::warn!(error = %e, "releasing page quota failed");
374            }
375            return;
376        }
377        self.limiter.refund(&self.account, self.day, self.amount);
378    }
379}
380
381pub(crate) fn quota_error(c: &Caller) -> ApiError {
382    ApiError::new(
383        ErrorCode::QuotaExceeded,
384        format!(
385            "daily page quota exceeded: {} pages per day for this API key",
386            c.limits.pages_per_day.unwrap_or(0)
387        ),
388    )
389    .with_retry_after(secs_to_midnight())
390}
391
392#[cfg(test)]
393mod tests {
394    use super::*;
395
396    fn caller(limits: KeyLimits) -> Caller {
397        Caller::new(Some("secret-key"), Arc::new(limits))
398    }
399
400    #[test]
401    fn token_bucket_allows_a_burst_then_429s() {
402        let l = Limiter::default();
403        let c = caller(KeyLimits {
404            requests_per_minute: Some(3),
405            ..Default::default()
406        });
407        for _ in 0..3 {
408            l.check_rate(&c).unwrap();
409        }
410        let e = l.check_rate(&c).unwrap_err();
411        assert_eq!(e.code, ErrorCode::RateLimitExceeded);
412        assert!(e.retry_after.unwrap() >= 1);
413        // 别的调用方有自己的桶
414        let other = Caller::new(Some("other"), c.limits.clone());
415        l.check_rate(&other).unwrap();
416        // 不限速的调用方不受影响
417        l.check_rate(&caller(KeyLimits::default())).unwrap();
418    }
419
420    #[test]
421    fn accounts_use_names_or_fingerprints() {
422        let named = Caller::new(
423            Some("k"),
424            Arc::new(KeyLimits {
425                name: Some("team-a".into()),
426                ..Default::default()
427            }),
428        );
429        assert_eq!(&*named.account, "team-a");
430        let anon = caller(KeyLimits::default());
431        assert!(anon.account.starts_with("key-") && !anon.account.contains("secret"));
432        assert_eq!(&*Caller::new(None, Arc::default()).account, "anonymous");
433    }
434
435    /// 网关：限额记在账本里，行为与单主机一致（桶先突发后 429，配额用完 QUOTA_EXCEEDED）。
436    #[tokio::test]
437    async fn ledger_backed_limits_match_the_local_ones() {
438        let ledger = Arc::new(
439            pagesift_cluster::Ledger::open(pagesift_cluster::LedgerConfig::new(
440                pagesift_cluster::MonoClock::new(),
441                false,
442            ))
443            .unwrap(),
444        );
445        let limits = Arc::new(KeyLimits {
446            requests_per_minute: Some(2),
447            pages_per_day: Some(2),
448            ..Default::default()
449        });
450        let c = Caller::new(Some("k"), limits.clone());
451        let accounts = HashMap::from([(c.account.to_string(), limits)]);
452        let l = Arc::new(Limiter::with_ledger(ledger, &accounts));
453        l.check_rate(&c).unwrap();
454        l.check_rate(&c).unwrap();
455        let e = l.check_rate(&c).unwrap_err();
456        assert_eq!(e.code, ErrorCode::RateLimitExceeded);
457        assert!(e.retry_after.unwrap() >= 1);
458        assert_eq!(l.remaining_pages(&c), Some(2));
459        l.reserve_page(&c).await.unwrap().commit().await;
460        assert!(l.check_quota(&c).is_ok());
461        l.reserve_page(&c).await.unwrap().commit().await;
462        assert_eq!(l.remaining_pages(&c), Some(0));
463        assert_eq!(
464            l.check_quota(&c).unwrap_err().code,
465            ErrorCode::QuotaExceeded
466        );
467        // 不限额的调用方不受影响
468        l.check_rate(&caller(KeyLimits::default())).unwrap();
469    }
470
471    #[tokio::test]
472    async fn concurrent_page_reservations_stop_at_quota() {
473        let c = caller(KeyLimits {
474            pages_per_day: Some(1),
475            ..Default::default()
476        });
477        let limiter = Arc::new(Limiter::new(None));
478        let (a, b) = tokio::join!(limiter.reserve_page(&c), limiter.reserve_page(&c));
479        assert_ne!(a.is_ok(), b.is_ok());
480        assert_eq!(limiter.remaining_pages(&c), Some(0));
481        // 丢掉未结算的预留：额度回来
482        drop(a);
483        drop(b);
484        assert_eq!(limiter.remaining_pages(&c), Some(1));
485    }
486
487    #[tokio::test]
488    async fn twenty_concurrent_reservations_allow_exactly_one() {
489        let c = Arc::new(caller(KeyLimits {
490            pages_per_day: Some(1),
491            ..Default::default()
492        }));
493        let limiter = Arc::new(Limiter::new(None));
494        let mut tasks = Vec::new();
495        for _ in 0..20 {
496            let (limiter, c) = (limiter.clone(), c.clone());
497            tasks.push(tokio::spawn(async move { limiter.reserve_page(&c).await }));
498        }
499        let mut granted = 0;
500        let mut holds = Vec::new();
501        for task in tasks {
502            match task.await.unwrap() {
503                Ok(hold) => {
504                    granted += 1;
505                    holds.push(hold);
506                }
507                Err(e) => assert_eq!(e.code, ErrorCode::QuotaExceeded),
508            }
509        }
510        assert_eq!(granted, 1);
511        assert_eq!(limiter.remaining_pages(&c), Some(0));
512        drop(holds);
513        assert_eq!(limiter.remaining_pages(&c), Some(1));
514    }
515
516    #[tokio::test]
517    async fn committed_pages_survive_restart_and_holds_refund() {
518        let dir = tempfile::tempdir().unwrap();
519        let c = caller(KeyLimits {
520            pages_per_day: Some(2),
521            ..Default::default()
522        });
523        let limiter = Arc::new(Limiter::new(Some(Store::open_dir(dir.path()).unwrap())));
524        let hold = limiter.reserve_page(&c).await.unwrap();
525        assert_eq!(limiter.remaining_pages(&c), Some(1));
526        drop(hold);
527        assert_eq!(limiter.remaining_pages(&c), Some(2));
528        let hold = limiter.reserve_page(&c).await.unwrap();
529        hold.commit().await;
530        assert_eq!(limiter.remaining_pages(&c), Some(1));
531        drop(limiter);
532        // 重启：已结算的 1 页还在，未结算的预留不落盘
533        let limiter = Limiter::new(Some(Store::open_dir(dir.path()).unwrap()));
534        assert_eq!(limiter.remaining_pages(&c), Some(1));
535    }
536
537    #[tokio::test]
538    async fn ledger_reservations_are_atomic_across_limiter_instances() {
539        let ledger = Arc::new(
540            pagesift_cluster::Ledger::open(pagesift_cluster::LedgerConfig::new(
541                pagesift_cluster::MonoClock::new(),
542                false,
543            ))
544            .unwrap(),
545        );
546        let limits = Arc::new(KeyLimits {
547            pages_per_day: Some(1),
548            ..Default::default()
549        });
550        let c = Arc::new(Caller::new(Some("k"), limits.clone()));
551        let accounts = HashMap::from([(c.account.to_string(), limits)]);
552        let a = Arc::new(Limiter::with_ledger(ledger.clone(), &accounts));
553        let b = Arc::new(Limiter::with_ledger(ledger.clone(), &accounts));
554        let mut tasks = Vec::new();
555        for i in 0..20 {
556            let (l, c) = (if i % 2 == 0 { a.clone() } else { b.clone() }, c.clone());
557            tasks.push(tokio::spawn(async move { l.reserve_page(&c).await }));
558        }
559        let mut holds = Vec::new();
560        let mut granted = 0;
561        for task in tasks {
562            match task.await.unwrap() {
563                Ok(hold) => {
564                    granted += 1;
565                    holds.push(hold);
566                }
567                Err(e) => assert_eq!(e.code, ErrorCode::QuotaExceeded),
568            }
569        }
570        assert_eq!(granted, 1);
571        assert_eq!(a.remaining_pages(&c), Some(0));
572        drop(holds);
573        assert_eq!(a.remaining_pages(&c), Some(1));
574    }
575
576    #[tokio::test]
577    async fn page_quota_survives_a_restart() {
578        let dir = tempfile::tempdir().unwrap();
579        let c = caller(KeyLimits {
580            pages_per_day: Some(5),
581            ..Default::default()
582        });
583        {
584            let store = Store::open_dir(dir.path()).unwrap();
585            let l = Arc::new(Limiter::new(Some(store)));
586            assert_eq!(l.remaining_pages(&c), Some(5));
587            for _ in 0..4 {
588                l.reserve_page(&c).await.unwrap().commit().await;
589            }
590            assert!(l.check_quota(&c).is_ok());
591            l.reserve_page(&c).await.unwrap().commit().await;
592            assert_eq!(
593                l.check_quota(&c).unwrap_err().code,
594                ErrorCode::QuotaExceeded
595            );
596        }
597        let l = Limiter::new(Some(Store::open_dir(dir.path()).unwrap()));
598        assert_eq!(l.remaining_pages(&c), Some(0));
599        let jobs = caller(KeyLimits {
600            max_concurrent_jobs: Some(1),
601            ..Default::default()
602        });
603        assert!(l.check_jobs(&jobs, 0).is_ok());
604        assert_eq!(
605            l.check_jobs(&jobs, 1).unwrap_err().code,
606            ErrorCode::ConcurrencyLimitExceeded
607        );
608    }
609}
pagesift_api/limits.rs

pagesift_api/
limits.rs
