From d3b7d246da5ce9a552ad852526d013621e35bea2 Mon Sep 17 00:00:00 2001 From: "a.tolmachev" Date: Sun, 3 May 2026 21:59:49 +0000 Subject: [PATCH] cache: share ingress rate limits across instances --- TASKS.md | 3 +- apps/admin-api/src/main.rs | 12 +- apps/admin-api/src/rate_limit.rs | 2 +- apps/mcp-server/src/app.rs | 14 +-- apps/mcp-server/src/main.rs | 12 +- crates/crank-runtime/src/rate_limit.rs | 158 ++++++++++++++++++++++--- 6 files changed, 171 insertions(+), 30 deletions(-) diff --git a/TASKS.md b/TASKS.md index 9dacdae..e4a03fc 100644 --- a/TASKS.md +++ b/TASKS.md @@ -52,8 +52,9 @@ Progress: - Community fallback in-memory implementations now exist for response cache, rate-limit state, replay guard, and coordination state - Community deployment manifest now includes optional `valkey` profile and cache env wiring without making cache mandatory - shared `Valkey/Redis` backend layer now exists in `crank-runtime` with config-driven store factory for managed / enterprise contours + - admin-api and mcp-server ingress rate limiting now use shared external cache state when `CRANK_CACHE_BACKEND` is `valkey` or `redis` - pending: - - wire selected cache stores into concrete runtime/admin/mcp paths so external backend reduces real load instead of existing only as reusable layer + - wire response cache / replay guard / coordination state into additional concrete runtime and transport paths ## Planned diff --git a/apps/admin-api/src/main.rs b/apps/admin-api/src/main.rs index d4ea324..61cacd7 100644 --- a/apps/admin-api/src/main.rs +++ b/apps/admin-api/src/main.rs @@ -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> { }, }; 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> { } 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> { 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, diff --git a/apps/admin-api/src/rate_limit.rs b/apps/admin-api/src/rate_limit.rs index 33a4180..0100eb2 100644 --- a/apps/admin-api/src/rate_limit.rs +++ b/apps/admin-api/src/rate_limit.rs @@ -14,7 +14,7 @@ pub async fn apply_api_rate_limit( next: Next, ) -> Result { 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), diff --git a/apps/mcp-server/src/app.rs b/apps/mcp-server/src/app.rs index 498b0a7..597a92f 100644 --- a/apps/mcp-server/src/app.rs +++ b/apps/mcp-server/src/app.rs @@ -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, 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, 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 { diff --git a/apps/mcp-server/src/main.rs b/apps/mcp-server/src/main.rs index 2a3f79f..42502ec 100644 --- a/apps/mcp-server/src/main.rs +++ b/apps/mcp-server/src/main.rs @@ -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> { 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> { 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> { 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, diff --git a/crates/crank-runtime/src/rate_limit.rs b/crates/crank-runtime/src/rate_limit.rs index a411942..44a5aeb 100644 --- a/crates/crank-runtime/src/rate_limit.rs +++ b/crates/crank-runtime/src/rate_limit.rs @@ -1,12 +1,14 @@ use std::{ collections::HashMap, sync::{Arc, Mutex}, - time::{Duration, Instant}, + time::{Duration, Instant, SystemTime, UNIX_EPOCH}, }; +use crank_core::{RateLimitBucketState, RateLimitStateStore}; use thiserror::Error; const STALE_KEY_TTL: Duration = Duration::from_secs(300); +const TOKEN_SCALE: u64 = 1_000_000; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct RequestRateLimitConfig { @@ -45,14 +47,24 @@ pub struct RateLimitRejection { pub retry_after_ms: u64, } -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct RequestRateLimiter { config: RequestRateLimitConfig, - states: Arc>>, + backend: RequestRateLimiterBackend, +} + +#[derive(Clone)] +enum RequestRateLimiterBackend { + Local { + states: Arc>>, + }, + Shared { + store: Arc, + }, } #[derive(Clone, Copy, Debug)] -struct BucketState { +struct LocalBucketState { tokens: f64, last_refill: Instant, last_seen: Instant, @@ -62,24 +74,42 @@ impl RequestRateLimiter { pub fn new(config: RequestRateLimitConfig) -> Self { Self { config, - states: Arc::new(Mutex::new(HashMap::new())), + backend: RequestRateLimiterBackend::Local { + states: Arc::new(Mutex::new(HashMap::new())), + }, } } - pub fn check(&self, key: &str) -> Result<(), RateLimitRejection> { - self.check_at(key, Instant::now()) + pub fn new_shared(config: RequestRateLimitConfig, store: Arc) -> Self { + Self { + config, + backend: RequestRateLimiterBackend::Shared { store }, + } } - pub fn check_at(&self, key: &str, now: Instant) -> Result<(), RateLimitRejection> { - let mut states = self - .states + pub async fn check(&self, key: &str) -> Result<(), RateLimitRejection> { + match &self.backend { + RequestRateLimiterBackend::Local { .. } => self.check_local_at(key, Instant::now()), + RequestRateLimiterBackend::Shared { store } => { + self.check_shared_at(store.as_ref(), key, now_unix_ms()) + .await + } + } + } + + fn check_local_at(&self, key: &str, now: Instant) -> Result<(), RateLimitRejection> { + let RequestRateLimiterBackend::Local { states } = &self.backend else { + panic!("check_local_at called for non-local limiter"); + }; + + let mut states = states .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()); states.retain(|_, state| now.duration_since(state.last_seen) <= STALE_KEY_TTL); let burst = self.config.burst as f64; let requests_per_second = self.config.requests_per_second as f64; - let state = states.entry(key.to_owned()).or_insert(BucketState { + let state = states.entry(key.to_owned()).or_insert(LocalBucketState { tokens: burst, last_refill: now, last_seen: now, @@ -101,11 +131,61 @@ impl RequestRateLimiter { .max(1.0) as u64; Err(RateLimitRejection { retry_after_ms }) } + + async fn check_shared_at( + &self, + store: &dyn RateLimitStateStore, + key: &str, + now_unix_ms: i64, + ) -> Result<(), RateLimitRejection> { + let burst_tokens = u64::from(self.config.burst) * TOKEN_SCALE; + let refill_per_second = u64::from(self.config.requests_per_second) * TOKEN_SCALE; + let mut state = + store + .get_bucket(key) + .await + .unwrap_or(None) + .unwrap_or(RateLimitBucketState { + tokens_micros: burst_tokens, + last_refill_unix_ms: now_unix_ms, + }); + + let elapsed_ms = (now_unix_ms - state.last_refill_unix_ms).max(0) as u64; + let replenished = + state.tokens_micros + (elapsed_ms.saturating_mul(refill_per_second) / 1000); + state.tokens_micros = replenished.min(burst_tokens); + state.last_refill_unix_ms = now_unix_ms; + + if state.tokens_micros >= TOKEN_SCALE { + state.tokens_micros -= TOKEN_SCALE; + let _ = store.put_bucket(key, state, STALE_KEY_TTL).await; + return Ok(()); + } + + let missing_tokens = TOKEN_SCALE.saturating_sub(state.tokens_micros); + let retry_after_ms = missing_tokens.div_ceil(refill_per_second).max(1); + let _ = store.put_bucket(key, state, STALE_KEY_TTL).await; + Err(RateLimitRejection { retry_after_ms }) + } +} + +fn now_unix_ms() -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_else(|_| Duration::from_secs(0)) + .as_millis() + .try_into() + .unwrap_or(i64::MAX) } #[cfg(test)] mod tests { - use std::time::{Duration, Instant}; + use std::{ + sync::Arc, + time::{Duration, Instant}, + }; + + use crate::InMemoryRateLimitStateStore; use super::{RequestRateLimitConfig, RequestRateLimitConfigError, RequestRateLimiter}; @@ -126,20 +206,64 @@ mod tests { } #[test] - fn allows_burst_then_rejects_until_refilled() { + fn allows_burst_then_rejects_until_refilled_for_local_limiter() { let limiter = RequestRateLimiter::new(RequestRateLimitConfig::new(2, 2).unwrap()); let start = Instant::now(); - assert!(limiter.check_at("key", start).is_ok()); - assert!(limiter.check_at("key", start).is_ok()); + assert!(limiter.check_local_at("key", start).is_ok()); + assert!(limiter.check_local_at("key", start).is_ok()); - let rejection = limiter.check_at("key", start).unwrap_err(); + let rejection = limiter.check_local_at("key", start).unwrap_err(); assert_eq!(rejection.retry_after_ms, 500); assert!( limiter - .check_at("key", start + Duration::from_millis(500)) + .check_local_at("key", start + Duration::from_millis(500)) .is_ok() ); } + + #[tokio::test] + async fn allows_burst_then_rejects_until_refilled_for_shared_limiter() { + let limiter = RequestRateLimiter::new_shared( + RequestRateLimitConfig::new(2, 2).unwrap(), + Arc::new(InMemoryRateLimitStateStore::default()), + ); + + assert!(limiter.check_shared_at_store("key", 0).await.is_ok()); + assert!(limiter.check_shared_at_store("key", 0).await.is_ok()); + + let rejection = limiter.check_shared_at_store("key", 0).await.unwrap_err(); + assert_eq!(rejection.retry_after_ms, 1); + + assert!(limiter.check_shared_at_store("key", 500).await.is_ok()); + } + + #[tokio::test] + async fn shared_limiter_state_is_visible_across_instances() { + let store = Arc::new(InMemoryRateLimitStateStore::default()); + let first = RequestRateLimiter::new_shared( + RequestRateLimitConfig::new(1, 1).unwrap(), + store.clone(), + ); + let second = + RequestRateLimiter::new_shared(RequestRateLimitConfig::new(1, 1).unwrap(), store); + + assert!(first.check_shared_at_store("shared", 0).await.is_ok()); + assert!(second.check_shared_at_store("shared", 0).await.is_err()); + assert!(second.check_shared_at_store("shared", 1000).await.is_ok()); + } + + impl RequestRateLimiter { + async fn check_shared_at_store( + &self, + key: &str, + now_unix_ms: i64, + ) -> Result<(), super::RateLimitRejection> { + let super::RequestRateLimiterBackend::Shared { store } = &self.backend else { + panic!("check_shared_at_store called for non-shared limiter"); + }; + self.check_shared_at(store.as_ref(), key, now_unix_ms).await + } + } }