use crate::application::config::RateLimiterConfig;
use crate::constants::{
DEFAULT_RATE_LIMIT_BURST_SIZE, FALLBACK_RATE_LIMIT_MAX_REQUESTS,
HISTORICAL_RATE_LIMIT_PER_SECOND, TRADING_HISTORICAL_BURST_SIZE, TRADING_RATE_LIMIT_PER_SECOND,
};
use governor::{
Quota, RateLimiter as GovernorRateLimiter,
clock::QuantaClock,
state::{InMemoryState, NotKeyed},
};
use std::num::NonZeroU32;
use std::sync::Arc;
use std::time::Duration;
type DirectLimiter = GovernorRateLimiter<NotKeyed, InMemoryState, QuantaClock>;
const NANOS_PER_SECOND: u64 = 1_000_000_000;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[repr(u8)]
pub enum RateLimitClass {
Trading,
Historical,
NonTrading,
}
#[derive(Clone)]
pub struct RateLimiter {
non_trading: Arc<DirectLimiter>,
trading: Arc<DirectLimiter>,
historical: Arc<DirectLimiter>,
}
#[must_use]
#[inline]
fn replenish_period(config: &RateLimiterConfig) -> Duration {
let max_requests = config.max_requests.max(FALLBACK_RATE_LIMIT_MAX_REQUESTS);
let period = if config.period_seconds == 0 {
Duration::from_secs(1)
} else {
Duration::from_secs(config.period_seconds)
};
let period_nanos = u64::try_from(period.as_nanos()).unwrap_or(u64::MAX);
let per_cell_nanos = period_nanos.div_ceil(u64::from(max_requests)).max(1);
Duration::from_nanos(per_cell_nanos)
}
#[must_use]
#[inline]
fn per_second_period(requests_per_second: u32) -> Duration {
let rps = u64::from(requests_per_second.max(1));
Duration::from_nanos((NANOS_PER_SECOND / rps).max(1))
}
#[must_use]
fn build_limiter(period_per_cell: Duration, burst_size: u32) -> Arc<DirectLimiter> {
let burst = NonZeroU32::new(burst_size)
.or_else(|| NonZeroU32::new(DEFAULT_RATE_LIMIT_BURST_SIZE))
.unwrap_or(NonZeroU32::MIN);
let quota = match Quota::with_period(period_per_cell) {
Some(quota) => quota.allow_burst(burst),
None => Quota::per_second(NonZeroU32::MIN).allow_burst(burst),
};
Arc::new(GovernorRateLimiter::direct(quota))
}
impl RateLimiter {
#[must_use]
pub fn new(config: &RateLimiterConfig) -> Self {
let non_trading = build_limiter(replenish_period(config), config.burst_size);
let trading = build_limiter(
per_second_period(TRADING_RATE_LIMIT_PER_SECOND),
TRADING_HISTORICAL_BURST_SIZE,
);
let historical = build_limiter(
per_second_period(HISTORICAL_RATE_LIMIT_PER_SECOND),
TRADING_HISTORICAL_BURST_SIZE,
);
Self {
non_trading,
trading,
historical,
}
}
#[inline]
fn limiter_for(&self, class: RateLimitClass) -> &DirectLimiter {
match class {
RateLimitClass::Trading => &self.trading,
RateLimitClass::Historical => &self.historical,
RateLimitClass::NonTrading => &self.non_trading,
}
}
pub async fn wait_for(&self, class: RateLimitClass) {
self.limiter_for(class).until_ready().await;
}
#[must_use]
pub fn check_for(&self, class: RateLimitClass) -> bool {
self.limiter_for(class).check().is_ok()
}
pub async fn wait(&self) {
self.wait_for(RateLimitClass::NonTrading).await;
}
#[must_use]
pub fn check(&self) -> bool {
self.check_for(RateLimitClass::NonTrading)
}
}
impl std::fmt::Debug for RateLimiter {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RateLimiter")
.field("non_trading", &"GovernorRateLimiter")
.field("trading", &"GovernorRateLimiter")
.field("historical", &"GovernorRateLimiter")
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_replenish_period_honors_max_requests() {
let config = RateLimiterConfig {
max_requests: 60,
period_seconds: 60,
burst_size: 3,
};
assert_eq!(replenish_period(&config), Duration::from_secs(1));
assert_ne!(replenish_period(&config), Duration::from_secs(60));
let config = RateLimiterConfig {
max_requests: 4,
period_seconds: 12,
burst_size: 3,
};
assert_eq!(replenish_period(&config), Duration::from_secs(3));
}
#[test]
fn test_replenish_period_zero_max_requests_falls_back() {
let config = RateLimiterConfig {
max_requests: 0,
period_seconds: 10,
burst_size: 1,
};
assert_eq!(replenish_period(&config), Duration::from_secs(10));
}
#[test]
fn test_replenish_period_zero_period_falls_back() {
let config = RateLimiterConfig {
max_requests: 5,
period_seconds: 0,
burst_size: 1,
};
assert_eq!(replenish_period(&config), Duration::from_millis(200));
}
#[test]
fn test_per_second_period_matches_rate() {
assert_eq!(per_second_period(1), Duration::from_secs(1));
assert_eq!(per_second_period(4), Duration::from_millis(250));
assert_eq!(per_second_period(0), Duration::from_secs(1));
}
#[test]
fn test_configured_quota_replenishes_at_max_requests_rate() {
use governor::clock::FakeRelativeClock;
let config = RateLimiterConfig {
max_requests: 20,
period_seconds: 1,
burst_size: 1,
};
let interval = replenish_period(&config);
assert_eq!(interval, Duration::from_millis(50));
let burst = NonZeroU32::new(config.burst_size).expect("burst is non-zero");
let quota = Quota::with_period(interval)
.expect("non-zero interval")
.allow_burst(burst);
let clock = FakeRelativeClock::default();
let limiter = GovernorRateLimiter::direct_with_clock(quota, clock.clone());
assert!(limiter.check().is_ok());
assert!(limiter.check().is_err(), "burst cell must be exhausted");
clock.advance(interval / 2);
assert!(
limiter.check().is_err(),
"must not replenish before the configured interval"
);
clock.advance(interval / 2);
assert!(
limiter.check().is_ok(),
"one cell must replenish after the configured 50ms interval"
);
}
#[tokio::test]
async fn test_trading_not_blocked_behind_saturated_non_trading() {
let config = RateLimiterConfig {
max_requests: 1,
period_seconds: 60,
burst_size: 1,
};
let limiter = RateLimiter::new(&config);
assert!(limiter.check_for(RateLimitClass::NonTrading));
assert!(
!limiter.check_for(RateLimitClass::NonTrading),
"non-trading bucket should be saturated"
);
assert!(
limiter.check_for(RateLimitClass::Trading),
"trading must not be blocked behind a saturated non-trading bucket"
);
assert!(
limiter.check_for(RateLimitClass::Historical),
"historical must not be blocked behind a saturated non-trading bucket"
);
}
#[tokio::test]
async fn test_wait_for_returns_without_parking_when_slot_available() {
let config = RateLimiterConfig {
max_requests: 10,
period_seconds: 1,
burst_size: 5,
};
let limiter = RateLimiter::new(&config);
assert!(
limiter.check_for(RateLimitClass::NonTrading),
"a burst slot must be available on a fresh bucket"
);
limiter.wait_for(RateLimitClass::NonTrading).await;
}
#[test]
fn test_check_delegates_to_non_trading() {
let config = RateLimiterConfig {
max_requests: 1,
period_seconds: 60,
burst_size: 1,
};
let limiter = RateLimiter::new(&config);
assert!(limiter.check());
assert!(
!limiter.check(),
"check() must delegate to the non-trading bucket"
);
assert!(
!limiter.check_for(RateLimitClass::NonTrading),
"check_for(NonTrading) must observe the same saturated bucket as check()"
);
}
}