Something went wrong. Try again.
fork of hydrant
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461use crate::config::RateTier;use parking_lot::Mutex;use scc::HashMap;use std::future::Future;use std::sync::Arc;use std::sync::atomic::{AtomicI64, AtomicUsize, Ordering};use std::time::{Duration, Instant};use tokio::sync::{Notify, Semaphore, SemaphorePermit};use url::Url;
/// default max concurrent in-flight requests per PDS *per egress IP* before we queue./// ref pds allows ~10 requests per second per IP. the effective per-PDS limit is this/// multiplied by the egress pool size (direct + proxies). overridable via/// `HYDRANT_PER_PDS_CONCURRENCY`; raise it so rate-tier pacing (not concurrency) is the/// binding limit when the pool has high round-trip latency.pub const DEFAULT_PER_PDS_CONCURRENCY_PER_IP: usize = 10;
// per second, hour and dayconst DURATIONS: [Duration; 3] = [ Duration::from_secs(1), Duration::from_secs(3600), Duration::from_secs(86400),];
#[derive(Clone)]pub struct Throttler { states: Arc<HashMap<Url, Arc<State>>>, /// number of egress IPs (direct + proxies) sharing this throttler. multiplies the /// per-PDS concurrency and rate-tier pacing so tiers describe a single IP's budget. pool_size: usize, /// per-egress-IP in-flight concurrency; the per-PDS semaphore is this times `pool_size`. per_pds_base: usize,}
#[derive(Debug, Clone, Default)]pub struct ThrottleSnapshot { pub consecutive_failures: usize, pub throttled_until: i64, pub last_failure: Option<FailureSnapshot>,}
impl ThrottleSnapshot { pub fn is_failing(&self) -> bool { self.consecutive_failures != 0 || self.throttled_until != 0 }
pub fn is_throttled(&self, now: i64) -> bool { self.throttled_until != 0 && now < self.throttled_until }
pub fn retry_in_secs(&self, now: i64) -> Option<u64> { self.is_throttled(now) .then_some((self.throttled_until - now) as u64) }}
#[derive(Debug, Clone)]pub struct FailureSnapshot { pub at: i64, pub kind: String, pub detail: String,}
/// Canonical throttle identity: scheme, normalized host, and effective port only./// PDS paths, queries, fragments, credentials, and trailing-dot spellings must/// never mint distinct rate-limit buckets.fn canonical_origin(url: &Url) -> Url { let mut origin = url.clone(); origin.set_path("/"); origin.set_query(None); origin.set_fragment(None); let _ = origin.set_username(""); let _ = origin.set_password(None); if let Some(host) = origin.host_str().map(str::to_owned) { let host = host.strip_suffix('.').unwrap_or(&host); let _ = origin.set_host(Some(host)); } if let Some(port) = origin.port_or_known_default() { let _ = origin.set_port(Some(port)); } origin}
impl Throttler { /// `pool_size` is the number of egress IPs (direct connection + proxies) that share this /// throttler; it multiplies per-PDS concurrency and rate-tier pacing. pass `1` when there /// are no proxies. pub fn new(pool_size: usize, per_pds_base: usize) -> Self { Self { states: Arc::new(HashMap::new()), pool_size: pool_size.max(1), per_pds_base: per_pds_base.max(1), } }
pub async fn get_handle(&self, url: &Url) -> ThrottleHandle { let pool_size = self.pool_size; let concurrency = self.per_pds_base * pool_size; let state = self .states .entry_async(canonical_origin(url)) .await .or_insert_with(|| Arc::new(State::new(concurrency))) .get() .clone();
ThrottleHandle { state, pool_size } }
#[cfg(feature = "indexer")] pub(crate) fn per_pds_concurrency(&self) -> usize { self.per_pds_base * self.pool_size }
/// drop entries with no active throttle and no consecutive failures. pub async fn evict_clean(&self) { self.states .retain_async(|_, v| { v.throttled_until.load(Ordering::Acquire) != 0 || v.consecutive_failures.load(Ordering::Acquire) != 0 }) .await; }
pub fn snapshot(&self, url: &Url) -> ThrottleSnapshot { self.states .read_sync(&canonical_origin(url), |_, state| ThrottleSnapshot { consecutive_failures: state.consecutive_failures.load(Ordering::Acquire), throttled_until: state.throttled_until.load(Ordering::Acquire), last_failure: state.last_failure.lock().clone(), }) .unwrap_or_default() }}
struct State { throttled_until: AtomicI64, consecutive_failures: AtomicUsize, consecutive_timeouts: AtomicUsize, last_failure: Mutex<Option<FailureSnapshot>>, /// only fires on hard failures (timeout, TLS, bad gateway, etc). /// ratelimits do NOT fire this — they just store `throttled_until` and /// let tasks exit naturally, deferring to the background retry loop. failure_notify: Notify, semaphore: Semaphore, rate_limiter: RateLimiter,}
impl State { fn new(concurrency: usize) -> Self { Self { throttled_until: AtomicI64::new(0), consecutive_failures: AtomicUsize::new(0), consecutive_timeouts: AtomicUsize::new(0), last_failure: Mutex::new(None), failure_notify: Notify::new(), semaphore: Semaphore::new(concurrency), rate_limiter: RateLimiter::new(), } }}
pub struct ThrottleHandle { state: Arc<State>, pool_size: usize,}
impl ThrottleHandle { pub fn is_throttled(&self) -> bool { let until = self.state.throttled_until.load(Ordering::Acquire); until != 0 && chrono::Utc::now().timestamp() < until }
/// the unix timestamp at which this throttle expires (0 if not throttled). pub fn throttled_until(&self) -> i64 { self.state.throttled_until.load(Ordering::Acquire) }
pub fn record_success(&self) { self.state.consecutive_failures.store(0, Ordering::Release); self.state.consecutive_timeouts.store(0, Ordering::Release); self.state.throttled_until.store(0, Ordering::Release); *self.state.last_failure.lock() = None; }
/// called on a 429 response. `retry_after_secs` comes from the `Retry-After` /// header if present; falls back to 60s. uses `fetch_max` so concurrent callers /// don't race each other back to a shorter window. pub fn record_ratelimit(&self, retry_after_secs: Option<u64>) { let secs = retry_after_secs.unwrap_or(60) as i64; let until = chrono::Utc::now().timestamp() + secs; self.state .throttled_until .fetch_max(until, Ordering::AcqRel); }
/// called on hard failures (timeout, TLS error, bad gateway, etc). /// always increments `consecutive_failures`. only sets a new `throttled_until` /// (and notifies waiters) if not already throttled. pub fn record_failure(&self) -> Option<u64> { self.record_failure_inner() }
pub fn record_failure_detail( &self, kind: impl Into<String>, detail: impl Into<String>, ) -> Option<u64> { *self.state.last_failure.lock() = Some(FailureSnapshot { at: chrono::Utc::now().timestamp(), kind: kind.into(), detail: detail.into(), }); self.record_failure_inner() }
fn record_failure_inner(&self) -> Option<u64> { let failures = self .state .consecutive_failures .fetch_add(1, Ordering::AcqRel) + 1;
if self.is_throttled() { return None; }
let base_secs = 15u64; let exponent = (failures as u32).saturating_sub(1); let secs = (base_secs * 2u64.pow(exponent.min(10))).min(300); #[cfg(debug_assertions)] let secs = secs.min(1);
let until = chrono::Utc::now().timestamp() + secs as i64;
self.state.throttled_until.store(until, Ordering::Release); self.state.failure_notify.notify_waiters();
Some(secs) }
/// returns current timeout duration — 3s, 6s, or 12s depending on prior timeouts. pub fn timeout(&self) -> Duration { let n = self.state.consecutive_timeouts.load(Ordering::Acquire); Duration::from_secs(3 * 2u64.pow(n.min(2) as u32)) }
pub fn consecutive_failures(&self) -> usize { self.state.consecutive_failures.load(Ordering::Acquire) }
/// returns whether the timeout attempts are exhausted pub fn record_timeout(&self) -> bool { let timeouts = self .state .consecutive_timeouts .fetch_add(1, Ordering::AcqRel) + 1; timeouts > 2 }
/// acquire a concurrency slot for this PDS. hold the returned permit /// for the duration of the request. pub async fn acquire(&self) -> SemaphorePermit<'_> { self.state .semaphore .acquire() .await .expect("throttle semaphore unexpectedly closed") }
/// resolves when this PDS gets a hard failure notification. pub async fn wait_for_failure(&self) { loop { let notified = self.state.failure_notify.notified(); if self.is_throttled() { return; } notified.await; } }
/// waits until the rate tier's limits allow more events for this PDS. /// sleeps precisely until the most restrictive window opens rather than polling. pub async fn wait_for_allow(&self, num_accounts: u64, tier: &RateTier) { let limits = limits_for(num_accounts, tier, self.pool_size); while let Some(wait) = self.state.rate_limiter.try_acquire(limits) { tokio::time::sleep(wait).await; } }}
/// computes the `[per_second, per_hour, per_day]` window limits for a PDS. the tier describes/// the budget for a single egress IP, so every window scales by `pool_size` (direct + proxies).fn limits_for(num_accounts: u64, tier: &RateTier, pool_size: usize) -> [u64; 3] { let pool = pool_size.max(1) as u64; let per_sec = tier .per_second_base .max((num_accounts as f64 * tier.per_second_account_mul) as u64); [ per_sec.saturating_mul(pool), tier.per_hour.saturating_mul(pool), tier.per_day.saturating_mul(pool), ]}
struct WindowState { count: u64, prev_count: u64, window_start: Instant,}
impl WindowState { fn new() -> Self { Self { count: 0, prev_count: 0, window_start: Instant::now(), } }
fn rotate(&mut self, dur: Duration) { let elapsed = self.window_start.elapsed(); if elapsed >= dur { let n = (elapsed.as_nanos() / dur.as_nanos()).max(1) as u32; self.prev_count = if n == 1 { self.count } else { 0 }; self.count = 0; self.window_start += dur * n; } }
/// returns how long to sleep before this window would allow one more event. /// Duration::ZERO means allow now. fn wait_needed(&self, dur: Duration, limit: u64) -> Duration { let elapsed = self.window_start.elapsed(); let remaining = dur.saturating_sub(elapsed); let weight = remaining.as_secs_f64() / dur.as_secs_f64(); let effective = self.count as f64 + self.prev_count as f64 * weight;
if effective < limit as f64 { return Duration::ZERO; }
if self.prev_count == 0 || self.count as f64 >= limit as f64 { // must wait for a full window rotation remaining + Duration::from_millis(1) } else { let secs = remaining.as_secs_f64() - dur.as_secs_f64() * (limit as f64 - self.count as f64) / self.prev_count as f64; Duration::from_secs_f64(secs.max(0.0)) + Duration::from_micros(500) } }}
struct RateLimiter { // parking_lot::Mutex — uncontended path never touches the kernel windows: Mutex<[WindowState; 3]>,}
impl RateLimiter { fn new() -> Self { Self { windows: Mutex::new([WindowState::new(), WindowState::new(), WindowState::new()]), } }
/// returns None if the slot was acquired, or Some(sleep_for) if limited. fn try_acquire(&self, limits: [u64; 3]) -> Option<Duration> { let mut windows = self.windows.lock();
windows .iter_mut() .zip(DURATIONS) .for_each(|(w, d)| w.rotate(d));
let max_wait = windows .iter() .zip(DURATIONS) .zip(limits) .map(|((w, dur), limit)| w.wait_needed(dur, limit)) .max() .unwrap_or(Duration::ZERO);
if max_wait.is_zero() { windows.iter_mut().for_each(|w| w.count += 1); None } else { Some(max_wait) } }}
/// adds a method for racing the future against a hard-failure notification.#[allow(async_fn_in_trait)]pub trait OrFailure<T, E>: Future<Output = Result<T, E>> { async fn or_failure( self, handle: &ThrottleHandle, on_throttle: impl FnOnce() -> E, ) -> Result<T, E> where Self: Sized, { tokio::select! { res = self => res, _ = handle.wait_for_failure() => Err(on_throttle()), } }}
impl<T, E, F: Future<Output = Result<T, E>>> OrFailure<T, E> for F {}
#[cfg(test)]mod tests { use super::*;
#[tokio::test] async fn throttle_keys_canonicalize_to_origin() { let throttler = Throttler::new(1, 10); let first = Url::parse("https://user:pass@PDS.example.:443/one?cursor=a#frag").unwrap(); let second = Url::parse("https://pds.example/two").unwrap(); throttler.get_handle(&first).await.record_failure(); assert!(throttler.snapshot(&second).is_failing()); }
#[test] fn limits_for_pool_size_one_is_unscaled() { let tier = RateTier::default_tier(); let base = limits_for(0, &tier, 1); assert_eq!(base[0], tier.per_second_base); assert_eq!(base[1], tier.per_hour); assert_eq!(base[2], tier.per_day); }
#[test] fn limits_for_scales_every_window_by_pool_size() { let tier = RateTier::default_tier(); let one = limits_for(0, &tier, 1); let eight = limits_for(0, &tier, 8); assert_eq!(eight[0], one[0] * 8); assert_eq!(eight[1], one[1] * 8); assert_eq!(eight[2], one[2] * 8); }
#[test] fn limits_for_account_floor_then_scales() { // per_second = max(base, accounts * mul), then multiplied by the pool. let tier = RateTier::default_tier(); let accounts = ((tier.per_second_base as f64 / tier.per_second_account_mul) as u64) + 100; let expected_per_sec = (accounts as f64 * tier.per_second_account_mul) as u64; let scaled = limits_for(accounts, &tier, 4); assert_eq!(scaled[0], expected_per_sec * 4); }
#[test] fn limits_for_pool_size_zero_treated_as_one() { let tier = RateTier::default_tier(); assert_eq!(limits_for(0, &tier, 0), limits_for(0, &tier, 1)); }}