use super::{
ConcurrencyLimiter, FixedWindowLimiter, Limiter, ShardedSlidingWindowLimiter,
TokenBucketLimiter,
};
use crate::config::LimiterConfig;
use crate::config::parse_window_size;
use crate::constants::{
MAX_CONCURRENT_REQUESTS, MAX_SLIDING_LOG_REQUESTS, MAX_TOKEN_BUCKET_CAPACITY,
MAX_TOKEN_BUCKET_REFILL_RATE, MAX_WINDOW_REQUESTS,
};
use crate::error::LimiteronError;
use crate::i18n::t;
use std::sync::Arc;
pub struct LimiterFactory;
impl LimiterFactory {
pub fn create(config: &LimiterConfig) -> Result<Arc<dyn Limiter>, LimiteronError> {
match config {
LimiterConfig::TokenBucket {
capacity,
refill_rate,
} => Ok(Arc::new(TokenBucketLimiter::new(*capacity, *refill_rate))),
LimiterConfig::SlidingWindow {
window_size,
max_requests,
} => {
let duration = Self::parse_window_size(window_size)?;
Ok(Arc::new(ShardedSlidingWindowLimiter::new(
duration,
*max_requests,
)))
}
LimiterConfig::FixedWindow {
window_size,
max_requests,
} => {
let duration = Self::parse_window_size(window_size)?;
Ok(Arc::new(FixedWindowLimiter::new(duration, *max_requests)))
}
LimiterConfig::Concurrency { max_concurrent } => {
Ok(Arc::new(ConcurrencyLimiter::new(*max_concurrent)))
}
LimiterConfig::LeakyBucket {
capacity,
leak_rate,
} => Ok(Arc::new(super::leaky_bucket::LeakyBucketLimiter::new(
*capacity, *leak_rate,
)?)),
LimiterConfig::SlidingWindowLog {
window_size,
max_requests,
} => {
let duration = Self::parse_window_size(window_size)?;
Ok(Arc::new(
super::sliding_window_log::SlidingWindowLogLimiter::new(
*max_requests,
duration,
)?,
))
}
#[cfg(feature = "priority-queue")]
LimiterConfig::PriorityQueue {
window_size,
total_per_window,
level_weights,
default_priority,
} => {
let window = Self::parse_window_size(window_size)?;
Ok(Arc::new(super::priority_queue::PriorityQueueLimiter::new(
super::priority_queue::PriorityQueueConfig {
window,
total_per_window: *total_per_window,
level_weights: level_weights.clone(),
default_priority: default_priority
.unwrap_or_else(|| level_weights.len().saturating_sub(1)),
},
)?))
}
#[cfg(not(feature = "priority-queue"))]
LimiterConfig::PriorityQueue { .. } => Err(LimiteronError::LimitError(t(
"limiter-priority-queue-feature-disabled",
&[],
))),
#[cfg(feature = "admission-control")]
LimiterConfig::AdmissionControl {
max_concurrent,
max_per_second,
} => Ok(Arc::new(
super::admission_control::AdmissionController::new(
super::admission_control::AdmissionControlConfig {
max_concurrent: *max_concurrent,
max_per_second: *max_per_second,
},
),
)),
#[cfg(not(feature = "admission-control"))]
LimiterConfig::AdmissionControl { .. } => Err(LimiteronError::LimitError(t(
"limiter-admission-control-feature-disabled",
&[],
))),
LimiterConfig::Quota {
quota_type: _,
limit: _limit,
window: _window,
alert_threshold: _,
overdraft: _,
} => {
Err(LimiteronError::LimitError(t(
"limiter-quota-requires-controller",
&[],
)))
}
LimiterConfig::Custom { .. } => {
Err(LimiteronError::LimitError(t(
"limiter-custom-requires-registry",
&[],
)))
}
}
}
#[cfg(all(feature = "distributed", feature = "lua-script"))]
pub fn create_with_redis(
config: &LimiterConfig,
cache: oxcache::Cache<String, String>,
) -> Result<Arc<dyn Limiter>, LimiteronError> {
use crate::limiters::distributed::{
RedisDistributedLimiter, RedisSlidingWindowLimiter, RedisTokenBucketLimiter,
};
const DEFAULT_KEY: &str = "_global";
match config {
LimiterConfig::TokenBucket {
capacity,
refill_rate,
} => {
let inner = Arc::new(RedisDistributedLimiter::new(
cache, *capacity,
1000, ));
Ok(Arc::new(RedisTokenBucketLimiter::new(
inner,
DEFAULT_KEY,
*capacity,
*refill_rate,
)))
}
LimiterConfig::SlidingWindow {
window_size,
max_requests,
} => {
let window_ms = Self::parse_window_size(window_size)?.as_millis() as u64;
let inner = Arc::new(RedisDistributedLimiter::new(
cache,
*max_requests,
window_ms,
));
Ok(Arc::new(RedisSlidingWindowLimiter::new(
inner,
DEFAULT_KEY,
window_ms,
*max_requests,
)))
}
LimiterConfig::FixedWindow {
window_size,
max_requests,
} => {
let window_ms = Self::parse_window_size(window_size)?.as_millis() as u64;
Ok(Arc::new(RedisDistributedLimiter::new(
cache,
*max_requests,
window_ms,
)))
}
other => {
tracing::warn!(
config = ?other,
"redis backend: limiter type has no distributed script, \
falling back to in-process limiter"
);
Self::create(other)
}
}
}
pub fn create_batch(
configs: &[LimiterConfig],
) -> Result<Vec<Arc<dyn Limiter>>, LimiteronError> {
let mut limiters = Vec::with_capacity(configs.len());
for (index, config) in configs.iter().enumerate() {
let limiter = Self::create(config).map_err(|e| {
LimiteronError::LimitError(crate::i18n::t(
"limiter-create-failed",
&[
("index", (index + 1).to_string()),
("reason", e.to_string()),
],
))
})?;
limiters.push(limiter);
}
Ok(limiters)
}
pub fn parse_window_size(window_size: &str) -> Result<std::time::Duration, LimiteronError> {
parse_window_size(window_size).map_err(LimiteronError::ConfigError)
}
fn validate_window_config(
window_size: &str,
max_requests: u64,
limiter_type: &str,
max_limit: u64,
) -> Result<(), LimiteronError> {
Self::parse_window_size(window_size)?;
if max_requests == 0 {
return Err(LimiteronError::ConfigError(t(
"limiter-max-requests-must-be-positive",
&[("limiter_type", limiter_type.to_string())],
)));
}
if max_requests > max_limit {
return Err(LimiteronError::ConfigError(t(
"limiter-max-requests-too-large",
&[
("limiter_type", limiter_type.to_string()),
("max", max_limit.to_string()),
],
)));
}
Ok(())
}
pub fn validate_config(config: &LimiterConfig) -> Result<(), LimiteronError> {
match config {
LimiterConfig::TokenBucket {
capacity,
refill_rate,
} => {
if *capacity == 0 {
return Err(LimiteronError::ConfigError(t(
"limiter-capacity-must-be-positive",
&[],
)));
}
if *refill_rate == 0 {
return Err(LimiteronError::ConfigError(t(
"limiter-refill-rate-must-be-positive",
&[],
)));
}
if *capacity > MAX_TOKEN_BUCKET_CAPACITY {
return Err(LimiteronError::ConfigError(t(
"limiter-capacity-too-large",
&[("max", MAX_TOKEN_BUCKET_CAPACITY.to_string())],
)));
}
if *refill_rate > MAX_TOKEN_BUCKET_REFILL_RATE {
return Err(LimiteronError::ConfigError(t(
"limiter-refill-rate-too-large",
&[("max", MAX_TOKEN_BUCKET_REFILL_RATE.to_string())],
)));
}
}
LimiterConfig::SlidingWindow {
window_size,
max_requests,
} => {
Self::validate_window_config(
window_size,
*max_requests,
"sliding window",
MAX_WINDOW_REQUESTS,
)?;
}
LimiterConfig::FixedWindow {
window_size,
max_requests,
} => {
Self::validate_window_config(
window_size,
*max_requests,
"fixed window",
MAX_WINDOW_REQUESTS,
)?;
}
LimiterConfig::Concurrency { max_concurrent } => {
if *max_concurrent == 0 {
return Err(LimiteronError::ConfigError(t(
"limiter-max-concurrent-must-be-positive",
&[],
)));
}
if *max_concurrent > MAX_CONCURRENT_REQUESTS {
return Err(LimiteronError::ConfigError(t(
"limiter-max-concurrent-too-large",
&[("max", MAX_CONCURRENT_REQUESTS.to_string())],
)));
}
}
LimiterConfig::PriorityQueue { .. } | LimiterConfig::AdmissionControl { .. } => {
config.validate().map_err(LimiteronError::ConfigError)?;
}
LimiterConfig::LeakyBucket {
capacity,
leak_rate,
} => {
if *capacity == 0 {
return Err(LimiteronError::ConfigError(t(
"limiter-leaky-capacity-must-be-positive",
&[],
)));
}
if *leak_rate == 0 {
return Err(LimiteronError::ConfigError(t(
"limiter-leak-rate-must-be-positive",
&[],
)));
}
if *capacity > MAX_TOKEN_BUCKET_CAPACITY {
return Err(LimiteronError::ConfigError(t(
"limiter-leaky-capacity-too-large",
&[("max", MAX_TOKEN_BUCKET_CAPACITY.to_string())],
)));
}
if *leak_rate > MAX_TOKEN_BUCKET_REFILL_RATE {
return Err(LimiteronError::ConfigError(t(
"limiter-leak-rate-too-large",
&[("max", MAX_TOKEN_BUCKET_REFILL_RATE.to_string())],
)));
}
}
LimiterConfig::SlidingWindowLog {
window_size,
max_requests,
} => {
Self::validate_window_config(
window_size,
*max_requests,
"sliding window log",
MAX_SLIDING_LOG_REQUESTS,
)?;
}
LimiterConfig::Quota { .. } => {
return Err(LimiteronError::LimitError(t(
"limiter-quota-requires-controller",
&[],
)));
}
LimiterConfig::Custom { .. } => {
return Err(LimiteronError::LimitError(t(
"limiter-custom-requires-registry",
&[],
)));
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::QuotaType;
use std::time::Duration;
#[test]
fn test_create_token_bucket() {
let config = LimiterConfig::TokenBucket {
capacity: 1000,
refill_rate: 100,
};
let limiter = LimiterFactory::create(&config);
assert!(limiter.is_ok());
}
#[test]
fn test_create_sliding_window() {
let config = LimiterConfig::SlidingWindow {
window_size: "1m".to_string(),
max_requests: 60,
};
let limiter = LimiterFactory::create(&config);
assert!(limiter.is_ok());
}
#[test]
fn test_create_fixed_window() {
let config = LimiterConfig::FixedWindow {
window_size: "30s".to_string(),
max_requests: 30,
};
let limiter = LimiterFactory::create(&config);
assert!(limiter.is_ok());
}
#[test]
fn test_create_concurrency() {
let config = LimiterConfig::Concurrency { max_concurrent: 50 };
let limiter = LimiterFactory::create(&config);
assert!(limiter.is_ok());
}
#[test]
fn test_create_leaky_bucket() {
let config = LimiterConfig::LeakyBucket {
capacity: 100,
leak_rate: 10,
};
let limiter = LimiterFactory::create(&config);
assert!(limiter.is_ok());
}
#[test]
fn test_create_sliding_window_log() {
let config = LimiterConfig::SlidingWindowLog {
window_size: "30s".to_string(),
max_requests: 100,
};
let limiter = LimiterFactory::create(&config);
assert!(limiter.is_ok());
let limiter = limiter.unwrap();
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
for _ in 0..100 {
assert!(Limiter::allow(limiter.as_ref(), 1).await.unwrap());
}
assert!(!Limiter::allow(limiter.as_ref(), 1).await.unwrap());
});
}
#[test]
fn test_create_batch() {
let configs = vec![
LimiterConfig::TokenBucket {
capacity: 1000,
refill_rate: 100,
},
LimiterConfig::Concurrency { max_concurrent: 50 },
];
let limiters = LimiterFactory::create_batch(&configs);
assert!(limiters.is_ok());
assert_eq!(limiters.unwrap().len(), 2);
}
#[test]
fn test_parse_window_size_seconds() {
let duration = LimiterFactory::parse_window_size("10s");
assert!(duration.is_ok());
assert_eq!(duration.unwrap(), Duration::from_secs(10));
}
#[test]
fn test_parse_window_size_minutes() {
let duration = LimiterFactory::parse_window_size("5m");
assert!(duration.is_ok());
assert_eq!(duration.unwrap(), Duration::from_secs(5 * 60));
}
#[test]
fn test_parse_window_size_hours() {
let duration = LimiterFactory::parse_window_size("2h");
assert!(duration.is_ok());
assert_eq!(duration.unwrap(), Duration::from_secs(2 * 3600));
}
#[test]
fn test_parse_window_size_days() {
let duration = LimiterFactory::parse_window_size("1d");
assert!(duration.is_ok());
assert_eq!(duration.unwrap(), Duration::from_secs(24 * 3600));
}
#[test]
fn test_parse_window_size_invalid() {
let duration = LimiterFactory::parse_window_size("invalid");
assert!(duration.is_err());
}
#[test]
fn test_parse_window_size_empty() {
let duration = LimiterFactory::parse_window_size("");
assert!(duration.is_err());
}
#[test]
fn test_parse_window_size_zero() {
let duration = LimiterFactory::parse_window_size("0s");
assert!(duration.is_err());
}
#[test]
fn test_validate_token_bucket_valid() {
let config = LimiterConfig::TokenBucket {
capacity: 1000,
refill_rate: 100,
};
let result = LimiterFactory::validate_config(&config);
assert!(result.is_ok());
}
#[test]
fn test_validate_token_bucket_invalid_capacity() {
let config = LimiterConfig::TokenBucket {
capacity: 0,
refill_rate: 100,
};
let result = LimiterFactory::validate_config(&config);
assert!(result.is_err());
}
#[test]
fn test_validate_token_bucket_invalid_refill() {
let config = LimiterConfig::TokenBucket {
capacity: 1000,
refill_rate: 0,
};
let result = LimiterFactory::validate_config(&config);
assert!(result.is_err());
}
#[test]
fn test_validate_concurrency_valid() {
let config = LimiterConfig::Concurrency { max_concurrent: 50 };
let result = LimiterFactory::validate_config(&config);
assert!(result.is_ok());
}
#[test]
fn test_validate_concurrency_invalid() {
let config = LimiterConfig::Concurrency { max_concurrent: 0 };
let result = LimiterFactory::validate_config(&config);
assert!(result.is_err());
}
#[test]
fn test_validate_leaky_bucket() {
assert!(
LimiterFactory::validate_config(&LimiterConfig::LeakyBucket {
capacity: 100,
leak_rate: 10,
})
.is_ok()
);
assert!(
LimiterFactory::validate_config(&LimiterConfig::LeakyBucket {
capacity: 0,
leak_rate: 10,
})
.is_err()
);
assert!(
LimiterFactory::validate_config(&LimiterConfig::LeakyBucket {
capacity: 100,
leak_rate: 0,
})
.is_err()
);
assert!(
LimiterFactory::validate_config(&LimiterConfig::LeakyBucket {
capacity: MAX_TOKEN_BUCKET_CAPACITY + 1,
leak_rate: 10,
})
.is_err()
);
assert!(
LimiterFactory::validate_config(&LimiterConfig::LeakyBucket {
capacity: 100,
leak_rate: MAX_TOKEN_BUCKET_REFILL_RATE + 1,
})
.is_err()
);
}
#[test]
fn test_validate_sliding_window_log() {
assert!(
LimiterFactory::validate_config(&LimiterConfig::SlidingWindowLog {
window_size: "1m".to_string(),
max_requests: 100,
})
.is_ok()
);
assert!(
LimiterFactory::validate_config(&LimiterConfig::SlidingWindowLog {
window_size: "bogus".to_string(),
max_requests: 100,
})
.is_err()
);
assert!(
LimiterFactory::validate_config(&LimiterConfig::SlidingWindowLog {
window_size: "1m".to_string(),
max_requests: 0,
})
.is_err()
);
}
#[test]
fn test_validate_window_config_valid() {
assert!(
LimiterFactory::validate_window_config("1m", 100, "test", MAX_WINDOW_REQUESTS).is_ok()
);
assert!(
LimiterFactory::validate_window_config("1h", 1000, "test", MAX_WINDOW_REQUESTS).is_ok()
);
}
#[test]
fn test_validate_window_config_invalid_size() {
assert!(
LimiterFactory::validate_window_config("", 100, "test", MAX_WINDOW_REQUESTS).is_err()
);
}
#[test]
fn test_validate_window_config_invalid_requests() {
assert!(
LimiterFactory::validate_window_config("1m", 0, "test", MAX_WINDOW_REQUESTS).is_err()
);
}
#[test]
fn test_validate_window_config_requests_exceeded() {
assert!(
LimiterFactory::validate_window_config("1m", 10_000_001, "test", MAX_WINDOW_REQUESTS)
.is_err()
);
assert!(
LimiterFactory::validate_window_config(
"1m",
MAX_WINDOW_REQUESTS,
"sliding window log",
MAX_SLIDING_LOG_REQUESTS
)
.is_err()
);
assert!(
LimiterFactory::validate_window_config(
"1m",
MAX_SLIDING_LOG_REQUESTS,
"sliding window log",
MAX_SLIDING_LOG_REQUESTS
)
.is_ok()
);
}
#[test]
fn test_create_quota_returns_error() {
let config = LimiterConfig::Quota {
quota_type: QuotaType::Count,
limit: 100,
window: "1m".to_string(),
alert_threshold: None,
overdraft: None,
};
assert!(LimiterFactory::create(&config).is_err());
}
#[test]
fn test_create_custom_returns_error() {
let config = LimiterConfig::Custom {
name: "my_limiter".to_string(),
config: serde_json::json!({}),
};
assert!(LimiterFactory::create(&config).is_err());
}
#[test]
fn test_create_batch_with_invalid_config() {
let configs = vec![
LimiterConfig::TokenBucket {
capacity: 1000,
refill_rate: 100,
},
LimiterConfig::Custom {
name: "bad".to_string(),
config: serde_json::json!({}),
},
];
let result = LimiterFactory::create_batch(&configs);
assert!(result.is_err());
}
#[test]
fn test_create_batch_empty() {
let result = LimiterFactory::create_batch(&[]);
assert!(result.is_ok());
assert_eq!(result.unwrap().len(), 0);
}
#[test]
fn test_validate_token_bucket_capacity_exceeded() {
let config = LimiterConfig::TokenBucket {
capacity: 10_000_001,
refill_rate: 100,
};
assert!(LimiterFactory::validate_config(&config).is_err());
}
#[test]
fn test_validate_token_bucket_refill_exceeded() {
let config = LimiterConfig::TokenBucket {
capacity: 1000,
refill_rate: 1_000_001,
};
assert!(LimiterFactory::validate_config(&config).is_err());
}
#[test]
fn test_validate_sliding_window_valid() {
let config = LimiterConfig::SlidingWindow {
window_size: "1m".to_string(),
max_requests: 100,
};
assert!(LimiterFactory::validate_config(&config).is_ok());
}
#[test]
fn test_validate_sliding_window_invalid_window() {
let config = LimiterConfig::SlidingWindow {
window_size: "".to_string(),
max_requests: 100,
};
assert!(LimiterFactory::validate_config(&config).is_err());
}
#[test]
fn test_validate_fixed_window_valid() {
let config = LimiterConfig::FixedWindow {
window_size: "1m".to_string(),
max_requests: 100,
};
assert!(LimiterFactory::validate_config(&config).is_ok());
}
#[test]
fn test_validate_fixed_window_invalid_window() {
let config = LimiterConfig::FixedWindow {
window_size: "".to_string(),
max_requests: 100,
};
assert!(LimiterFactory::validate_config(&config).is_err());
}
#[test]
fn test_validate_concurrency_exceeded() {
let config = LimiterConfig::Concurrency {
max_concurrent: 100_001,
};
assert!(LimiterFactory::validate_config(&config).is_err());
}
#[test]
fn test_validate_quota_returns_limit_error() {
let config = LimiterConfig::Quota {
quota_type: QuotaType::Count,
limit: 100,
window: "1m".to_string(),
alert_threshold: None,
overdraft: None,
};
assert!(LimiterFactory::validate_config(&config).is_err());
}
#[test]
fn test_validate_custom_returns_limit_error() {
let config = LimiterConfig::Custom {
name: "my_limiter".to_string(),
config: serde_json::json!({}),
};
assert!(LimiterFactory::validate_config(&config).is_err());
}
}