crates/holger-public-site/src/limit.rs
//! **Crawlers never take over the box**: a small budget per client, per kind of
//! client and for the whole site, and a clear 429 with `Retry-After` beyond it.
//!
//! Token buckets, in memory. A request must find a token in all three buckets
//! (its client's, its class's and the global one) or it takes none of them, so
//! one client that is refused does not spend the budget of the others.
use std::collections::HashMap;
use std::sync::Mutex;
use std::time::{Duration, Instant};
/// The kind of client, read off its `User-Agent`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Class {
/// The AI crawlers robots.txt names: GPTBot, ClaudeBot, PerplexityBot,
/// Google-Extended and their kin.
AiCrawler,
/// The classic search engines.
SearchCrawler,
/// Everybody else: browsers, scripts, unknown bots.
Other,
}
/// The AI crawlers robots.txt speaks to by name.
pub const AI_CRAWLERS: &[&str] = &[
"GPTBot",
"ChatGPT-User",
"OAI-SearchBot",
"ClaudeBot",
"Claude-User",
"Claude-SearchBot",
"anthropic-ai",
"PerplexityBot",
"Perplexity-User",
"Google-Extended",
"Applebot-Extended",
"CCBot",
"Meta-ExternalAgent",
"Amazonbot",
"DuckAssistBot",
"MistralAI-User",
];
/// The classic search engines robots.txt speaks to by name.
pub const SEARCH_CRAWLERS: &[&str] = &["Googlebot", "Bingbot", "DuckDuckBot", "Applebot", "YandexBot", "Baiduspider", "SeznamBot"];
/// Which [`Class`] a `User-Agent` falls in.
pub fn classify(user_agent: &str) -> Class {
let ua = user_agent.to_ascii_lowercase();
if AI_CRAWLERS.iter().any(|b| ua.contains(&b.to_ascii_lowercase())) {
Class::AiCrawler
} else if SEARCH_CRAWLERS.iter().any(|b| ua.contains(&b.to_ascii_lowercase())) {
Class::SearchCrawler
} else {
Class::Other
}
}
/// One bucket's size and refill.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Budget {
pub burst: u32,
pub per_second: f64,
}
/// The three budgets.
#[derive(Debug, Clone, PartialEq)]
pub struct RateConfig {
pub per_client: Budget,
pub per_class: Budget,
pub global: Budget,
/// Clients remembered at most; beyond it the idle ones are forgotten.
pub max_clients: usize,
}
impl Default for RateConfig {
fn default() -> RateConfig {
RateConfig {
// A person clicking about, or a crawler that honours Crawl-delay.
per_client: Budget { burst: 20, per_second: 1.0 },
// All AI crawlers together, all search engines together.
per_class: Budget { burst: 60, per_second: 5.0 },
// The whole public site.
global: Budget { burst: 200, per_second: 25.0 },
max_clients: 10_000,
}
}
}
#[derive(Debug, Clone, Copy)]
struct Bucket {
tokens: f64,
last: Instant,
}
impl Bucket {
fn full(budget: Budget, now: Instant) -> Bucket {
Bucket { tokens: f64::from(budget.burst), last: now }
}
fn refill(&mut self, budget: Budget, now: Instant) {
let elapsed = now.saturating_duration_since(self.last).as_secs_f64();
self.tokens = (self.tokens + elapsed * budget.per_second).min(f64::from(budget.burst));
self.last = now;
}
/// Seconds until one token is there, rounded up, at least 1.
fn wait(&self, budget: Budget) -> u64 {
if budget.per_second <= 0.0 {
return 3600;
}
let missing = (1.0 - self.tokens).max(0.0);
((missing / budget.per_second).ceil() as u64).max(1)
}
}
struct Buckets {
clients: HashMap<String, Bucket>,
classes: HashMap<Class, Bucket>,
global: Bucket,
}
/// The limiter itself.
pub struct Limiter {
config: RateConfig,
inner: Mutex<Buckets>,
}
impl Limiter {
pub fn new(config: RateConfig) -> Limiter {
let now = Instant::now();
let global = Bucket::full(config.global, now);
Limiter { config, inner: Mutex::new(Buckets { clients: HashMap::new(), classes: HashMap::new(), global }) }
}
/// Take one token for `client` of `class` at `now`, or say how many
/// seconds to wait.
pub fn check(&self, client: &str, class: Class, now: Instant) -> Result<(), u64> {
let cfg = &self.config;
let mut b = self.inner.lock().unwrap_or_else(|p| p.into_inner());
if b.clients.len() >= cfg.max_clients && !b.clients.contains_key(client) {
// Forget the clients that have been quiet long enough to be full again.
let idle = Duration::from_secs_f64(f64::from(cfg.per_client.burst) / cfg.per_client.per_second.max(0.001));
b.clients.retain(|_, k| now.saturating_duration_since(k.last) < idle);
if b.clients.len() >= cfg.max_clients {
b.clients.clear();
}
}
let mut mine = b.clients.get(client).copied().unwrap_or_else(|| Bucket::full(cfg.per_client, now));
let mut kind = b.classes.get(&class).copied().unwrap_or_else(|| Bucket::full(cfg.per_class, now));
let mut all = b.global;
mine.refill(cfg.per_client, now);
kind.refill(cfg.per_class, now);
all.refill(cfg.global, now);
let refused = [(mine, cfg.per_client), (kind, cfg.per_class), (all, cfg.global)]
.iter()
.filter(|(k, _)| k.tokens < 1.0)
.map(|(k, budget)| k.wait(*budget))
.max();
if refused.is_none() {
mine.tokens -= 1.0;
kind.tokens -= 1.0;
all.tokens -= 1.0;
}
b.clients.insert(client.to_owned(), mine);
b.classes.insert(class, kind);
b.global = all;
match refused {
None => Ok(()),
Some(wait) => Err(wait),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn small() -> RateConfig {
RateConfig {
per_client: Budget { burst: 3, per_second: 1.0 },
per_class: Budget { burst: 5, per_second: 1.0 },
global: Budget { burst: 8, per_second: 1.0 },
max_clients: 4,
}
}
#[test]
fn the_crawlers_are_told_apart_by_their_user_agent() {
assert_eq!(classify("Mozilla/5.0 AppleWebKit/537.36 (KHTML, like Gecko; compatible; GPTBot/1.2; +https://openai.com/gptbot)"), Class::AiCrawler);
assert_eq!(classify("Mozilla/5.0 (compatible; ClaudeBot/1.0; +claudebot@anthropic.com)"), Class::AiCrawler);
assert_eq!(classify("Mozilla/5.0 (compatible; PerplexityBot/1.0)"), Class::AiCrawler);
assert_eq!(classify("Mozilla/5.0 (compatible; Googlebot/2.1; +http://www.google.com/bot.html)"), Class::SearchCrawler);
assert_eq!(classify("Mozilla/5.0 (compatible; bingbot/2.0)"), Class::SearchCrawler);
assert_eq!(classify("Mozilla/5.0 (X11; Linux x86_64) Firefox/131.0"), Class::Other);
assert_eq!(classify(""), Class::Other);
}
#[test]
fn a_burst_from_one_client_is_refused_with_a_wait() {
let l = Limiter::new(small());
let now = Instant::now();
for _ in 0..3 {
assert_eq!(l.check("a", Class::Other, now), Ok(()));
}
assert_eq!(l.check("a", Class::Other, now), Err(1));
// Another client still gets in: the refusal spent nobody else's budget.
assert_eq!(l.check("b", Class::Other, now), Ok(()));
// And a second later the first is let in again.
assert_eq!(l.check("a", Class::Other, now + Duration::from_secs(1)), Ok(()));
}
#[test]
fn many_clients_of_one_class_share_its_budget() {
let l = Limiter::new(small());
let now = Instant::now();
for i in 0..5 {
assert_eq!(l.check(&format!("bot{i}"), Class::AiCrawler, now), Ok(()), "bot{i}");
}
assert!(l.check("bot9", Class::AiCrawler, now).is_err(), "the class budget is spent");
assert_eq!(l.check("person", Class::Other, now), Ok(()), "another class is not");
}
#[test]
fn the_whole_site_has_a_ceiling() {
let l = Limiter::new(small());
let now = Instant::now();
let mut ok = 0;
for i in 0..20 {
let class = [Class::AiCrawler, Class::SearchCrawler, Class::Other][i % 3];
if l.check(&format!("c{i}"), class, now).is_ok() {
ok += 1;
}
}
assert_eq!(ok, 8, "the global burst caps every class together");
}
#[test]
fn the_client_table_is_bounded() {
let l = Limiter::new(RateConfig { global: Budget { burst: 1000, per_second: 1.0 }, per_class: Budget { burst: 1000, per_second: 1.0 }, ..small() });
let now = Instant::now();
for i in 0..50 {
let _ = l.check(&format!("c{i}"), Class::Other, now);
}
assert!(l.inner.lock().unwrap().clients.len() <= 4);
}
}