mod policy;
use std::collections::HashMap;
use std::collections::hash_map::RandomState;
use std::hash::{BuildHasher, Hasher};
use std::num::NonZeroU32;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use governor::clock::DefaultClock;
use governor::state::{InMemoryState, NotKeyed};
use governor::{Jitter, Quota, RateLimiter};
use tokio::sync::{Mutex, OwnedSemaphorePermit, Semaphore};
use tokio::time::Instant;
pub use policy::{
CAUTIOUS_LIMIT, DEFAULT_RETRY_AFTER, HostLimit, MAX_RETRY_AFTER, RESOLVER_HOST, RESOLVER_LIMIT,
clamp_retry_after, published_limit, starting_limit,
};
const PROBE_DEADLINE: Duration = Duration::from_secs(30);
type TokenBucket = RateLimiter<NotKeyed, InMemoryState, DefaultClock>;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Refusal {
Throttled { retry_after: Option<Duration> },
Blocked,
Dropped,
}
impl Refusal {
#[must_use]
fn base_wait(&self) -> Duration {
match self {
Self::Throttled { retry_after } => clamp_retry_after(*retry_after),
Self::Blocked => MAX_RETRY_AFTER,
Self::Dropped => DEFAULT_RETRY_AFTER,
}
}
#[must_use]
const fn has_stated_wait(&self) -> bool {
matches!(
self,
Self::Throttled {
retry_after: Some(_)
}
)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PacingLimits {
pub total_concurrency: usize,
pub refusal_threshold: u32,
pub recovery_threshold: u32,
pub jitter: Duration,
pub max_backoff: Duration,
pub is_cautious: bool,
pub per_registry: Option<usize>,
pub rate: Option<u32>,
}
impl Default for PacingLimits {
fn default() -> Self {
Self {
total_concurrency: 24,
refusal_threshold: 3,
recovery_threshold: 8,
jitter: Duration::from_millis(120),
max_backoff: Duration::from_secs(120),
is_cautious: false,
per_registry: None,
rate: None,
}
}
}
impl PacingLimits {
#[must_use]
pub fn cautious() -> Self {
Self {
total_concurrency: 8,
refusal_threshold: 2,
recovery_threshold: 16,
jitter: Duration::from_millis(300),
is_cautious: true,
..Self::default()
}
}
fn limit_for(&self, host: &str) -> HostLimit {
let mut limit = if self.is_cautious {
CAUTIOUS_LIMIT
} else {
starting_limit(host)
};
if let Some(concurrency) = self.per_registry {
limit.concurrency = concurrency.max(1);
}
if let Some(queries) = self.rate {
limit.queries = queries.max(1);
limit.window = Duration::from_secs(1);
}
if self.is_cautious {
limit.concurrency = limit.concurrency.min(CAUTIOUS_LIMIT.concurrency);
if limit.per_second_rate() > CAUTIOUS_LIMIT.per_second_rate() {
limit.queries = CAUTIOUS_LIMIT.queries;
limit.window = CAUTIOUS_LIMIT.window;
}
}
limit
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BreakerState {
Closed,
Open,
HalfOpen,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PausedHost {
pub host: String,
pub remaining_wait: Duration,
pub refusals: u32,
}
#[derive(Debug)]
pub struct RequestPermit {
host: String,
is_probe: bool,
_global_slot: OwnedSemaphorePermit,
_host_slot: OwnedSemaphorePermit,
}
impl RequestPermit {
#[must_use]
pub const fn is_probe(&self) -> bool {
self.is_probe
}
#[must_use]
pub fn host(&self) -> &str {
&self.host
}
}
#[derive(Debug)]
struct HostState {
limiter: Arc<TokenBucket>,
slots: Arc<Semaphore>,
concurrency: usize,
max_concurrency: usize,
refusals: u32,
successes: u32,
open_until: Option<Instant>,
is_probing: bool,
}
fn usable_permits(requested: usize) -> usize {
requested.clamp(1, Semaphore::MAX_PERMITS)
}
impl HostState {
fn new(limit: HostLimit) -> Self {
Self {
limiter: Arc::new(RateLimiter::direct(quota_for(limit))),
slots: Arc::new(Semaphore::new(usable_permits(limit.concurrency))),
concurrency: limit.concurrency.max(1),
max_concurrency: limit.concurrency.max(1),
refusals: 0,
successes: 0,
open_until: None,
is_probing: false,
}
}
fn breaker(&self, now: Instant) -> BreakerState {
match self.open_until {
Some(until) if now < until => BreakerState::Open,
Some(_) => BreakerState::HalfOpen,
None => {
if self.is_probing {
BreakerState::HalfOpen
} else {
BreakerState::Closed
}
}
}
}
}
fn quota_for(limit: HostLimit) -> Quota {
let rate = limit.per_second_rate().max(0.05);
let interval = Duration::from_secs_f64(1.0 / rate);
let burst = NonZeroU32::new(limit.queries.max(1)).unwrap_or(NonZeroU32::MIN);
Quota::with_period(interval).map_or_else(
|| Quota::per_second(NonZeroU32::MIN),
|quota| quota.allow_burst(burst),
)
}
#[derive(Debug)]
struct JitterRng(AtomicU64);
impl JitterRng {
fn new() -> Self {
let seed = RandomState::new().build_hasher().finish() | 1;
Self(AtomicU64::new(seed))
}
fn next_u64(&self) -> u64 {
let mut x = self.0.load(Ordering::Relaxed);
x ^= x << 13;
x ^= x >> 7;
x ^= x << 17;
self.0.store(x, Ordering::Relaxed);
x
}
fn sample_up_to(&self, ceiling: Duration) -> Duration {
let nanos = u64::try_from(ceiling.as_nanos()).unwrap_or(u64::MAX);
if nanos == 0 {
return Duration::ZERO;
}
Duration::from_nanos(self.next_u64() % nanos.saturating_add(1))
}
}
#[derive(Debug)]
pub struct Pacer {
pacing: PacingLimits,
global_slots: Arc<Semaphore>,
hosts: Mutex<HashMap<String, HostState>>,
jitter: JitterRng,
}
impl Pacer {
#[must_use]
pub fn new(pacing: PacingLimits) -> Self {
Self {
global_slots: Arc::new(Semaphore::new(usable_permits(pacing.total_concurrency))),
hosts: Mutex::new(HashMap::new()),
jitter: JitterRng::new(),
pacing,
}
}
#[must_use]
pub const fn pacing(&self) -> &PacingLimits {
&self.pacing
}
pub async fn acquire_patiently(
&self,
host: &str,
budget: Duration,
) -> Result<RequestPermit, PausedHost> {
match self.acquire(host).await {
Ok(permit) => Ok(permit),
Err(paused) => {
let wait = paused.remaining_wait;
if wait.is_zero() || wait > budget {
return Err(paused);
}
tokio::time::sleep(wait).await;
self.acquire(host).await
}
}
}
pub async fn acquire(&self, host: &str) -> Result<RequestPermit, PausedHost> {
let key = normalize_host(host);
let (limiter, slots, probe) = {
let mut hosts = self.hosts.lock().await;
let state = hosts
.entry(key.clone())
.or_insert_with(|| HostState::new(self.pacing.limit_for(&key)));
let now = Instant::now();
match state.breaker(now) {
BreakerState::Open => {
let remaining = state
.open_until
.map_or(Duration::ZERO, |until| until.saturating_duration_since(now));
return Err(PausedHost {
host: key,
remaining_wait: remaining,
refusals: state.refusals,
});
}
BreakerState::HalfOpen => {
if state.is_probing {
return Err(PausedHost {
host: key,
remaining_wait: Duration::ZERO,
refusals: state.refusals,
});
}
state.is_probing = true;
state.open_until = Some(now + PROBE_DEADLINE);
(Arc::clone(&state.limiter), Arc::clone(&state.slots), true)
}
BreakerState::Closed => {
(Arc::clone(&state.limiter), Arc::clone(&state.slots), false)
}
}
};
if self.pacing.jitter.is_zero() {
limiter.until_ready().await;
} else {
limiter
.until_ready_with_jitter(Jitter::up_to(self.pacing.jitter))
.await;
}
let host_permit = slots
.acquire_owned()
.await
.map_err(|_| self.shutdown_pause(&key))?;
let global = Arc::clone(&self.global_slots)
.acquire_owned()
.await
.map_err(|_| self.shutdown_pause(&key))?;
Ok(RequestPermit {
host: key,
is_probe: probe,
_global_slot: global,
_host_slot: host_permit,
})
}
fn shutdown_pause(&self, host: &str) -> PausedHost {
PausedHost {
host: host.to_owned(),
remaining_wait: Duration::ZERO,
refusals: 0,
}
}
pub async fn record_success(&self, host: &str) {
let key = normalize_host(host);
let mut hosts = self.hosts.lock().await;
let Some(state) = hosts.get_mut(&key) else {
return;
};
state.refusals = 0;
state.open_until = None;
state.is_probing = false;
state.successes = state.successes.saturating_add(1);
if state.successes >= self.pacing.recovery_threshold
&& state.concurrency < state.max_concurrency
{
state.successes = 0;
state.concurrency = state
.concurrency
.saturating_add(1)
.min(state.max_concurrency);
state.slots.add_permits(1);
tracing::debug!(host = %key, concurrency = state.concurrency, "registry concurrency raised");
}
}
pub async fn record_refusal(&self, host: &str, refusal: &Refusal) -> Duration {
let key = normalize_host(host);
let mut hosts = self.hosts.lock().await;
let state = hosts
.entry(key.clone())
.or_insert_with(|| HostState::new(self.pacing.limit_for(&key)));
state.refusals = state.refusals.saturating_add(1);
state.successes = 0;
state.is_probing = false;
let wait = if refusal.has_stated_wait() {
refusal.base_wait()
} else {
let exponent = state.refusals.saturating_sub(1).min(6);
let ceiling = refusal
.base_wait()
.saturating_mul(1_u32 << exponent)
.min(self.pacing.max_backoff);
self.jitter.sample_up_to(ceiling)
};
let wait = wait.min(self.pacing.max_backoff);
if state.concurrency > 1 {
let target = state
.concurrency
.saturating_mul(3)
.div_ceil(4)
.min(state.concurrency.saturating_sub(1))
.max(1);
let surplus = state.concurrency.saturating_sub(target);
if surplus > 0 {
let forgotten = state.slots.forget_permits(surplus);
state.concurrency = state.concurrency.saturating_sub(forgotten);
tracing::debug!(host = %key, concurrency = state.concurrency, "registry concurrency cut");
}
}
let should_open =
matches!(refusal, Refusal::Blocked) || state.refusals >= self.pacing.refusal_threshold;
if should_open {
state.open_until = Some(Instant::now() + wait);
tracing::debug!(host = %key, refusals = state.refusals, ?wait, "registry paused");
}
wait
}
pub async fn breaker(&self, host: &str) -> BreakerState {
let key = normalize_host(host);
let hosts = self.hosts.lock().await;
hosts
.get(&key)
.map_or(BreakerState::Closed, |state| state.breaker(Instant::now()))
}
pub async fn paused(&self, host: &str) -> Option<PausedHost> {
let key = normalize_host(host);
let mut hosts = self.hosts.lock().await;
let state = hosts.get_mut(&key)?;
let open_until = state.open_until?;
let now = Instant::now();
if now >= open_until {
return None;
}
Some(PausedHost {
host: key,
remaining_wait: open_until.saturating_duration_since(now),
refusals: state.refusals,
})
}
pub async fn paused_hosts(&self) -> Vec<PausedHost> {
let now = Instant::now();
let hosts = self.hosts.lock().await;
let mut paused: Vec<PausedHost> = hosts
.iter()
.filter_map(|(host, state)| {
let open_until = state.open_until?;
(open_until > now).then(|| PausedHost {
host: host.clone(),
remaining_wait: open_until.saturating_duration_since(now),
refusals: state.refusals,
})
})
.collect();
paused.sort_by(|a, b| a.host.cmp(&b.host));
paused
}
pub async fn host_concurrency(&self, host: &str) -> usize {
let key = normalize_host(host);
let hosts = self.hosts.lock().await;
hosts.get(&key).map_or(0, |state| state.concurrency)
}
pub async fn prune_settled_hosts(&self) {
let now = Instant::now();
let mut hosts = self.hosts.lock().await;
hosts.retain(|_, state| {
state.refusals > 0
|| state.is_probing
|| state.concurrency < state.max_concurrency
|| state.open_until.is_some_and(|until| until > now)
});
}
}
fn normalize_host(host: &str) -> String {
host.trim().trim_end_matches('.').to_lowercase()
}
#[cfg(test)]
mod tests {
#[test]
fn cautious_cannot_be_raised_by_a_flag_that_asks_for_more() {
let reckless = PacingLimits {
rate: Some(500),
per_registry: Some(64),
..PacingLimits::cautious()
};
let limit = reckless.limit_for("rdap.example");
assert!(
limit.per_second_rate() <= CAUTIOUS_LIMIT.per_second_rate(),
"cautious means the rate can only go down, never up"
);
assert!(limit.concurrency <= CAUTIOUS_LIMIT.concurrency);
}
#[test]
fn cautious_still_lets_a_flag_ask_for_less() {
let slower = PacingLimits {
rate: Some(1),
per_registry: Some(1),
..PacingLimits::cautious()
};
let limit = slower.limit_for("rdap.example");
assert_eq!(limit.queries, 1);
assert_eq!(limit.concurrency, 1);
}
use super::*;
fn throttled(seconds: u64) -> Refusal {
Refusal::Throttled {
retry_after: Some(Duration::from_secs(seconds)),
}
}
const UNSTATED: Refusal = Refusal::Throttled { retry_after: None };
fn fast() -> PacingLimits {
PacingLimits {
jitter: Duration::ZERO,
..PacingLimits::default()
}
}
#[tokio::test(start_paused = true)]
async fn a_lease_is_granted_and_released() {
let pacer = Pacer::new(fast());
let lease = pacer
.acquire("rdap.example.test")
.await
.expect("first lease");
assert_eq!(lease.host(), "rdap.example.test");
assert!(!lease.is_probe());
drop(lease);
assert!(pacer.acquire("rdap.example.test").await.is_ok());
}
#[tokio::test(start_paused = true)]
async fn hosts_are_keyed_case_and_dot_insensitively() {
let pacer = Pacer::new(fast());
pacer
.record_refusal("RDAP.Example.Test.", &throttled(30))
.await;
pacer
.record_refusal("rdap.example.test", &throttled(30))
.await;
pacer
.record_refusal("rdap.example.test", &throttled(30))
.await;
assert!(pacer.paused("rdap.example.test").await.is_some());
}
#[tokio::test(start_paused = true)]
async fn a_published_registry_starts_with_its_published_allowance() {
let pacer = Pacer::new(fast());
let _lease = pacer
.acquire("rdap.identitydigital.services")
.await
.unwrap();
assert_eq!(
pacer
.host_concurrency("rdap.identitydigital.services")
.await,
4
);
}
#[tokio::test(start_paused = true)]
async fn gentle_pacing_ignores_a_generous_published_allowance() {
let pacer = Pacer::new(PacingLimits {
jitter: Duration::ZERO,
..PacingLimits::cautious()
});
let _lease = pacer
.acquire("rdap.identitydigital.services")
.await
.unwrap();
assert_eq!(
pacer
.host_concurrency("rdap.identitydigital.services")
.await,
CAUTIOUS_LIMIT.concurrency
);
}
#[tokio::test(start_paused = true)]
async fn the_breaker_stays_shut_until_the_threshold_is_reached() {
let pacer = Pacer::new(PacingLimits {
refusal_threshold: 3,
..fast()
});
pacer.record_refusal("slow.test", &UNSTATED).await;
assert_eq!(pacer.breaker("slow.test").await, BreakerState::Closed);
pacer.record_refusal("slow.test", &UNSTATED).await;
assert_eq!(pacer.breaker("slow.test").await, BreakerState::Closed);
pacer.record_refusal("slow.test", &UNSTATED).await;
assert_eq!(pacer.breaker("slow.test").await, BreakerState::Open);
}
#[tokio::test(start_paused = true)]
async fn an_outright_block_opens_the_breaker_on_the_first_refusal() {
let pacer = Pacer::new(PacingLimits {
refusal_threshold: 99,
..fast()
});
pacer
.record_refusal("blocked.test", &Refusal::Blocked)
.await;
assert_eq!(pacer.breaker("blocked.test").await, BreakerState::Open);
}
#[tokio::test(start_paused = true)]
async fn a_dropped_connection_counts_as_backpressure() {
let pacer = Pacer::new(PacingLimits {
refusal_threshold: 1,
..fast()
});
pacer.record_refusal("quiet.test", &Refusal::Dropped).await;
assert_eq!(pacer.breaker("quiet.test").await, BreakerState::Open);
}
#[tokio::test(start_paused = true)]
async fn an_open_breaker_refuses_a_lease_instead_of_blocking() {
let pacer = Pacer::new(PacingLimits {
refusal_threshold: 1,
..fast()
});
pacer.record_refusal("busy.test", &throttled(5)).await;
match pacer.acquire("busy.test").await {
Err(paused) => {
assert_eq!(paused.host, "busy.test");
assert!(paused.remaining_wait <= Duration::from_secs(5));
}
Ok(_) => panic!("an open breaker must refuse the lease"),
}
}
#[tokio::test(start_paused = true)]
async fn the_cooldown_lets_exactly_one_probe_through() {
let pacer = Pacer::new(PacingLimits {
refusal_threshold: 1,
..fast()
});
pacer.record_refusal("recovering.test", &throttled(2)).await;
tokio::time::advance(Duration::from_secs(3)).await;
assert_eq!(
pacer.breaker("recovering.test").await,
BreakerState::HalfOpen
);
let probe = pacer
.acquire("recovering.test")
.await
.expect("probe allowed");
assert!(probe.is_probe());
assert!(
pacer.acquire("recovering.test").await.is_err(),
"only one probe may be in flight"
);
}
#[tokio::test(start_paused = true)]
async fn a_successful_probe_closes_the_breaker() {
let pacer = Pacer::new(PacingLimits {
refusal_threshold: 1,
..fast()
});
pacer.record_refusal("healing.test", &throttled(2)).await;
tokio::time::advance(Duration::from_secs(3)).await;
let probe = pacer.acquire("healing.test").await.expect("probe");
drop(probe);
pacer.record_success("healing.test").await;
assert_eq!(pacer.breaker("healing.test").await, BreakerState::Closed);
assert!(pacer.acquire("healing.test").await.is_ok());
}
#[tokio::test(start_paused = true)]
async fn a_failed_probe_reopens_the_breaker() {
let pacer = Pacer::new(PacingLimits {
refusal_threshold: 1,
..fast()
});
pacer.record_refusal("stubborn.test", &throttled(2)).await;
tokio::time::advance(Duration::from_secs(3)).await;
let probe = pacer.acquire("stubborn.test").await.expect("probe");
drop(probe);
pacer.record_refusal("stubborn.test", &throttled(4)).await;
assert_eq!(pacer.breaker("stubborn.test").await, BreakerState::Open);
}
#[tokio::test(start_paused = true)]
async fn a_stated_wait_is_honored_exactly() {
let pacer = Pacer::new(fast());
let wait = pacer.record_refusal("polite.test", &throttled(7)).await;
assert_eq!(wait, Duration::from_secs(7));
}
#[tokio::test(start_paused = true)]
async fn an_absurd_stated_wait_is_capped() {
let pacer = Pacer::new(fast());
let wait = pacer
.record_refusal("hostile.test", &throttled(86_400))
.await;
assert_eq!(wait, MAX_RETRY_AFTER);
}
#[tokio::test(start_paused = true)]
async fn an_unstated_wait_stays_inside_the_growing_ceiling() {
let pacer = Pacer::new(PacingLimits {
refusal_threshold: 99,
..fast()
});
for _ in 0..12 {
let wait = pacer.record_refusal("steep.test", &UNSTATED).await;
assert!(
wait <= pacer.pacing().max_backoff,
"{wait:?} exceeded the cap"
);
}
}
#[tokio::test(start_paused = true)]
async fn a_refusal_cuts_the_allowance_and_success_earns_it_back() {
let pacer = Pacer::new(PacingLimits {
refusal_threshold: 99,
recovery_threshold: 2,
..fast()
});
let host = "rdap.identitydigital.services";
let lease = pacer.acquire(host).await.unwrap();
drop(lease);
assert_eq!(pacer.host_concurrency(host).await, 4);
pacer.record_refusal(host, &UNSTATED).await;
assert_eq!(pacer.host_concurrency(host).await, 3);
for _ in 0..2 {
pacer.record_success(host).await;
}
assert_eq!(pacer.host_concurrency(host).await, 4);
}
#[tokio::test(start_paused = true)]
async fn a_refusal_still_cuts_a_host_that_starts_at_two() {
let pacer = Pacer::new(PacingLimits {
refusal_threshold: 99,
..fast()
});
let host = "unpublished.test";
let lease = pacer.acquire(host).await.unwrap();
drop(lease);
assert_eq!(pacer.host_concurrency(host).await, 2);
pacer.record_refusal(host, &UNSTATED).await;
assert_eq!(
pacer.host_concurrency(host).await,
1,
"backpressure must reach a host that starts at the cautious limit"
);
}
#[tokio::test(start_paused = true)]
async fn the_allowance_never_climbs_past_where_it_started() {
let pacer = Pacer::new(PacingLimits {
recovery_threshold: 1,
..fast()
});
let host = "rdap.identitydigital.services";
let lease = pacer.acquire(host).await.unwrap();
drop(lease);
for _ in 0..50 {
pacer.record_success(host).await;
}
assert_eq!(pacer.host_concurrency(host).await, 4);
}
#[tokio::test(start_paused = true)]
async fn paused_hosts_lists_only_the_registries_actually_on_hold() {
let pacer = Pacer::new(PacingLimits {
refusal_threshold: 1,
..fast()
});
pacer.record_refusal("one.test", &throttled(10)).await;
pacer.record_success("two.test").await;
let paused = pacer.paused_hosts().await;
assert_eq!(paused.len(), 1);
assert_eq!(paused.first().map(|p| p.host.as_str()), Some("one.test"));
}
#[tokio::test(start_paused = true)]
async fn tidy_keeps_troubled_hosts_and_drops_settled_ones() {
let pacer = Pacer::new(PacingLimits {
refusal_threshold: 1,
..fast()
});
let settled = pacer.acquire("calm.test").await.unwrap();
drop(settled);
pacer.record_success("calm.test").await;
pacer.record_refusal("angry.test", &throttled(60)).await;
pacer.prune_settled_hosts().await;
assert_eq!(pacer.host_concurrency("calm.test").await, 0);
assert!(pacer.paused("angry.test").await.is_some());
}
#[test]
fn full_jitter_draws_inside_the_window_and_actually_varies() {
let jitterer = JitterRng::new();
let ceiling = Duration::from_secs(10);
let draws: Vec<Duration> = (0..64).map(|_| jitterer.sample_up_to(ceiling)).collect();
assert!(draws.iter().all(|d| *d <= ceiling));
let unique = draws
.iter()
.collect::<std::collections::BTreeSet<_>>()
.len();
assert!(
unique > 32,
"jitter is not spreading: {unique} distinct draws"
);
}
#[test]
fn a_zero_window_yields_no_wait_rather_than_panicking() {
let jitterer = JitterRng::new();
assert_eq!(jitterer.sample_up_to(Duration::ZERO), Duration::ZERO);
}
#[test]
fn a_quota_survives_an_extremely_slow_published_limit() {
let quota = quota_for(HostLimit::per_minute(1, 1));
assert!(quota.burst_size().get() >= 1);
}
}