pub const DEFAULT_QUOTA_LIMIT: u64 = 1000;
pub const DEFAULT_WINDOW_SIZE_SECS: u64 = 3600;
pub const DEFAULT_DEDUP_WINDOW_SECS: u64 = 300;
pub const DEFAULT_OVERDRAFT_LIMIT_PERCENT: u8 = 20;
use crate::error::{ConsumeResult, FlowGuardError};
use crate::storage::QuotaStorage;
use chrono::{DateTime, Duration, Utc};
use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::time::Duration as StdDuration;
#[cfg(feature = "quota-control")]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum QuotaType {
Token,
Money,
Count,
}
impl QuotaType {
pub fn parse(s: &str) -> Option<Self> {
match s.to_lowercase().as_str() {
"token" => Some(QuotaType::Token),
"money" => Some(QuotaType::Money),
"count" => Some(QuotaType::Count),
_ => None,
}
}
pub fn as_str(&self) -> &'static str {
match self {
QuotaType::Token => "token",
QuotaType::Money => "money",
QuotaType::Count => "count",
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg(feature = "quota-control")]
pub struct QuotaConfig {
pub quota_type: QuotaType,
pub limit: u64,
pub window_size: u64,
pub allow_overdraft: bool,
pub overdraft_limit_percent: u8,
pub alert_config: AlertConfig,
}
impl Default for QuotaConfig {
fn default() -> Self {
Self {
quota_type: QuotaType::Count,
limit: DEFAULT_QUOTA_LIMIT,
window_size: DEFAULT_WINDOW_SIZE_SECS,
allow_overdraft: false,
overdraft_limit_percent: DEFAULT_OVERDRAFT_LIMIT_PERCENT,
alert_config: AlertConfig::default(),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg(feature = "quota-control")]
pub struct AlertConfig {
pub enabled: bool,
pub thresholds: Vec<u8>,
pub channels: Vec<AlertChannel>,
pub dedup_window: u64,
}
impl Default for AlertConfig {
fn default() -> Self {
Self {
enabled: true,
thresholds: vec![80, 90, 100],
channels: vec![AlertChannel::Log],
dedup_window: DEFAULT_DEDUP_WINDOW_SECS,
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[cfg(feature = "quota-control")]
pub enum AlertChannel {
Log,
Webhook { url: String },
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg(feature = "quota-control")]
pub struct AlertInfo {
pub user_id: String,
pub resource: String,
pub quota_type: QuotaType,
pub threshold: u8,
pub current_usage: u64,
pub limit: u64,
pub triggered_at: DateTime<Utc>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg(feature = "quota-control")]
pub struct QuotaState {
pub consumed: u64,
pub window_start: DateTime<Utc>,
pub window_end: DateTime<Utc>,
}
#[cfg(feature = "quota-control")]
pub struct QuotaController<S: QuotaStorage> {
storage: Arc<S>,
config: QuotaConfig,
alert_dedup: Arc<DashMap<String, DateTime<Utc>>>,
}
impl<S: QuotaStorage + Clone + 'static> Clone for QuotaController<S> {
fn clone(&self) -> Self {
Self {
storage: self.storage.clone(),
config: self.config.clone(),
alert_dedup: self.alert_dedup.clone(),
}
}
}
impl<S: QuotaStorage + 'static> QuotaController<S> {
pub fn new(storage: S, config: QuotaConfig) -> Self {
Self {
storage: Arc::new(storage),
config,
alert_dedup: Arc::new(DashMap::new()),
}
}
pub async fn consume(
&self,
user_id: &str,
resource: &str,
cost: u64,
) -> Result<ConsumeResult, FlowGuardError> {
if cost == 0 {
return Ok(ConsumeResult {
allowed: true,
remaining: self.config.limit,
alert_triggered: false,
});
}
let quota_state = self.get_or_create_quota_state(user_id, resource).await?;
let updated_state = self.check_and_reset_window(quota_state).await?;
let overdraft_limit = if self.config.allow_overdraft {
self.config
.limit
.checked_mul(self.config.overdraft_limit_percent as u64)
.and_then(|v| v.checked_div(100))
.unwrap_or(u64::MAX / 2) } else {
0
};
let total_limit = self
.config
.limit
.checked_add(overdraft_limit)
.unwrap_or(u64::MAX / 2);
if updated_state.consumed + cost > total_limit {
return Ok(ConsumeResult {
allowed: false,
remaining: total_limit.saturating_sub(updated_state.consumed),
alert_triggered: false,
});
}
let new_consumed = updated_state.consumed + cost;
self.save_quota_state(user_id, resource, &updated_state, new_consumed)
.await?;
let remaining = total_limit.saturating_sub(new_consumed);
let alert_triggered = self
.check_and_trigger_alert(user_id, resource, new_consumed)
.await?;
Ok(ConsumeResult {
allowed: true,
remaining,
alert_triggered,
})
}
pub async fn get_quota(
&self,
user_id: &str,
resource: &str,
) -> Result<Option<QuotaState>, FlowGuardError> {
let quota_info = self
.storage
.get_quota(user_id, resource)
.await
.map_err(FlowGuardError::StorageError)?;
if let Some(info) = quota_info {
Ok(Some(QuotaState {
consumed: info.consumed,
window_start: info.window_start,
window_end: info.window_end,
}))
} else {
Ok(None)
}
}
pub async fn reset_quota(&self, user_id: &str, resource: &str) -> Result<(), FlowGuardError> {
self.storage
.reset(
user_id,
resource,
self.config.limit,
StdDuration::from_secs(self.config.window_size),
)
.await
.map_err(FlowGuardError::StorageError)?;
Ok(())
}
async fn get_or_create_quota_state(
&self,
user_id: &str,
resource: &str,
) -> Result<QuotaState, FlowGuardError> {
if let Some(state) = self.get_quota(user_id, resource).await? {
return Ok(state);
}
let now = Utc::now();
let window_start = now;
let window_end = now + Duration::seconds(self.config.window_size as i64);
Ok(QuotaState {
consumed: 0,
window_start,
window_end,
})
}
async fn check_and_reset_window(
&self,
state: QuotaState,
) -> Result<QuotaState, FlowGuardError> {
let now = Utc::now();
if now < state.window_end {
return Ok(state);
}
let window_duration = Duration::seconds(self.config.window_size as i64);
let elapsed = now.signed_duration_since(state.window_start);
let windows_passed = (elapsed.num_seconds() / window_duration.num_seconds()) as u64;
let safe_windows_passed = windows_passed.min(i32::MAX as u64) as i32;
let new_window_start = state.window_start + window_duration * safe_windows_passed;
let new_window_end = new_window_start + window_duration;
let window_elapsed = now.signed_duration_since(state.window_start);
let window_progress = (window_elapsed.num_milliseconds() as f64
/ window_duration.num_milliseconds() as f64)
.min(1.0);
let retained_consumed = if windows_passed >= 1 {
0
} else {
(state.consumed as f64 * (1.0 - window_progress)) as u64
};
Ok(QuotaState {
consumed: retained_consumed,
window_start: new_window_start,
window_end: new_window_end,
})
}
async fn save_quota_state(
&self,
user_id: &str,
resource: &str,
state: &QuotaState,
new_consumed: u64,
) -> Result<(), FlowGuardError> {
let overdraft_limit = if self.config.allow_overdraft {
self.config
.limit
.checked_mul(self.config.overdraft_limit_percent as u64)
.and_then(|v| v.checked_div(100))
.unwrap_or(u64::MAX / 2)
} else {
0
};
let total_limit = self
.config
.limit
.checked_add(overdraft_limit)
.unwrap_or(u64::MAX / 2);
let _result = self
.storage
.consume(
user_id,
resource,
new_consumed.saturating_sub(state.consumed),
total_limit,
StdDuration::from_secs(self.config.window_size),
)
.await
.map_err(FlowGuardError::StorageError)?;
Ok(())
}
async fn check_and_trigger_alert(
&self,
user_id: &str,
resource: &str,
consumed: u64,
) -> Result<bool, FlowGuardError> {
if !self.config.alert_config.enabled {
return Ok(false);
}
let usage_percent = if self.config.limit > 0 {
(consumed as f64 / self.config.limit as f64 * 100.0) as u8
} else {
100
};
let mut alert_triggered = false;
for &threshold in &self.config.alert_config.thresholds {
if usage_percent >= threshold {
let dedup_key = format!("{}:{}:{}", user_id, resource, threshold);
let should_alert = {
if let Some(last_alert_time) = self.alert_dedup.get(&dedup_key) {
let elapsed = Utc::now().signed_duration_since(*last_alert_time);
elapsed.num_seconds() as u64 >= self.config.alert_config.dedup_window
} else {
true
}
};
if should_alert {
let alert_info = AlertInfo {
user_id: user_id.to_string(),
resource: resource.to_string(),
quota_type: self.config.quota_type,
threshold,
current_usage: consumed,
limit: self.config.limit,
triggered_at: Utc::now(),
};
self.send_alert(alert_info).await;
self.alert_dedup.insert(dedup_key, Utc::now());
alert_triggered = true;
}
}
}
Ok(alert_triggered)
}
async fn send_alert(&self, alert_info: AlertInfo) {
for channel in &self.config.alert_config.channels {
let channel = channel.clone();
let alert_info = alert_info.clone();
tokio::spawn(async move {
match channel {
AlertChannel::Log => {
tracing::warn!(
user_id = %alert_info.user_id,
resource = %alert_info.resource,
quota_type = %alert_info.quota_type.as_str(),
threshold = alert_info.threshold,
current_usage = alert_info.current_usage,
limit = alert_info.limit,
triggered_at = %alert_info.triggered_at.format("%Y-%m-%d %H:%M:%S UTC"),
"配额告警触发"
);
}
AlertChannel::Webhook { url } => {
if let Err(e) = send_webhook_alert(&url, &alert_info).await {
tracing::error!(error = %e, "发送 Webhook 告警失败");
}
}
}
});
}
}
pub fn config(&self) -> &QuotaConfig {
&self.config
}
pub fn update_config(&mut self, config: QuotaConfig) {
self.config = config;
}
pub fn cleanup_alert_dedup(&self) {
let now = Utc::now();
let dedup_window = Duration::seconds(self.config.alert_config.dedup_window as i64);
self.alert_dedup.retain(|_, last_alert_time| {
now.signed_duration_since(*last_alert_time) < dedup_window
});
}
}
#[cfg(feature = "webhook")]
async fn send_webhook_alert(
url: &str,
alert_info: &AlertInfo,
) -> Result<(), Box<dyn std::error::Error>> {
let client = reqwest::Client::new();
let response = client
.post(url)
.json(alert_info)
.timeout(StdDuration::from_secs(5))
.send()
.await?;
if response.status().is_success() {
Ok(())
} else {
Err(format!("Webhook 返回错误状态码: {}", response.status()).into())
}
}
#[cfg(not(feature = "webhook"))]
async fn send_webhook_alert(
_url: &str,
_alert_info: &AlertInfo,
) -> Result<(), Box<dyn std::error::Error>> {
Err("Webhook 功能未启用,请启用 'webhook' feature".into())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::error::StorageError;
use crate::storage::{QuotaInfo, QuotaStorage};
use ahash::AHashMap as HashMap;
use async_trait::async_trait;
use std::sync::Mutex;
struct TestQuotaStorage {
quotas: Mutex<HashMap<String, QuotaInfo>>,
}
impl TestQuotaStorage {
fn new() -> Self {
Self {
quotas: Mutex::new(HashMap::new()),
}
}
}
#[async_trait]
impl QuotaStorage for TestQuotaStorage {
async fn get_quota(
&self,
user_id: &str,
resource: &str,
) -> Result<Option<QuotaInfo>, StorageError> {
let key = format!("{}:{}", user_id, resource);
Ok(self.quotas.lock().unwrap().get(&key).cloned())
}
async fn consume(
&self,
user_id: &str,
resource: &str,
cost: u64,
limit: u64,
window: StdDuration,
) -> Result<ConsumeResult, StorageError> {
let key = format!("{}:{}", user_id, resource);
let mut quotas = self.quotas.lock().unwrap();
let quota_info = quotas.entry(key.clone()).or_insert_with(|| {
let now = Utc::now();
QuotaInfo {
consumed: 0,
limit,
window_start: now,
window_end: now
+ Duration::from_std(window)
.unwrap_or(Duration::seconds(DEFAULT_WINDOW_SIZE_SECS as i64)),
}
});
let now = Utc::now();
if now >= quota_info.window_end {
quota_info.consumed = 0;
quota_info.window_start = now;
quota_info.window_end = now
+ Duration::from_std(window)
.unwrap_or(Duration::seconds(DEFAULT_WINDOW_SIZE_SECS as i64));
quota_info.limit = limit;
}
if quota_info.consumed + cost > quota_info.limit {
return Ok(ConsumeResult {
allowed: false,
remaining: quota_info.limit - quota_info.consumed,
alert_triggered: false,
});
}
quota_info.consumed += cost;
Ok(ConsumeResult {
allowed: true,
remaining: quota_info.limit - quota_info.consumed,
alert_triggered: false,
})
}
async fn reset(
&self,
user_id: &str,
resource: &str,
limit: u64,
window: StdDuration,
) -> Result<(), StorageError> {
let key = format!("{}:{}", user_id, resource);
let mut quotas = self.quotas.lock().unwrap();
if let Some(quota_info) = quotas.get_mut(&key) {
quota_info.consumed = 0;
quota_info.limit = limit;
let now = Utc::now();
quota_info.window_start = now;
quota_info.window_end = now
+ Duration::from_std(window)
.unwrap_or(Duration::seconds(DEFAULT_WINDOW_SIZE_SECS as i64));
}
Ok(())
}
}
#[test]
fn test_quota_type_parse() {
assert_eq!(QuotaType::parse("token"), Some(QuotaType::Token));
assert_eq!(QuotaType::parse("money"), Some(QuotaType::Money));
assert_eq!(QuotaType::parse("count"), Some(QuotaType::Count));
assert_eq!(QuotaType::parse("unknown"), None);
}
#[test]
fn test_quota_type_as_str() {
assert_eq!(QuotaType::Token.as_str(), "token");
assert_eq!(QuotaType::Money.as_str(), "money");
assert_eq!(QuotaType::Count.as_str(), "count");
}
#[test]
fn test_quota_config_default() {
let config = QuotaConfig::default();
assert_eq!(config.quota_type, QuotaType::Count);
assert_eq!(config.limit, 1000);
assert_eq!(config.window_size, 3600);
assert!(!config.allow_overdraft);
assert_eq!(config.overdraft_limit_percent, 20);
assert!(config.alert_config.enabled);
}
#[test]
fn test_alert_config_default() {
let config = AlertConfig::default();
assert!(config.enabled);
assert_eq!(config.thresholds, vec![80, 90, 100]);
assert_eq!(config.channels, vec![AlertChannel::Log]);
assert_eq!(config.dedup_window, 300);
}
#[test]
fn test_quota_controller_new() {
let storage = TestQuotaStorage::new();
let config = QuotaConfig::default();
let controller = QuotaController::new(storage, config);
assert_eq!(controller.config().limit, 1000);
}
#[tokio::test]
async fn test_consume_basic() {
let storage = TestQuotaStorage::new();
let config = QuotaConfig {
quota_type: QuotaType::Count,
limit: 100,
window_size: 3600,
allow_overdraft: false,
overdraft_limit_percent: 0,
alert_config: AlertConfig {
enabled: false,
..Default::default()
},
};
let controller = QuotaController::new(storage, config);
let result = controller.consume("user1", "resource1", 10).await.unwrap();
assert!(result.allowed);
assert_eq!(result.remaining, 90);
assert!(!result.alert_triggered);
let result = controller.consume("user1", "resource1", 20).await.unwrap();
assert!(result.allowed);
assert_eq!(result.remaining, 70);
}
#[tokio::test]
async fn test_consume_exceeds_limit() {
let storage = TestQuotaStorage::new();
let config = QuotaConfig {
quota_type: QuotaType::Count,
limit: 100,
window_size: 3600,
allow_overdraft: false,
overdraft_limit_percent: 0,
alert_config: AlertConfig {
enabled: false,
..Default::default()
},
};
let controller = QuotaController::new(storage, config);
let result = controller.consume("user1", "resource1", 100).await.unwrap();
assert!(result.allowed);
assert_eq!(result.remaining, 0);
let result = controller.consume("user1", "resource1", 1).await.unwrap();
assert!(!result.allowed);
assert_eq!(result.remaining, 0);
}
#[tokio::test]
async fn test_overdraft() {
let storage = TestQuotaStorage::new();
let config = QuotaConfig {
quota_type: QuotaType::Count,
limit: 100,
window_size: 3600,
allow_overdraft: true,
overdraft_limit_percent: 20,
alert_config: AlertConfig {
enabled: false,
..Default::default()
},
};
let controller = QuotaController::new(storage, config);
let result = controller.consume("user1", "resource1", 100).await.unwrap();
assert!(result.allowed);
assert_eq!(result.remaining, 20);
let result = controller.consume("user1", "resource1", 10).await.unwrap();
assert!(result.allowed);
assert_eq!(result.remaining, 10);
let result = controller.consume("user1", "resource1", 11).await.unwrap();
assert!(!result.allowed);
}
#[tokio::test]
async fn test_sliding_window_reset() {
let storage = TestQuotaStorage::new();
let config = QuotaConfig {
quota_type: QuotaType::Count,
limit: 100,
window_size: 1, allow_overdraft: false,
overdraft_limit_percent: 0,
alert_config: AlertConfig {
enabled: false,
..Default::default()
},
};
let controller = QuotaController::new(storage, config);
let result = controller.consume("user1", "resource1", 50).await.unwrap();
assert!(result.allowed);
assert_eq!(result.remaining, 50);
tokio::time::sleep(tokio::time::Duration::from_millis(1100)).await;
let result = controller.consume("user1", "resource1", 30).await.unwrap();
assert!(result.allowed);
assert!(result.allowed);
}
#[tokio::test]
async fn test_alert_trigger() {
let storage = TestQuotaStorage::new();
let config = QuotaConfig {
quota_type: QuotaType::Count,
limit: 100,
window_size: 3600,
allow_overdraft: false,
overdraft_limit_percent: 0,
alert_config: AlertConfig {
enabled: true,
thresholds: vec![80, 90, 100],
channels: vec![AlertChannel::Log],
dedup_window: DEFAULT_DEDUP_WINDOW_SECS,
},
};
let controller = QuotaController::new(storage, config);
let result = controller.consume("user1", "resource1", 80).await.unwrap();
assert!(result.allowed);
assert!(result.alert_triggered);
let result = controller.consume("user1", "resource1", 10).await.unwrap();
assert!(result.allowed);
assert!(result.alert_triggered);
let result = controller.consume("user1", "resource1", 10).await.unwrap();
assert!(result.allowed);
assert!(result.alert_triggered);
}
#[tokio::test]
async fn test_alert_dedup() {
let storage = TestQuotaStorage::new();
let config = QuotaConfig {
quota_type: QuotaType::Count,
limit: 100,
window_size: 3600,
allow_overdraft: false,
overdraft_limit_percent: 0,
alert_config: AlertConfig {
enabled: true,
thresholds: vec![80],
channels: vec![AlertChannel::Log],
dedup_window: 5, },
};
let controller = QuotaController::new(storage, config);
let result = controller.consume("user1", "resource1", 80).await.unwrap();
assert!(result.allowed);
assert!(result.alert_triggered);
let result = controller.consume("user1", "resource1", 10).await.unwrap();
assert!(result.allowed);
assert!(!result.alert_triggered);
tokio::time::sleep(tokio::time::Duration::from_millis(5100)).await;
controller.cleanup_alert_dedup();
let result = controller.consume("user1", "resource1", 5).await.unwrap();
assert!(result.allowed);
assert!(result.alert_triggered);
}
#[tokio::test]
async fn test_get_quota() {
let storage = TestQuotaStorage::new();
let config = QuotaConfig {
quota_type: QuotaType::Count,
limit: 100,
window_size: 3600,
allow_overdraft: false,
overdraft_limit_percent: 0,
alert_config: AlertConfig {
enabled: false,
..Default::default()
},
};
let controller = QuotaController::new(storage, config);
controller.consume("user1", "resource1", 50).await.unwrap();
let state = controller.get_quota("user1", "resource1").await.unwrap();
assert!(state.is_some());
assert_eq!(state.unwrap().consumed, 50);
}
#[tokio::test]
async fn test_reset_quota() {
let storage = TestQuotaStorage::new();
let config = QuotaConfig {
quota_type: QuotaType::Count,
limit: 100,
window_size: 3600,
allow_overdraft: false,
overdraft_limit_percent: 0,
alert_config: AlertConfig {
enabled: false,
..Default::default()
},
};
let controller = QuotaController::new(storage, config);
controller.consume("user1", "resource1", 50).await.unwrap();
controller.reset_quota("user1", "resource1").await.unwrap();
let state = controller.get_quota("user1", "resource1").await.unwrap();
assert!(state.is_some());
assert_eq!(state.unwrap().consumed, 0);
}
#[test]
fn test_quota_types() {
let token_config = QuotaConfig {
quota_type: QuotaType::Token,
..Default::default()
};
assert_eq!(token_config.quota_type.as_str(), "token");
let money_config = QuotaConfig {
quota_type: QuotaType::Money,
..Default::default()
};
assert_eq!(money_config.quota_type.as_str(), "money");
let count_config = QuotaConfig {
quota_type: QuotaType::Count,
..Default::default()
};
assert_eq!(count_config.quota_type.as_str(), "count");
}
#[tokio::test]
async fn test_concurrent_consume() {
let storage = TestQuotaStorage::new();
let config = QuotaConfig {
quota_type: QuotaType::Count,
limit: 100,
window_size: 3600,
allow_overdraft: false,
overdraft_limit_percent: 0,
alert_config: AlertConfig {
enabled: false,
..Default::default()
},
};
let controller = Arc::new(QuotaController::new(storage, config));
let mut handles = vec![];
for _ in 0..10 {
let controller_clone = Arc::clone(&controller);
handles.push(tokio::spawn(async move {
controller_clone.consume("user1", "resource1", 10).await
}));
}
let mut total_consumed = 0;
for handle in handles {
let result = handle.await.unwrap().unwrap();
if result.allowed {
total_consumed += 10;
}
}
assert!(total_consumed <= 100);
}
#[tokio::test]
async fn test_consume_zero() {
let storage = TestQuotaStorage::new();
let config = QuotaConfig::default();
let controller = QuotaController::new(storage, config);
let result = controller.consume("user1", "resource1", 0).await.unwrap();
assert!(result.allowed);
assert_eq!(result.remaining, 1000);
assert!(!result.alert_triggered);
}
#[test]
fn test_update_config() {
let storage = TestQuotaStorage::new();
let config = QuotaConfig::default();
let mut controller = QuotaController::new(storage, config);
assert_eq!(controller.config().limit, 1000);
let new_config = QuotaConfig {
limit: 500,
..Default::default()
};
controller.update_config(new_config);
assert_eq!(controller.config().limit, 500);
}
}