#[cfg(feature = "quota-control")]
mod quota_limiter;
use crate::constants::MAX_COST;
use crate::constants::MAX_SPIN_ITERATIONS;
use crate::error::FlowGuardError;
use std::collections::VecDeque;
use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
fn validate_cost(cost: u64) -> Result<u64, FlowGuardError> {
if cost == 0 {
return Err(FlowGuardError::ConfigError(
"Cost cannot be zero".to_string(),
));
}
if cost > MAX_COST {
return Err(FlowGuardError::ConfigError(format!(
"Cost exceeds maximum limit ({})",
MAX_COST
)));
}
Ok(cost)
}
pub trait Limiter: Send + Sync {
fn allow(
&self,
cost: u64,
) -> Pin<Box<dyn Future<Output = Result<bool, FlowGuardError>> + Send + '_>>;
fn check(
&self,
_key: &str,
) -> Pin<Box<dyn Future<Output = Result<(), FlowGuardError>> + Send + '_>> {
Box::pin(async move {
self.allow(1).await?;
Ok(())
})
}
}
pub struct TokenBucketLimiter {
capacity: u64,
tokens: std::sync::atomic::AtomicU64,
refill_rate: u64,
last_refill: std::sync::atomic::AtomicU64,
}
impl TokenBucketLimiter {
pub fn new(capacity: u64, refill_rate: u64) -> Self {
Self {
capacity,
tokens: std::sync::atomic::AtomicU64::new(capacity),
refill_rate,
last_refill: std::sync::atomic::AtomicU64::new(
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos() as u64,
),
}
}
fn refill_tokens(&self) {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos() as u64;
loop {
let last = self.last_refill.load(std::sync::atomic::Ordering::Acquire);
let elapsed_nanos = now.saturating_sub(last);
if elapsed_nanos < 1_000_000 {
break;
}
let elapsed_seconds = elapsed_nanos as f64 / 1_000_000_000.0;
let tokens_to_add = (elapsed_seconds * self.refill_rate as f64) as u64;
if tokens_to_add == 0 {
break;
}
if self
.last_refill
.compare_exchange(
last,
now,
std::sync::atomic::Ordering::Release,
std::sync::atomic::Ordering::Relaxed,
)
.is_ok()
{
loop {
let current = self.tokens.load(std::sync::atomic::Ordering::Acquire);
let new_tokens = current.saturating_add(tokens_to_add).min(self.capacity);
if self
.tokens
.compare_exchange(
current,
new_tokens,
std::sync::atomic::Ordering::Release,
std::sync::atomic::Ordering::Relaxed,
)
.is_ok()
{
break;
}
}
break;
}
}
}
fn try_consume(&self, cost: u64) -> bool {
let mut retry_count = 0u32;
const MAX_RETRY: u32 = 3;
loop {
let current = self.tokens.load(std::sync::atomic::Ordering::Acquire);
if current < cost {
return false;
}
match self.tokens.compare_exchange(
current,
current - cost,
std::sync::atomic::Ordering::Release,
std::sync::atomic::Ordering::Relaxed,
) {
Ok(_) => return true,
Err(_) => {
retry_count += 1;
if retry_count >= MAX_RETRY {
return false;
}
if retry_count > 1 {
let backoff = 1u64 << (retry_count - 2);
for _ in 0..backoff.min(MAX_SPIN_ITERATIONS) {
std::hint::spin_loop();
}
}
}
}
}
}
#[cfg(test)]
fn get_tokens(&self) -> u64 {
self.tokens.load(std::sync::atomic::Ordering::SeqCst)
}
}
impl Limiter for TokenBucketLimiter {
fn allow(
&self,
cost: u64,
) -> Pin<Box<dyn Future<Output = Result<bool, FlowGuardError>> + Send + '_>> {
Box::pin(async move {
let cost = validate_cost(cost)?;
self.refill_tokens();
Ok(self.try_consume(cost))
})
}
}
pub struct SlidingWindowLimiter {
window_size: Duration,
max_requests: u64,
requests: Arc<Mutex<VecDeque<Instant>>>,
}
impl SlidingWindowLimiter {
pub fn new(window_size: Duration, max_requests: u64) -> Self {
let capacity = (max_requests as usize).min(10_000);
Self {
window_size,
max_requests,
requests: Arc::new(Mutex::new(VecDeque::with_capacity(capacity))),
}
}
fn cleanup_expired_requests(&self) {
let mut requests = self.requests.lock().unwrap();
let now = Instant::now();
while let Some(&front) = requests.front() {
if now.duration_since(front) > self.window_size {
requests.pop_front();
} else {
break;
}
}
}
#[cfg(test)]
fn get_request_count(&self) -> usize {
self.cleanup_expired_requests();
self.requests.lock().unwrap().len()
}
}
impl Limiter for SlidingWindowLimiter {
fn allow(
&self,
cost: u64,
) -> Pin<Box<dyn Future<Output = Result<bool, FlowGuardError>> + Send + '_>> {
Box::pin(async move {
let cost = validate_cost(cost)?;
self.cleanup_expired_requests();
let mut requests = self.requests.lock().unwrap();
let current_count = requests.len() as u64;
if current_count + cost > self.max_requests {
return Ok(false);
}
let now = Instant::now();
for _ in 0..cost {
requests.push_back(now);
}
Ok(true)
})
}
}
pub struct FixedWindowLimiter {
window_size: Duration,
max_requests: u64,
count: std::sync::atomic::AtomicU64,
window_start: std::sync::atomic::AtomicU64,
}
impl FixedWindowLimiter {
pub fn new(window_size: Duration, max_requests: u64) -> Self {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos() as u64;
Self {
window_size,
max_requests,
count: std::sync::atomic::AtomicU64::new(0),
window_start: std::sync::atomic::AtomicU64::new(now),
}
}
fn check_and_reset_window(&self) {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos() as u64;
let window_size_nanos = self.window_size.as_nanos() as u64;
loop {
let current_start = self.window_start.load(std::sync::atomic::Ordering::Acquire);
let window_end = current_start.saturating_add(window_size_nanos);
if now < window_end {
break;
}
let elapsed = now.saturating_sub(current_start);
let windows_passed = elapsed / window_size_nanos;
let new_start = current_start.saturating_add(windows_passed * window_size_nanos);
match self.window_start.compare_exchange(
current_start,
new_start,
std::sync::atomic::Ordering::Release,
std::sync::atomic::Ordering::Relaxed,
) {
Ok(_) => {
self.count.store(0, std::sync::atomic::Ordering::Release);
break;
}
Err(_) => continue,
}
}
}
#[cfg(test)]
fn get_count(&self) -> u64 {
self.check_and_reset_window();
self.count.load(std::sync::atomic::Ordering::Acquire)
}
}
impl Limiter for FixedWindowLimiter {
fn allow(
&self,
cost: u64,
) -> Pin<Box<dyn Future<Output = Result<bool, FlowGuardError>> + Send + '_>> {
Box::pin(async move {
let cost = validate_cost(cost)?;
self.check_and_reset_window();
loop {
let current = self.count.load(std::sync::atomic::Ordering::Acquire);
if current + cost > self.max_requests {
return Ok(false);
}
match self.count.compare_exchange(
current,
current + cost,
std::sync::atomic::Ordering::Release,
std::sync::atomic::Ordering::Relaxed,
) {
Ok(_) => return Ok(true),
Err(_) => continue, }
}
})
}
}
pub struct ConcurrencyLimiter {
semaphore: Arc<tokio::sync::Semaphore>,
timeout: Option<Duration>,
}
impl ConcurrencyLimiter {
pub fn new(max_concurrent: u64) -> Self {
Self {
semaphore: Arc::new(tokio::sync::Semaphore::new(max_concurrent as usize)),
timeout: None,
}
}
pub fn with_timeout(max_concurrent: u64, timeout: Duration) -> Self {
Self {
semaphore: Arc::new(tokio::sync::Semaphore::new(max_concurrent as usize)),
timeout: Some(timeout),
}
}
pub async fn acquire(
&self,
cost: u64,
) -> Result<tokio::sync::SemaphorePermit<'_>, FlowGuardError> {
let cost_u32 = cost as u32;
if cost_u32 as u64 != cost {
return Err(FlowGuardError::LimitError(
"许可数量超出 u32 范围".to_string(),
));
}
let permit = match self.timeout {
Some(timeout) => tokio::time::timeout(timeout, self.semaphore.acquire_many(cost_u32))
.await
.map_err(|_| FlowGuardError::LimitError("获取许可超时".to_string()))?
.map_err(|_| FlowGuardError::LimitError("信号量已关闭".to_string()))?,
None => self
.semaphore
.acquire_many(cost_u32)
.await
.map_err(|_| FlowGuardError::LimitError("信号量已关闭".to_string()))?,
};
Ok(permit)
}
#[cfg(test)]
fn available_permits(&self) -> usize {
self.semaphore.available_permits()
}
#[cfg(test)]
fn try_acquire(&self, cost: u64) -> Result<tokio::sync::SemaphorePermit<'_>, FlowGuardError> {
let cost_u32 = cost as u32;
if cost_u32 as u64 != cost {
return Err(FlowGuardError::LimitError(
"许可数量超出 u32 范围".to_string(),
));
}
self.semaphore
.try_acquire_many(cost_u32)
.map_err(|e| FlowGuardError::LimitError(format!("获取许可失败: {:?}", e)))
}
}
impl Limiter for ConcurrencyLimiter {
fn allow(
&self,
cost: u64,
) -> Pin<Box<dyn Future<Output = Result<bool, FlowGuardError>> + Send + '_>> {
Box::pin(async move {
let cost_u32 = cost as u32;
if cost_u32 as u64 != cost {
return Err(FlowGuardError::LimitError(
"许可数量超出 u32 范围".to_string(),
));
}
match self.semaphore.try_acquire_many(cost_u32) {
Ok(_permit) => {
Ok(true)
}
Err(_) => Ok(false),
}
})
}
}
#[cfg(feature = "quota-control")]
pub use quota_limiter::QuotaLimiter;
#[cfg(test)]
mod tests {
use super::*;
use std::time::Duration;
use tokio::time::sleep;
#[tokio::test]
async fn test_token_bucket_basic() {
let limiter = TokenBucketLimiter::new(100, 10);
assert!(limiter.allow(10).await.unwrap());
assert_eq!(limiter.get_tokens(), 90);
}
#[tokio::test]
async fn test_token_bucket_insufficient_tokens() {
let limiter = TokenBucketLimiter::new(10, 1);
assert!(limiter.allow(10).await.unwrap());
assert!(!limiter.allow(1).await.unwrap());
}
#[tokio::test]
async fn test_token_bucket_refill() {
let limiter = TokenBucketLimiter::new(10, 100); limiter.allow(10).await.unwrap();
assert_eq!(limiter.get_tokens(), 0);
sleep(Duration::from_millis(20)).await; limiter.allow(1).await.unwrap(); assert!(limiter.get_tokens() >= 1);
}
#[tokio::test]
async fn test_token_bucket_concurrent() {
let limiter = Arc::new(TokenBucketLimiter::new(100, 10));
let mut handles = vec![];
for _ in 0..10 {
let limiter_clone = Arc::clone(&limiter);
handles.push(tokio::spawn(async move {
for _ in 0..10 {
limiter_clone.allow(1).await.unwrap();
}
}));
}
for handle in handles {
handle.await.unwrap();
}
assert_eq!(limiter.get_tokens(), 0);
}
#[tokio::test]
async fn test_token_bucket_no_overconsumption() {
let limiter = Arc::new(TokenBucketLimiter::new(10, 1));
let mut handles = vec![];
for _ in 0..100 {
let limiter_clone = Arc::clone(&limiter);
handles.push(tokio::spawn(async move {
limiter_clone.allow(1).await.unwrap()
}));
}
let mut allowed_count = 0;
for handle in handles {
if handle.await.unwrap() {
allowed_count += 1;
}
}
assert!(allowed_count <= 10);
}
#[tokio::test]
async fn test_sliding_window_basic() {
let limiter = SlidingWindowLimiter::new(Duration::from_secs(1), 10);
assert!(limiter.allow(1).await.unwrap());
assert_eq!(limiter.get_request_count(), 1);
}
#[tokio::test]
async fn test_sliding_window_exceeds_limit() {
let limiter = SlidingWindowLimiter::new(Duration::from_secs(1), 10);
for _ in 0..10 {
assert!(limiter.allow(1).await.unwrap());
}
assert!(!limiter.allow(1).await.unwrap());
}
#[tokio::test]
async fn test_sliding_window_sliding() {
let limiter = SlidingWindowLimiter::new(Duration::from_millis(100), 5);
for _ in 0..5 {
assert!(limiter.allow(1).await.unwrap());
}
assert!(!limiter.allow(1).await.unwrap());
sleep(Duration::from_millis(101)).await;
assert!(limiter.allow(1).await.unwrap());
}
#[tokio::test]
async fn test_sliding_window_concurrent() {
let limiter = Arc::new(SlidingWindowLimiter::new(Duration::from_secs(1), 10));
let mut handles = vec![];
for _ in 0..20 {
let limiter_clone = Arc::clone(&limiter);
handles.push(tokio::spawn(async move {
limiter_clone.allow(1).await.unwrap()
}));
}
let mut allowed_count = 0;
for handle in handles {
if handle.await.unwrap() {
allowed_count += 1;
}
}
assert!(allowed_count <= 10);
}
#[tokio::test]
async fn test_sliding_window_cost() {
let limiter = SlidingWindowLimiter::new(Duration::from_secs(1), 10);
assert!(limiter.allow(5).await.unwrap());
assert!(limiter.allow(5).await.unwrap());
assert!(!limiter.allow(1).await.unwrap());
}
#[tokio::test]
async fn test_fixed_window_basic() {
let limiter = FixedWindowLimiter::new(Duration::from_secs(1), 10);
assert!(limiter.allow(1).await.unwrap());
assert_eq!(limiter.get_count(), 1);
}
#[tokio::test]
async fn test_fixed_window_exceeds_limit() {
let limiter = FixedWindowLimiter::new(Duration::from_secs(1), 10);
for _ in 0..10 {
assert!(limiter.allow(1).await.unwrap());
}
assert!(!limiter.allow(1).await.unwrap());
}
#[tokio::test]
async fn test_fixed_window_reset() {
let limiter = FixedWindowLimiter::new(Duration::from_millis(100), 5);
for _ in 0..5 {
assert!(limiter.allow(1).await.unwrap());
}
assert!(!limiter.allow(1).await.unwrap());
sleep(Duration::from_millis(101)).await;
assert!(limiter.allow(1).await.unwrap());
}
#[tokio::test]
async fn test_fixed_window_concurrent() {
let limiter = Arc::new(FixedWindowLimiter::new(Duration::from_secs(1), 10));
let mut handles = vec![];
for _ in 0..20 {
let limiter_clone = Arc::clone(&limiter);
handles.push(tokio::spawn(async move {
limiter_clone.allow(1).await.unwrap()
}));
}
let mut allowed_count = 0;
for handle in handles {
if handle.await.unwrap() {
allowed_count += 1;
}
}
assert!(allowed_count <= 10);
}
#[tokio::test]
async fn test_fixed_window_cost() {
let limiter = FixedWindowLimiter::new(Duration::from_secs(1), 10);
assert!(limiter.allow(5).await.unwrap());
assert!(limiter.allow(5).await.unwrap());
assert!(!limiter.allow(1).await.unwrap());
}
#[tokio::test]
async fn test_concurrency_limiter_basic() {
let limiter = ConcurrencyLimiter::new(10);
assert!(limiter.allow(1).await.unwrap());
assert_eq!(limiter.available_permits(), 10);
}
#[tokio::test]
async fn test_concurrency_limiter_exceeds_limit() {
let limiter = ConcurrencyLimiter::new(5);
for _ in 0..10 {
assert!(limiter.allow(1).await.unwrap());
}
}
#[tokio::test]
async fn test_concurrency_limiter_with_timeout() {
let limiter = ConcurrencyLimiter::with_timeout(1, Duration::from_millis(100));
assert!(limiter.allow(1).await.unwrap());
assert!(limiter.allow(1).await.unwrap());
}
#[tokio::test]
async fn test_concurrency_limiter_acquire_release() {
let limiter = Arc::new(ConcurrencyLimiter::new(2));
let permit1 = limiter.acquire(1).await.unwrap();
assert_eq!(limiter.available_permits(), 1);
let _permit2 = limiter.acquire(1).await.unwrap();
assert_eq!(limiter.available_permits(), 0);
assert!(limiter.try_acquire(1).is_err());
drop(permit1);
assert_eq!(limiter.available_permits(), 1);
let _permit3 = limiter.acquire(1).await.unwrap();
assert_eq!(limiter.available_permits(), 0);
}
#[tokio::test]
async fn test_concurrency_limiter_concurrent_acquire() {
let limiter = Arc::new(ConcurrencyLimiter::new(5));
let mut handles = vec![];
let barrier = Arc::new(tokio::sync::Barrier::new(10));
let start_signal = Arc::new(std::sync::atomic::AtomicBool::new(false));
for _ in 0..10 {
let limiter_clone = Arc::clone(&limiter);
let barrier_clone = Arc::clone(&barrier);
let start_signal_clone = Arc::clone(&start_signal);
handles.push(tokio::spawn(async move {
barrier_clone.wait().await;
loop {
if start_signal_clone.load(std::sync::atomic::Ordering::SeqCst) {
break;
}
}
match limiter_clone.try_acquire(1) {
Ok(_permit) => {
sleep(Duration::from_millis(50)).await;
true
}
Err(_) => false,
}
}));
}
start_signal.store(true, std::sync::atomic::Ordering::SeqCst);
let mut allowed_count = 0;
for handle in handles {
if handle.await.unwrap() {
allowed_count += 1;
}
}
assert!(allowed_count <= 5);
}
#[tokio::test]
async fn test_concurrency_limiter_allow_does_not_hold() {
let limiter = Arc::new(ConcurrencyLimiter::new(2));
assert!(limiter.allow(1).await.unwrap());
assert!(limiter.allow(1).await.unwrap());
assert!(limiter.allow(1).await.unwrap());
let _permit1 = limiter.acquire(1).await.unwrap();
assert_eq!(limiter.available_permits(), 1);
let _permit2 = limiter.acquire(1).await.unwrap();
assert_eq!(limiter.available_permits(), 0);
assert!(limiter.try_acquire(1).is_err());
}
}