use crate::error::FlowGuardError;
#[cfg(feature = "quota-control")]
use crate::QuotaConfig;
use dashmap::DashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::time::{Duration, Instant};
#[derive(Debug, Clone)]
struct QuotaRecord {
usage: u64,
window_start: Instant,
}
pub struct QuotaLimiter {
config: QuotaConfig,
usage: Arc<DashMap<String, QuotaRecord>>,
}
impl QuotaLimiter {
pub fn new(config: QuotaConfig) -> Self {
Self {
config,
usage: Arc::new(DashMap::new()),
}
}
async fn check_and_consume(&self, key: &str) -> Result<bool, FlowGuardError> {
let now = Instant::now();
let window_duration = Duration::from_secs(self.config.window_size);
let mut record = self
.usage
.entry(key.to_string())
.or_insert_with(|| QuotaRecord {
usage: 0,
window_start: now,
});
if now.duration_since(record.window_start) >= window_duration {
record.usage = 0;
record.window_start = now;
}
let max_usage = if self.config.allow_overdraft {
let overdraft_limit =
self.config.limit * self.config.overdraft_limit_percent as u64 / 100;
self.config.limit + overdraft_limit
} else {
self.config.limit
};
if record.usage >= max_usage {
return Err(FlowGuardError::QuotaExceeded(format!(
"Quota exceeded for key '{}': used {}/{}",
key, record.usage, max_usage
)));
}
record.usage += 1;
Ok(true)
}
}
impl crate::limiters::Limiter for QuotaLimiter {
fn allow(
&self,
_cost: u64,
) -> Pin<Box<dyn Future<Output = Result<bool, FlowGuardError>> + Send + '_>> {
Box::pin(async move {
Ok(true)
})
}
fn check(
&self,
key: &str,
) -> Pin<Box<dyn Future<Output = Result<(), FlowGuardError>> + Send + '_>> {
let key = key.to_string();
Box::pin(async move {
self.check_and_consume(&key).await?;
Ok(())
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::limiters::Limiter;
use crate::QuotaType;
fn create_test_config() -> QuotaConfig {
QuotaConfig {
quota_type: QuotaType::Count,
limit: 10,
window_size: 60,
allow_overdraft: false,
overdraft_limit_percent: 0,
alert_config: Default::default(),
}
}
#[tokio::test]
async fn test_quota_limiter_allows_within_limit() {
let config = create_test_config();
let limiter = QuotaLimiter::new(config);
for i in 0..10 {
let result = limiter.check("user1").await;
assert!(result.is_ok(), "Request {} should be allowed", i);
}
}
#[tokio::test]
async fn test_quota_limiter_rejects_over_limit() {
let config = create_test_config();
let limiter = QuotaLimiter::new(config);
for _ in 0..10 {
let _ = limiter.check("user1").await;
}
let result = limiter.check("user1").await;
assert!(result.is_err());
assert!(matches!(result, Err(FlowGuardError::QuotaExceeded(_))));
}
#[tokio::test]
async fn test_quota_limiter_independent_keys() {
let config = create_test_config();
let limiter = QuotaLimiter::new(config);
for _ in 0..10 {
let _ = limiter.check("user1").await;
}
let result = limiter.check("user2").await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_quota_limiter_with_overdraft() {
let mut config = create_test_config();
config.allow_overdraft = true;
config.overdraft_limit_percent = 20;
let limiter = QuotaLimiter::new(config);
for i in 0..12 {
let result = limiter.check("user1").await;
assert!(result.is_ok(), "Request {} should be allowed", i);
}
let result = limiter.check("user1").await;
assert!(result.is_err());
}
}