use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::{Arc, OnceLock};
use std::time::Duration;
use crate::hold::Hold;
use crate::hooks::{AttemptInfo, Charge, Refused, RequestMeta, ResponseMeta, Throttle};
use governor::clock::Clock;
use governor::Quota;
use http::Method;
use polyoxide_rt::{Instant, Runtime};
struct RuntimeClock {
runtime: Arc<dyn Runtime>,
origin: Instant,
}
impl RuntimeClock {
fn new(runtime: Arc<dyn Runtime>) -> Self {
let origin = runtime.now();
Self { runtime, origin }
}
}
impl Clock for RuntimeClock {
type Instant = Duration;
fn now(&self) -> Duration {
self.runtime.now().saturating_duration_since(self.origin)
}
}
type DirectLimiter = governor::RateLimiter<
governor::state::NotKeyed,
governor::state::InMemoryState,
RuntimeClock,
governor::middleware::NoOpMiddleware<Duration>,
>;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Matching {
Prefix,
Exact,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct RateSpec {
count: u32,
period: Duration,
}
struct Bucket {
id: BucketId,
spec: RateSpec,
quota: Quota,
limiter: OnceLock<DirectLimiter>,
}
impl Bucket {
fn new(id: BucketId, count: u32, period: Duration) -> Arc<Self> {
let quota = quota(count, period);
Arc::new(Self {
id,
spec: RateSpec { count, period },
quota,
limiter: OnceLock::new(),
})
}
fn effective(&self) -> EffectiveQuota {
EffectiveQuota {
bucket: self.id,
count: self.spec.count,
period: self.spec.period,
quota: self.quota,
}
}
async fn until_ready(&self, runtime: &Arc<dyn Runtime>) {
let limiter = self.limiter.get_or_init(|| {
DirectLimiter::direct_with_clock(self.quota, RuntimeClock::new(Arc::clone(runtime)))
});
loop {
let wait = match limiter.check() {
Ok(()) => return,
Err(not_until) => not_until.wait_time_from(limiter.clock().now()),
};
runtime.sleep(wait).await;
}
}
}
struct EndpointLimit {
path_prefix: &'static str,
method: Option<Method>,
match_mode: Matching,
buckets: Vec<Arc<Bucket>>,
}
impl EndpointLimit {
fn matches(&self, path: &str, method: Option<&Method>) -> bool {
let path_matches = match self.match_mode {
Matching::Exact => path == self.path_prefix,
Matching::Prefix => {
match path.strip_prefix(self.path_prefix) {
Some(rest) => rest.is_empty() || rest.starts_with('/') || rest.starts_with('?'),
None => false,
}
}
};
if !path_matches {
return false;
}
match &self.method {
Some(expected) => method == Some(expected),
None => true,
}
}
}
#[derive(Clone)]
pub struct RateLimiter {
inner: Arc<RateLimiterInner>,
}
impl std::fmt::Debug for RateLimiter {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RateLimiter")
.field("endpoints", &self.inner.limits.len())
.finish()
}
}
struct RateLimiterInner {
limits: Vec<EndpointLimit>,
general: Arc<Bucket>,
hold: Hold,
}
fn quota(count: u32, period: Duration) -> Quota {
Quota::with_period(WindowQuotaTable::paced_interval(count, period))
.expect("quota interval must be non-zero")
}
fn admitted_in_one_window(q: &Quota, period: Duration) -> u128 {
let refilled = period.as_nanos() / q.replenish_interval().as_nanos();
u128::from(q.burst_size().get()) + refilled
}
fn sustained_slots(count: u32) -> u32 {
let target = count.saturating_sub(count.div_ceil(RESERVED_FRACTION));
target.max(2) - 1
}
const RESERVED_FRACTION: u32 = 10;
#[cfg(test)]
mod quota_arithmetic {
use super::*;
const PUBLISHED_SHAPES: &[(u32, u64)] = &[
(9_000, 10), (4_000, 10), (1_000, 10), (25, 60), (150, 10), (200, 10), (300, 10), (350, 10), (500, 10), (100, 10), (50, 10), (5_000, 10), (120_000, 600), ];
#[test]
fn no_quota_admits_more_than_its_published_count_in_one_window() {
for &(count, secs) in PUBLISHED_SHAPES {
let period = Duration::from_secs(secs);
let q = quota(count, period);
let admitted = admitted_in_one_window(&q, period);
assert!(
admitted <= u128::from(count),
"{count}/{secs}s admits {admitted} in one window \
({} burst + {} refilled) — the published quota is spent twice",
q.burst_size(),
admitted - u128::from(q.burst_size().get()),
);
}
}
#[test]
fn every_quota_reserves_headroom_below_the_published_count() {
for &(count, secs) in PUBLISHED_SHAPES {
let period = Duration::from_secs(secs);
let admitted = admitted_in_one_window("a(count, period), period);
let ceiling = u128::from(count - count.div_ceil(RESERVED_FRACTION));
assert!(
admitted <= ceiling,
"{count}/{secs}s admits {admitted} in one window, above the {ceiling} \
the reserve allows — no headroom under the published cap"
);
}
}
#[test]
fn a_published_150_per_10s_admits_135_per_window() {
let period = Duration::from_secs(10);
assert_eq!(admitted_in_one_window("a(150, period), period), 135);
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct BucketId {
table: u32,
index: u32,
}
static TABLES: AtomicU32 = AtomicU32::new(0);
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct EffectiveQuota {
pub bucket: BucketId,
pub count: u32,
pub period: Duration,
quota: Quota,
}
impl EffectiveQuota {
pub fn interval(&self) -> Duration {
self.quota.replenish_interval()
}
pub fn admitted_in_one_window(&self) -> u128 {
admitted_in_one_window(&self.quota, self.period)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct QuotaRow<'a> {
pub pattern: &'static str,
pub method: Option<&'a Method>,
pub matching: Matching,
pub buckets: Vec<EffectiveQuota>,
}
pub struct WindowQuotaTable {
id: u32,
buckets: Vec<Arc<Bucket>>,
rows: Vec<EndpointLimit>,
hold: Hold,
hold_shared: bool,
}
impl WindowQuotaTable {
pub fn new(count: u32, period: Duration) -> Self {
let mut table = Self {
id: TABLES.fetch_add(1, Ordering::Relaxed),
buckets: Vec::new(),
rows: Vec::new(),
hold: Hold::unbounded(),
hold_shared: false,
};
table.bucket(count, period);
table
}
pub fn bucket(&mut self, count: u32, period: Duration) -> BucketId {
let id = BucketId {
table: self.id,
index: u32::try_from(self.buckets.len()).expect("fewer than 2^32 buckets"),
};
self.buckets.push(Bucket::new(id, count, period));
id
}
pub fn row(
&mut self,
matching: Matching,
pattern: &'static str,
method: Option<Method>,
buckets: &[BucketId],
) -> &mut Self {
let buckets = buckets
.iter()
.map(|id| {
assert_eq!(id.table, self.id, "{id:?} was made by another table");
self.buckets[id.index as usize].clone()
})
.collect();
self.rows.push(EndpointLimit {
path_prefix: pattern,
method,
match_mode: matching,
buckets,
});
self
}
pub fn prefix(
&mut self,
pattern: &'static str,
method: Option<Method>,
count: u32,
period: Duration,
) -> &mut Self {
let bucket = self.bucket(count, period);
self.row(Matching::Prefix, pattern, method, &[bucket])
}
pub fn with_hold(&mut self, hold: Hold) -> &mut Self {
self.hold = hold;
self.hold_shared = true;
self
}
pub fn with_runtime(&mut self, runtime: Arc<dyn Runtime>) -> &mut Self {
debug_assert!(
!self.hold_shared,
"WindowQuotaTable::with_runtime called after with_hold: it would replace the shared hold; call with_runtime first"
);
self.hold = Hold::unbounded_on(runtime);
self
}
pub fn build(self) -> RateLimiter {
RateLimiter {
inner: Arc::new(RateLimiterInner {
general: self.buckets[0].clone(),
limits: self.rows,
hold: self.hold,
}),
}
}
pub fn paced_interval(count: u32, period: Duration) -> Duration {
period / sustained_slots(count)
}
}
impl RateLimiter {
pub fn begin_cooldown(&self, delay: Duration) {
self.inner.hold.extend(delay);
}
pub fn hold(&self) -> &Hold {
&self.inner.hold
}
pub async fn acquire(&self, path: &str, method: Option<&Method>) {
let runtime = self.inner.hold.runtime();
self.inner.hold.wait().await;
self.inner.general.until_ready(runtime).await;
if let Some(limit) = self.row_for(path, method) {
for bucket in &limit.buckets {
bucket.until_ready(runtime).await;
}
}
self.inner.hold.wait().await;
}
pub fn effective_quota(&self, method: &Method, path: &str) -> Vec<EffectiveQuota> {
std::iter::once(&self.inner.general)
.chain(
self.row_for(path, Some(method))
.into_iter()
.flat_map(|limit| &limit.buckets),
)
.map(|bucket| bucket.effective())
.collect()
}
pub fn rows(&self) -> Vec<QuotaRow<'_>> {
self.inner
.limits
.iter()
.map(|limit| QuotaRow {
pattern: limit.path_prefix,
method: limit.method.as_ref(),
matching: limit.match_mode,
buckets: limit.buckets.iter().map(|b| b.effective()).collect(),
})
.collect()
}
fn row_for(&self, path: &str, method: Option<&Method>) -> Option<&EndpointLimit> {
self.inner.limits.iter().find(|l| l.matches(path, method))
}
}
impl Throttle for RateLimiter {
async fn acquire(&self, meta: &RequestMeta<'_>) -> Result<Charge, Refused> {
RateLimiter::acquire(self, meta.path, Some(meta.method)).await;
Ok(Charge::none())
}
fn observe(&self, _charge: &Charge, _response: &ResponseMeta<'_>, _attempt: &AttemptInfo) {}
fn hold(&self, delay: Duration) {
self.inner.hold.extend(delay);
}
}
#[cfg(test)]
mod window_table {
use super::*;
const TEN_SECONDS: Duration = Duration::from_secs(10);
const TEN_MINUTES: Duration = Duration::from_secs(600);
fn shape(quotas: &[EffectiveQuota]) -> Vec<(u32, Duration)> {
quotas.iter().map(|q| (q.count, q.period)).collect()
}
#[test]
fn a_request_awaits_the_general_bucket_then_its_rows_buckets_in_order() {
let mut table = WindowQuotaTable::new(1_000, TEN_SECONDS);
let burst = table.bucket(500, TEN_SECONDS);
let sustained = table.bucket(6_000, TEN_MINUTES);
table.row(
Matching::Prefix,
"/order",
Some(Method::POST),
&[burst, sustained],
);
let limiter = table.build();
let quotas = limiter.effective_quota(&Method::POST, "/order");
assert_eq!(
shape("as),
[
(1_000, TEN_SECONDS),
(500, TEN_SECONDS),
(6_000, TEN_MINUTES)
]
);
assert_eq!(quotas[1].bucket, burst);
assert_eq!(quotas[2].bucket, sustained);
assert_eq!(
shape(&limiter.effective_quota(&Method::GET, "/elsewhere")),
[(1_000, TEN_SECONDS)],
"an unmatched request is held to the general bucket alone"
);
}
#[tokio::test]
async fn two_rows_naming_one_bucket_share_one_allowance() {
let mut table = WindowQuotaTable::new(100_000, TEN_SECONDS);
let shared = table.bucket(100, TEN_SECONDS);
table
.row(Matching::Prefix, "/a", None, &[shared])
.row(Matching::Prefix, "/b", None, &[shared])
.prefix("/c", None, 100, TEN_SECONDS);
let limiter = table.build();
let a = limiter.effective_quota(&Method::GET, "/a")[1].bucket;
let b = limiter.effective_quota(&Method::GET, "/b")[1].bucket;
let c = limiter.effective_quota(&Method::GET, "/c")[1].bucket;
assert_eq!(a, b);
assert_ne!(a, c, "a row's own bucket is shared with nobody");
limiter.acquire("/a", None).await;
let start = std::time::Instant::now();
limiter.acquire("/b", None).await;
assert!(
start.elapsed() >= Duration::from_millis(80),
"/b went out after {:?}, so /a and /b are not one allowance",
start.elapsed()
);
}
#[test]
fn the_first_matching_row_wins() {
let mut table = WindowQuotaTable::new(1_000, TEN_SECONDS);
table
.prefix("/data/orders", None, 50, TEN_SECONDS)
.prefix("/data", None, 500, TEN_SECONDS)
.prefix("/data/trades", None, 60, TEN_SECONDS);
let limiter = table.build();
let count = |path| limiter.effective_quota(&Method::GET, path)[1].count;
assert_eq!(count("/data/orders"), 50);
assert_eq!(count("/data"), 500);
assert_eq!(
count("/data/trades"),
500,
"a row after a shorter pattern that matches it is never reached"
);
}
#[test]
fn a_prefix_row_matches_only_at_a_segment_boundary() {
let mut table = WindowQuotaTable::new(1_000, TEN_SECONDS);
table.prefix("/price", None, 100, TEN_SECONDS);
let limiter = table.build();
for path in ["/price", "/price/1", "/price?token=1"] {
assert_eq!(
limiter.effective_quota(&Method::GET, path).len(),
2,
"{path}"
);
}
for path in ["/prices-history", "/pricing", "/midpoint"] {
assert_eq!(
limiter.effective_quota(&Method::GET, path).len(),
1,
"{path}"
);
}
}
#[test]
fn an_exact_row_matches_only_its_own_path() {
let mut table = WindowQuotaTable::new(1_000, TEN_SECONDS);
let bucket = table.bucket(100, TEN_SECONDS);
table.row(Matching::Exact, "/trades", None, &[bucket]);
let limiter = table.build();
assert_eq!(limiter.effective_quota(&Method::GET, "/trades").len(), 2);
for path in ["/trades/1", "/trades?limit=10", "/traded"] {
assert_eq!(
limiter.effective_quota(&Method::GET, path).len(),
1,
"{path}"
);
}
}
#[test]
fn a_method_scoped_row_ignores_other_methods() {
let mut table = WindowQuotaTable::new(1_000, TEN_SECONDS);
table
.prefix("/order", Some(Method::POST), 500, TEN_SECONDS)
.prefix("/order", None, 90, TEN_SECONDS);
let limiter = table.build();
let count = |method| limiter.effective_quota(&method, "/order")[1].count;
assert_eq!(count(Method::POST), 500);
assert_eq!(count(Method::GET), 90);
assert_eq!(count(Method::DELETE), 90);
}
#[test]
fn rows_reports_every_row_in_match_order() {
let mut table = WindowQuotaTable::new(1_000, TEN_SECONDS);
let shared = table.bucket(900, TEN_SECONDS);
table
.prefix("/balance-allowance/update", None, 50, TEN_SECONDS)
.row(Matching::Exact, "/trades", Some(Method::GET), &[shared]);
let limiter = table.build();
let rows = limiter.rows();
assert_eq!(rows.len(), 2);
assert_eq!(rows[0].pattern, "/balance-allowance/update");
assert_eq!(rows[0].method, None);
assert_eq!(rows[0].matching, Matching::Prefix);
assert_eq!(shape(&rows[0].buckets), [(50, TEN_SECONDS)]);
assert_eq!(rows[1].pattern, "/trades");
assert_eq!(rows[1].method, Some(&Method::GET));
assert_eq!(rows[1].matching, Matching::Exact);
assert_eq!(rows[1].buckets[0].bucket, shared);
assert_eq!(format!("{limiter:?}"), "RateLimiter { endpoints: 2 }");
}
#[test]
fn every_bucket_paces_at_its_paced_interval_and_keeps_a_tenth_back() {
for (count, period) in [
(25, TEN_SECONDS * 6),
(150, TEN_SECONDS),
(9_000, TEN_SECONDS),
] {
let mut table = WindowQuotaTable::new(count, period);
table.prefix("/row", None, count, period);
let limiter = table.build();
for quota in limiter.effective_quota(&Method::GET, "/row") {
assert_eq!(
quota.interval(),
WindowQuotaTable::paced_interval(count, period)
);
assert!(
quota.admitted_in_one_window() <= u128::from(count - count.div_ceil(10)),
"{count}/{period:?} admits {} in one window",
quota.admitted_in_one_window()
);
}
}
assert_eq!(
WindowQuotaTable::paced_interval(150, TEN_SECONDS),
TEN_SECONDS / 134
);
}
#[tokio::test]
async fn a_hold_set_during_a_bucket_wait_is_honoured() {
let mut table = WindowQuotaTable::new(100_000, TEN_SECONDS);
table.prefix("/slow", None, 2, Duration::from_millis(300));
let limiter = table.build();
limiter.acquire("/slow", None).await;
let start = std::time::Instant::now();
let holder = {
let limiter = limiter.clone();
tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(50)).await;
limiter.begin_cooldown(Duration::from_millis(600));
})
};
limiter.acquire("/slow", None).await;
holder.await.unwrap();
assert!(
start.elapsed() >= Duration::from_millis(600),
"the request went out after {:?}, inside a hold set while it waited for its \
token",
start.elapsed()
);
}
#[test]
#[should_panic(expected = "was made by another table")]
fn a_bucket_from_another_table_is_refused() {
let mut other = WindowQuotaTable::new(1_000, TEN_SECONDS);
let foreign = other.bucket(100, TEN_SECONDS);
let mut table = WindowQuotaTable::new(1_000, TEN_SECONDS);
table.bucket(100, TEN_SECONDS);
table.row(Matching::Prefix, "/a", None, &[foreign]);
}
}
#[derive(Debug, Clone)]
pub struct RetryConfig {
pub max_retries: u32,
pub initial_backoff_ms: u64,
pub max_backoff_ms: u64,
}
impl Default for RetryConfig {
fn default() -> Self {
Self {
max_retries: 3,
initial_backoff_ms: 500,
max_backoff_ms: 10_000,
}
}
}
impl RetryConfig {
pub fn backoff(&self, attempt: u32) -> Duration {
let base = self
.initial_backoff_ms
.saturating_mul(1u64 << attempt.min(10));
let capped = base.min(self.max_backoff_ms);
let jitter_factor = 0.75 + (fastrand::f64() * 0.5);
let ms = (capped as f64 * jitter_factor) as u64;
Duration::from_millis(ms.max(1))
}
pub fn retry_delay(&self, attempt: u32, retry_after: Option<&str>) -> Duration {
let clamp = Duration::from_millis(self.max_backoff_ms);
let requested = retry_after.and_then(|v| polyoxide_venue::parse_retry_after(v, clamp));
polyoxide_venue::retry_delay(requested, self.backoff(attempt))
}
pub fn attempt_info(&self, attempt: u32) -> AttemptInfo {
AttemptInfo {
attempt,
retries_left: self.max_retries.saturating_sub(attempt),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_retry_config_default() {
let cfg = RetryConfig::default();
assert_eq!(cfg.max_retries, 3);
assert_eq!(cfg.initial_backoff_ms, 500);
assert_eq!(cfg.max_backoff_ms, 10_000);
}
#[test]
fn test_backoff_attempt_zero() {
let cfg = RetryConfig::default();
let d = cfg.backoff(0);
let ms = d.as_millis() as u64;
assert!(
(375..=625).contains(&ms),
"attempt 0: {ms}ms not in [375, 625]"
);
}
#[test]
fn test_backoff_exponential_growth() {
let cfg = RetryConfig::default();
let d0 = cfg.backoff(0);
let d1 = cfg.backoff(1);
let d2 = cfg.backoff(2);
assert!(d0 < d1, "d0={d0:?} should be < d1={d1:?}");
assert!(d1 < d2, "d1={d1:?} should be < d2={d2:?}");
}
#[test]
fn test_backoff_jitter_bounds() {
let cfg = RetryConfig::default();
for attempt in 0..20 {
let d = cfg.backoff(attempt);
let base = cfg
.initial_backoff_ms
.saturating_mul(1u64 << attempt.min(10));
let capped = base.min(cfg.max_backoff_ms);
let lower = (capped as f64 * 0.75) as u64;
let upper = (capped as f64 * 1.25) as u64;
let ms = d.as_millis() as u64;
assert!(
ms >= lower.max(1) && ms <= upper,
"attempt {attempt}: {ms}ms not in [{lower}, {upper}]"
);
}
}
#[test]
fn test_backoff_max_capping() {
let cfg = RetryConfig::default();
for attempt in 5..=10 {
let d = cfg.backoff(attempt);
let ceiling = (cfg.max_backoff_ms as f64 * 1.25) as u64;
assert!(
d.as_millis() as u64 <= ceiling,
"attempt {attempt}: {:?} exceeded ceiling {ceiling}ms",
d
);
}
}
#[test]
fn test_backoff_very_high_attempt() {
let cfg = RetryConfig::default();
let d = cfg.backoff(100);
let ceiling = (cfg.max_backoff_ms as f64 * 1.25) as u64;
assert!(d.as_millis() as u64 <= ceiling);
assert!(d.as_millis() >= 1);
}
#[test]
fn test_backoff_jitter_distribution() {
let cfg = RetryConfig::default();
let midpoint = cfg.initial_backoff_ms; let (mut below, mut above) = (0u32, 0u32);
for _ in 0..200 {
let ms = cfg.backoff(0).as_millis() as u64;
if ms < midpoint {
below += 1;
} else {
above += 1;
}
}
assert!(
below >= 20 && above >= 20,
"jitter looks degenerate: {below} below midpoint, {above} above"
);
}
#[test]
fn the_one_retry_after_parser_settles_the_old_disagreements() {
let cfg = RetryConfig::default();
assert_eq!(cfg.retry_delay(0, Some(" 2 ")), Duration::from_secs(2));
assert_eq!(
cfg.retry_delay(0, Some("604800")),
Duration::from_millis(cfg.max_backoff_ms)
);
for junk in [
"Wed, 21 Oct 2026 07:28:00 GMT",
"abc",
"NaN",
"inf",
"-1",
"",
] {
let ms = cfg.retry_delay(0, Some(junk)).as_millis();
assert!((375..=625).contains(&ms), "{junk:?} waited {ms}ms");
}
}
#[test]
fn test_quota_creation() {
let _ = quota(100, Duration::from_secs(10));
let _ = quota(1, Duration::from_secs(60));
let _ = quota(9_000, Duration::from_secs(10));
}
#[test]
fn test_quota_edge_zero_count() {
let _ = quota(0, Duration::from_secs(10));
let _ = quota(1, Duration::from_secs(10));
}
#[test]
fn test_match_mode_prefix_segment_boundary() {
let pattern = "/price";
let check = |path: &str| -> bool {
match path.strip_prefix(pattern) {
Some(rest) => rest.is_empty() || rest.starts_with('/') || rest.starts_with('?'),
None => false,
}
};
assert!(check("/price"), "exact match");
assert!(check("/price/foo"), "sub-path");
assert!(check("/price?token=abc"), "query params");
assert!(!check("/prices-history"), "partial word /prices-history");
assert!(!check("/pricelist"), "partial word /pricelist");
assert!(!check("/pricing"), "partial word /pricing");
assert!(!check("/midpoint"), "different prefix");
}
#[test]
fn test_match_mode_exact() {
let pattern = "/trades";
let check = |path: &str| -> bool { path == pattern };
assert!(check("/trades"), "exact match");
assert!(!check("/trades/123"), "sub-path should not match");
assert!(!check("/trades?limit=10"), "query params should not match");
assert!(!check("/traded"), "different word should not match");
}
#[test]
#[cfg(debug_assertions)]
#[should_panic(expected = "with_runtime called after with_hold")]
fn with_runtime_after_with_hold_is_refused() {
let mut table = WindowQuotaTable::new(10, Duration::from_secs(1));
table.with_hold(Hold::unbounded());
table.with_runtime(polyoxide_rt::default_runtime());
}
}