cache: share ingress rate limits across instances
This commit is contained in:
@@ -22,7 +22,8 @@ use crate::{
|
||||
state::AppState,
|
||||
};
|
||||
use crank_runtime::{
|
||||
RequestRateLimitConfig, RequestRateLimiter, RuntimeExecutor, RuntimeLimits, SecretCrypto,
|
||||
RequestRateLimitConfig, RequestRateLimiter, RuntimeCacheConfig, RuntimeCacheStores,
|
||||
RuntimeExecutor, RuntimeLimits, SecretCrypto,
|
||||
};
|
||||
|
||||
#[tokio::main]
|
||||
@@ -61,6 +62,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
},
|
||||
};
|
||||
let runtime_limits = RuntimeLimits::from_env()?;
|
||||
let cache_config = RuntimeCacheConfig::from_env()?;
|
||||
let cache_stores = RuntimeCacheStores::from_config(&cache_config).await?;
|
||||
let api_rate_limit = admin_api_rate_limit_config_from_env()?;
|
||||
let secret_crypto = SecretCrypto::new(&env::var("CRANK_MASTER_KEY")?)?;
|
||||
let runtime = RuntimeExecutor::with_limits(runtime_limits);
|
||||
@@ -77,7 +80,11 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
}
|
||||
let state = AppState {
|
||||
service,
|
||||
api_rate_limiter: RequestRateLimiter::new(api_rate_limit),
|
||||
api_rate_limiter: if cache_config.backend.is_external() {
|
||||
RequestRateLimiter::new_shared(api_rate_limit, cache_stores.rate_limit.clone())
|
||||
} else {
|
||||
RequestRateLimiter::new(api_rate_limit)
|
||||
},
|
||||
};
|
||||
let app = build_app(state);
|
||||
let listener = TcpListener::bind(socket_addr).await?;
|
||||
@@ -89,6 +96,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
runtime_max_concurrent_jobs = runtime_limits.max_concurrent_jobs,
|
||||
admin_rate_limit_rps = api_rate_limit.requests_per_second,
|
||||
admin_rate_limit_burst = api_rate_limit.burst,
|
||||
cache_backend = %cache_config.backend,
|
||||
max_connections = pool_config.max_connections,
|
||||
min_connections = pool_config.min_connections,
|
||||
acquire_timeout_ms = pool_config.acquire_timeout_ms,
|
||||
|
||||
@@ -14,7 +14,7 @@ pub async fn apply_api_rate_limit(
|
||||
next: Next,
|
||||
) -> Result<Response, ApiError> {
|
||||
let key = rate_limit_key(request.headers(), request.uri().path());
|
||||
if let Err(rejection) = state.api_rate_limiter.check(&key) {
|
||||
if let Err(rejection) = state.api_rate_limiter.check(&key).await {
|
||||
return Err(ApiError::rate_limited_with_context(
|
||||
"request rate limit exceeded",
|
||||
rejection_context(rejection),
|
||||
|
||||
@@ -183,7 +183,7 @@ async fn mcp_get(
|
||||
return status.into_response();
|
||||
}
|
||||
|
||||
if let Err(rejection) = enforce_transport_rate_limit(&state, &path, &headers) {
|
||||
if let Err(rejection) = enforce_transport_rate_limit(&state, &path, &headers).await {
|
||||
return rate_limited_status_response(rejection);
|
||||
}
|
||||
|
||||
@@ -231,7 +231,7 @@ async fn mcp_delete(
|
||||
return status.into_response();
|
||||
}
|
||||
|
||||
if let Err(rejection) = enforce_transport_rate_limit(&state, &path, &headers) {
|
||||
if let Err(rejection) = enforce_transport_rate_limit(&state, &path, &headers).await {
|
||||
return rate_limited_status_response(rejection);
|
||||
}
|
||||
|
||||
@@ -293,7 +293,7 @@ async fn mcp_post(
|
||||
}
|
||||
};
|
||||
|
||||
if let Err(rejection) = enforce_post_rate_limit(&state, &path, &headers) {
|
||||
if let Err(rejection) = enforce_post_rate_limit(&state, &path, &headers).await {
|
||||
return with_request_id_header(
|
||||
rate_limited_jsonrpc_response(&message, response_mode, &protocol_version, rejection),
|
||||
&transport_request_id,
|
||||
@@ -1908,21 +1908,21 @@ fn bearer_token(headers: &HeaderMap) -> Option<&str> {
|
||||
Some(token)
|
||||
}
|
||||
|
||||
fn enforce_post_rate_limit(
|
||||
async fn enforce_post_rate_limit(
|
||||
state: &Arc<AppState>,
|
||||
path: &AgentRoutePath,
|
||||
headers: &HeaderMap,
|
||||
) -> Result<(), RateLimitRejection> {
|
||||
enforce_transport_rate_limit(state, path, headers)
|
||||
enforce_transport_rate_limit(state, path, headers).await
|
||||
}
|
||||
|
||||
fn enforce_transport_rate_limit(
|
||||
async fn enforce_transport_rate_limit(
|
||||
state: &Arc<AppState>,
|
||||
path: &AgentRoutePath,
|
||||
headers: &HeaderMap,
|
||||
) -> Result<(), RateLimitRejection> {
|
||||
let key = rate_limit_key(path, headers);
|
||||
state.api_rate_limiter.check(&key)
|
||||
state.api_rate_limiter.check(&key).await
|
||||
}
|
||||
|
||||
fn rate_limit_key(path: &AgentRoutePath, headers: &HeaderMap) -> String {
|
||||
|
||||
@@ -8,7 +8,8 @@ use std::{env, net::SocketAddr, time::Duration};
|
||||
|
||||
use crank_registry::{PostgresPoolConfig, PostgresRegistry};
|
||||
use crank_runtime::{
|
||||
RequestRateLimitConfig, RequestRateLimiter, RuntimeExecutor, RuntimeLimits, SecretCrypto,
|
||||
RequestRateLimitConfig, RequestRateLimiter, RuntimeCacheConfig, RuntimeCacheStores,
|
||||
RuntimeExecutor, RuntimeLimits, SecretCrypto,
|
||||
};
|
||||
use sqlx::postgres::PgConnectOptions;
|
||||
use tokio::net::TcpListener;
|
||||
@@ -38,6 +39,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let socket_addr: SocketAddr = bind_addr.parse()?;
|
||||
let pool_config = PostgresPoolConfig::from_env()?;
|
||||
let runtime_limits = RuntimeLimits::from_env()?;
|
||||
let cache_config = RuntimeCacheConfig::from_env()?;
|
||||
let cache_stores = RuntimeCacheStores::from_config(&cache_config).await?;
|
||||
let api_rate_limit = mcp_api_rate_limit_config_from_env()?;
|
||||
let database_options = database_options_from_env()?;
|
||||
let registry = PostgresRegistry::connect_with_options_and_pool_config(
|
||||
@@ -58,7 +61,11 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
base_url,
|
||||
secret_crypto,
|
||||
runtime,
|
||||
RequestRateLimiter::new(api_rate_limit),
|
||||
if cache_config.backend.is_external() {
|
||||
RequestRateLimiter::new_shared(api_rate_limit, cache_stores.rate_limit.clone())
|
||||
} else {
|
||||
RequestRateLimiter::new(api_rate_limit)
|
||||
},
|
||||
std::sync::Arc::new(session_store),
|
||||
std::sync::Arc::new(CommunityMachineCredentialVerifier),
|
||||
);
|
||||
@@ -71,6 +78,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
runtime_max_concurrent_jobs = runtime_limits.max_concurrent_jobs,
|
||||
mcp_rate_limit_rps = api_rate_limit.requests_per_second,
|
||||
mcp_rate_limit_burst = api_rate_limit.burst,
|
||||
cache_backend = %cache_config.backend,
|
||||
max_connections = pool_config.max_connections,
|
||||
min_connections = pool_config.min_connections,
|
||||
acquire_timeout_ms = pool_config.acquire_timeout_ms,
|
||||
|
||||
Reference in New Issue
Block a user