#[cfg(feature = "rate-limiting")]
use async_trait::async_trait;
#[cfg(feature = "rate-limiting")]
pub use governor::{
clock::DefaultClock,
state::{keyed::DefaultKeyedStateStore, InMemoryState},
DefaultDirectRateLimiter, DefaultKeyedRateLimiter, Quota, RateLimiter,
};
#[cfg(feature = "rate-limiting")]
#[async_trait]
pub trait CacheAwareRateLimiter: Send + Sync + 'static {
async fn until_key_ready(&self, key: &str);
fn check_key(&self, key: &str) -> bool;
}
#[cfg(feature = "rate-limiting")]
#[derive(Debug)]
pub struct DomainRateLimiter {
limiter: DefaultKeyedRateLimiter<String>,
}
#[cfg(feature = "rate-limiting")]
impl DomainRateLimiter {
pub fn new(quota: Quota) -> Self {
Self { limiter: DefaultKeyedRateLimiter::keyed(quota) }
}
}
#[cfg(feature = "rate-limiting")]
#[async_trait]
impl CacheAwareRateLimiter for DomainRateLimiter {
async fn until_key_ready(&self, key: &str) {
self.limiter.until_key_ready(&key.to_string()).await;
}
fn check_key(&self, key: &str) -> bool {
self.limiter.check_key(&key.to_string()).is_ok()
}
}
#[cfg(feature = "rate-limiting")]
#[derive(Debug)]
pub struct DirectRateLimiter {
limiter: DefaultDirectRateLimiter,
}
#[cfg(feature = "rate-limiting")]
impl DirectRateLimiter {
pub fn direct(quota: Quota) -> DirectRateLimiter {
DirectRateLimiter { limiter: DefaultDirectRateLimiter::direct(quota) }
}
}
#[cfg(feature = "rate-limiting")]
#[async_trait]
impl CacheAwareRateLimiter for DirectRateLimiter {
async fn until_key_ready(&self, _key: &str) {
self.limiter.until_ready().await;
}
fn check_key(&self, _key: &str) -> bool {
self.limiter.check().is_ok()
}
}