diff --git a/src/code_indexer/server/middleware/admission_control.py b/src/code_indexer/server/middleware/admission_control.py new file mode 100644 index 000000000..1a1e1ade9 --- /dev/null +++ b/src/code_indexer/server/middleware/admission_control.py @@ -0,0 +1,200 @@ +"""Request admission control / backpressure middleware. + +Without this, every request queues in the anyio threadpool under overload until +it times out to a 504 -- the documented shared-queue failure mode. This +middleware caps the number of in-flight (non-exempt) requests PER WORKER PROCESS +and sheds excess immediately with ``429 Too Many Requests`` + ``Retry-After`` so +clients back off and retry instead of piling up toward the queue-collapse point. + +Scope is per worker process (each uvicorn worker has its own counter), which is +the natural unit: a worker sheds load when *its* in-flight set is full, so total +pod capacity is ``workers x max_inflight_requests``. + +Health/docs endpoints are always exempt so readiness probes (which must keep +returning 200 or k8s pulls the pod from the Service) and schema fetches are +never rejected. +""" + +from __future__ import annotations + +import hashlib +import math +import threading +from typing import Iterable, Optional, Tuple + +from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint +from starlette.requests import Request +from starlette.responses import JSONResponse, Response + +from code_indexer.server.auth.token_bucket import TokenBucketManager + +# Paths never subject to admission control: liveness/readiness + API schema. +# Prefix match, so e.g. "/docs" also covers "/docs/oauth2-redirect". +_DEFAULT_EXEMPT_PREFIXES = ("/health", "/docs", "/openapi.json", "/redoc") + +# JWT session cookie name. Mirrors auth.dependencies.CIDX_SESSION_COOKIE; kept as +# a stable local literal so this request-path middleware doesn't import the heavy +# auth module (which would risk an import cycle at app-wiring time). +_SESSION_COOKIE = "cidx_session" + +# Callers presenting no credential share a single bucket, so an unauthenticated +# flood is throttled as one group (auth rejects those requests downstream anyway). +_ANON_CONSUMER = "anon" + + +class AdmissionController: + """Thread-safe per-process in-flight counter with a hard cap. + + ``try_enter`` returns False (shed) when the cap is reached; callers that + entered must always ``leave`` (in a finally) so the slot is released. + """ + + def __init__(self, max_inflight: int, retry_after_seconds: int) -> None: + self._max = max_inflight + self.retry_after_seconds = retry_after_seconds + self._inflight = 0 + self._lock = threading.Lock() + + def try_enter(self) -> bool: + with self._lock: + # max <= 0 means "no cap" -- admit everything. + if self._max > 0 and self._inflight >= self._max: + return False + self._inflight += 1 + return True + + def leave(self) -> None: + with self._lock: + if self._inflight > 0: + self._inflight -= 1 + + @property + def inflight(self) -> int: + with self._lock: + return self._inflight + + +class PerConsumerRateLimiter: + """Per-consumer token-bucket rate limiter (dealer-wide fairness). + + The global :class:`AdmissionController` cap is fleet-wide: one abusive client + (e.g. a single dealer hammering the support agent) can consume every in-flight + slot and get everyone else shed. This limiter throttles each caller + INDIVIDUALLY, keyed by a hash of its presented credential, so a noisy consumer + is capped without starving the rest. It reuses the tested ``TokenBucketManager`` + (auth/token_bucket) — ``capacity`` is the burst allowance and ``refill_per_second`` + the sustained rate; idle buckets are reclaimed after ``cleanup_seconds``. + """ + + def __init__( + self, + capacity: int, + refill_per_second: float, + cleanup_seconds: int = 3600, + ) -> None: + self._manager = TokenBucketManager( + capacity=capacity, + refill_rate=refill_per_second, + cleanup_seconds=cleanup_seconds, + ) + + @staticmethod + def consumer_key(request: Request) -> str: + """Derive a stable per-consumer key from the request credential. + + Prefers the ``Authorization`` header (Bearer JWT or ``cidx_sk_`` API key), + then the JWT session cookie. The raw credential is never stored — only its + SHA-256 hash — so buckets can't leak secrets. Credential-less requests map + to a single shared anon bucket. + """ + cred = request.headers.get("authorization") + if not cred: + cred = request.cookies.get(_SESSION_COOKIE) + if not cred: + return _ANON_CONSUMER + return hashlib.sha256(cred.encode("utf-8")).hexdigest()[:32] + + def check(self, request: Request) -> Tuple[bool, float]: + """Consume one token for this request's consumer. + + Returns ``(allowed, retry_after_seconds)``; ``retry_after`` is meaningful + only when ``allowed`` is False. + """ + allowed, retry_after = self._manager.consume(self.consumer_key(request)) + return allowed, retry_after + + +class AdmissionControlMiddleware(BaseHTTPMiddleware): + """Sheds excess load with 429 + Retry-After. + + Two independent gates, either or both may be active: + * ``rate_limiter`` (PerConsumerRateLimiter) — per-client fairness, checked + FIRST so an abusive consumer is shed before it takes a global slot. + * ``controller`` (AdmissionController) — global per-worker in-flight cap. + """ + + def __init__( + self, + app, + controller: Optional[AdmissionController] = None, + rate_limiter: Optional[PerConsumerRateLimiter] = None, + exempt_prefixes: Iterable[str] = _DEFAULT_EXEMPT_PREFIXES, + ) -> None: + super().__init__(app) + self._controller = controller + self._rate_limiter = rate_limiter + self._exempt = tuple(exempt_prefixes) + + def _is_exempt(self, path: str) -> bool: + return path.startswith(self._exempt) + + async def dispatch( + self, request: Request, call_next: RequestResponseEndpoint + ) -> Response: + if self._is_exempt(request.url.path): + return await call_next(request) + + # Per-consumer fairness FIRST: shed an over-rate client before it takes a + # global in-flight slot from well-behaved clients. + if self._rate_limiter is not None: + allowed, retry_after = self._rate_limiter.check(request) + if not allowed: + # retry_after is inf when the bucket never refills (refill rate 0); + # cap it so math.ceil() can't raise OverflowError on a misconfig. + retry_secs = ( + max(1, math.ceil(retry_after)) + if math.isfinite(retry_after) + else 3600 + ) + return JSONResponse( + status_code=429, + headers={"Retry-After": str(retry_secs)}, + content={ + "detail": ( + "Per-client rate limit exceeded; retry after a short delay." + ), + "retry_after_seconds": retry_secs, + }, + ) + + # Global per-worker in-flight cap. + if self._controller is not None and not self._controller.try_enter(): + retry_after_s = self._controller.retry_after_seconds + return JSONResponse( + status_code=429, + headers={"Retry-After": str(retry_after_s)}, + content={ + "detail": ( + "Server at capacity (too many concurrent requests); " + "retry after a short delay." + ), + "retry_after_seconds": retry_after_s, + }, + ) + try: + return await call_next(request) + finally: + # Only release a slot we actually took (controller present + entered; + # a shed request returned above before reaching here). + if self._controller is not None: + self._controller.leave() diff --git a/src/code_indexer/server/startup/app_wiring.py b/src/code_indexer/server/startup/app_wiring.py index bc207d518..b21b37161 100644 --- a/src/code_indexer/server/startup/app_wiring.py +++ b/src/code_indexer/server/startup/app_wiring.py @@ -98,6 +98,52 @@ def create_fastapi_app(services: Dict[str, Any], lifespan: Callable) -> FastAPI: app.add_middleware(CorrelationBridgeMiddleware) + # Request admission control / backpressure. Added last so it is the OUTERMOST + # middleware -- it sheds excess load with 429 + Retry-After before any + # downstream processing. Opt-in; only wired when enabled so there is zero + # overhead by default. + _admission_cfg = server_config.admission_control_config + if _admission_cfg is not None and ( + _admission_cfg.enabled or _admission_cfg.per_consumer_enabled + ): + from code_indexer.server.middleware.admission_control import ( + AdmissionControlMiddleware, + AdmissionController, + PerConsumerRateLimiter, + ) + + _controller = ( + AdmissionController( + max_inflight=_admission_cfg.max_inflight_requests, + retry_after_seconds=_admission_cfg.retry_after_seconds, + ) + if _admission_cfg.enabled + else None + ) + _rate_limiter = ( + PerConsumerRateLimiter( + capacity=_admission_cfg.per_consumer_burst, + refill_per_second=_admission_cfg.per_consumer_refill_per_second, + cleanup_seconds=_admission_cfg.per_consumer_cleanup_seconds, + ) + if _admission_cfg.per_consumer_enabled + else None + ) + app.add_middleware( + AdmissionControlMiddleware, + controller=_controller, + rate_limiter=_rate_limiter, + ) + logging.getLogger(__name__).info( + "Admission control ENABLED: global_cap=%s (max_inflight=%d/worker), " + "per_consumer=%s (burst=%d, %.1f req/s)", + _admission_cfg.enabled, + _admission_cfg.max_inflight_requests, + _admission_cfg.per_consumer_enabled, + _admission_cfg.per_consumer_burst, + _admission_cfg.per_consumer_refill_per_second, + ) + # Add exception handlers for validation errors that FastAPI catches before middleware @app.exception_handler(RequestValidationError) def validation_exception_handler(request: Request, exc: RequestValidationError): diff --git a/src/code_indexer/server/utils/config_manager.py b/src/code_indexer/server/utils/config_manager.py index f1578ecbb..8461f7192 100644 --- a/src/code_indexer/server/utils/config_manager.py +++ b/src/code_indexer/server/utils/config_manager.py @@ -1157,6 +1157,43 @@ class QueryEmbeddingCacheConfig: query_embedding_cache_cohere_audit_sample_rate: float = 0.0 +@dataclass +class AdmissionControlConfig: + """Request admission control / backpressure. + + Under overload the server otherwise queues every request in the anyio + threadpool until it times out to a 504. When enabled, a middleware caps the + number of in-flight (non-exempt) requests PER WORKER PROCESS and/or + rate-limits each caller individually, shedding excess with an immediate + ``429 Too Many Requests`` + ``Retry-After`` so clients back off and retry + instead of piling up toward queue collapse. Health/docs endpoints are always + exempt so readiness probes and schema fetches are never rejected. + + Opt-in (both switches default False) because the right cap is + deployment-specific (roughly timeout x per-worker service rate). The two + gates are independent — enable either or both. Read at app-wiring time from + the merged server config (not a BOOTSTRAP key); a restart applies changes. + """ + + # Global per-worker in-flight cap. + enabled: bool = False + # Max concurrent non-exempt requests per worker before shedding. 0 = no cap. + max_inflight_requests: int = 100 + # Retry-After header value (seconds) sent on a shed request. + retry_after_seconds: int = 1 + + # Per-consumer (per-API-key/session) token-bucket rate limiting. INDEPENDENT + # of the global cap; keyed by a hash of the caller's credential so one client + # cannot starve the fleet. + per_consumer_enabled: bool = False + # Burst: tokens a single consumer may spend back-to-back. + per_consumer_burst: int = 30 + # Sustained rate (requests/second) a single consumer's bucket refills at. + per_consumer_refill_per_second: float = 10.0 + # Idle seconds before an unused consumer bucket is reclaimed (memory bound). + per_consumer_cleanup_seconds: int = 3600 + + @dataclass class ServerConfig: """ @@ -1244,6 +1281,7 @@ class ServerConfig: # Story #1105 - Query embedding cache configuration (runtime only, not bootstrap) query_embedding_cache_config: Optional[QueryEmbeddingCacheConfig] = None + admission_control_config: Optional[AdmissionControlConfig] = None # Story #885 - Lifecycle analysis subprocess timeout configuration lifecycle_analysis_config: Optional[LifecycleAnalysisConfig] = None @@ -1509,6 +1547,8 @@ def __post_init__(self): # Story #1105 - Initialize query embedding cache config if self.query_embedding_cache_config is None: self.query_embedding_cache_config = QueryEmbeddingCacheConfig() + if self.admission_control_config is None: + self.admission_control_config = AdmissionControlConfig() # Story #885 - Initialize lifecycle analysis config if self.lifecycle_analysis_config is None: self.lifecycle_analysis_config = LifecycleAnalysisConfig() @@ -2114,6 +2154,17 @@ def _dict_to_server_config(self, config_dict: dict) -> "ServerConfig": **{k: v for k, v in _qec_dict.items() if k in _qec_allowed} ) + # Convert admission_control_config dict to AdmissionControlConfig. Unknown + # keys are filtered for rolling-upgrade safety (same fields() pattern). + if "admission_control_config" in config_dict and isinstance( + config_dict["admission_control_config"], dict + ): + _adm_dict = config_dict["admission_control_config"] + _adm_allowed = {f.name for f in fields(AdmissionControlConfig)} + config_dict["admission_control_config"] = AdmissionControlConfig( + **{k: v for k, v in _adm_dict.items() if k in _adm_allowed} + ) + # Story #885 Phase 5b (A7d): Convert lifecycle_analysis_config dict to # LifecycleAnalysisConfig so that _merge_runtime_config produces proper # dataclass instances rather than plain dicts when loading from SQLite/PG. diff --git a/tests/unit/server/middleware/test_admission_control.py b/tests/unit/server/middleware/test_admission_control.py new file mode 100644 index 000000000..535e97dcf --- /dev/null +++ b/tests/unit/server/middleware/test_admission_control.py @@ -0,0 +1,138 @@ +"""Unit tests for request admission control / backpressure middleware. + +Covers the global in-flight cap, per-consumer rate limiting, credential keying, +exempt paths, and 429 + Retry-After shedding. Also a config-roundtrip test. +""" + +from dataclasses import fields + +from starlette.applications import Starlette +from starlette.responses import PlainTextResponse +from starlette.routing import Route +from starlette.testclient import TestClient + +from code_indexer.server.middleware.admission_control import ( + AdmissionControlMiddleware, + AdmissionController, + PerConsumerRateLimiter, +) +from code_indexer.server.utils.config_manager import AdmissionControlConfig + + +def _app(controller=None, rate_limiter=None): + async def ok(request): + return PlainTextResponse("ok") + + async def health(request): + return PlainTextResponse("healthy") + + app = Starlette(routes=[Route("/x", ok), Route("/health", health)]) + app.add_middleware( + AdmissionControlMiddleware, controller=controller, rate_limiter=rate_limiter + ) + return TestClient(app) + + +# ---------------- AdmissionController ---------------- + + +def test_controller_caps_inflight_and_releases(): + c = AdmissionController(max_inflight=1, retry_after_seconds=1) + assert c.try_enter() is True + assert c.try_enter() is False # at cap + c.leave() + assert c.try_enter() is True # slot freed + + +def test_controller_zero_means_no_cap(): + c = AdmissionController(max_inflight=0, retry_after_seconds=1) + assert all(c.try_enter() for _ in range(50)) + + +# ---------------- PerConsumerRateLimiter ---------------- + + +def test_rate_limiter_sheds_after_burst(): + rl = PerConsumerRateLimiter(capacity=2, refill_per_second=0.0) + + class _Req: + headers = {"authorization": "Bearer abc"} + cookies: dict = {} + + allowed = [rl.check(_Req())[0] for _ in range(4)] + assert allowed[:2] == [True, True] + assert allowed[2] is False # burst exhausted, no refill + + +def test_rate_limiter_keys_by_credential_not_shared(): + rl = PerConsumerRateLimiter(capacity=1, refill_per_second=0.0) + + class _A: + headers = {"authorization": "Bearer A"} + cookies: dict = {} + + class _B: + headers = {"authorization": "Bearer B"} + cookies: dict = {} + + assert rl.check(_A())[0] is True + assert rl.check(_B())[0] is True # different consumer -> own bucket + assert rl.check(_A())[0] is False # A exhausted + + +def test_consumer_key_hashes_and_never_returns_raw_credential(): + class _Req: + headers = {"authorization": "Bearer super-secret"} + cookies: dict = {} + + key = PerConsumerRateLimiter.consumer_key(_Req()) + assert "super-secret" not in key + assert len(key) == 32 + + +# ---------------- Middleware dispatch ---------------- + + +def test_middleware_sheds_with_429_and_retry_after(): + # capacity 1 then a slow refill: the 2nd request is shed with a finite + # Retry-After (the inf-refill path is covered by the isfinite guard). + rl = PerConsumerRateLimiter(capacity=1, refill_per_second=0.01) + client = _app(rate_limiter=rl) + assert client.get("/x").status_code == 200 + resp = client.get("/x") + assert resp.status_code == 429 + assert int(resp.headers["Retry-After"]) >= 1 + + +def test_middleware_shed_retry_after_survives_inf_refill(): + # refill 0 -> retry_after inf -> guard caps it instead of crashing (500). + rl = PerConsumerRateLimiter(capacity=1, refill_per_second=0.0) + client = _app(rate_limiter=rl) + assert client.get("/x").status_code == 200 + resp = client.get("/x") + assert resp.status_code == 429 + assert int(resp.headers["Retry-After"]) >= 1 + + +def test_middleware_exempts_health(): + rl = PerConsumerRateLimiter(capacity=0, refill_per_second=0.0) # sheds everything + client = _app(rate_limiter=rl) + # /health is exempt -> never shed even when the limiter would reject + assert client.get("/health").status_code == 200 + + +# ---------------- Config roundtrip ---------------- + + +def test_config_defaults_are_off(): + cfg = AdmissionControlConfig() + assert cfg.enabled is False + assert cfg.per_consumer_enabled is False + + +def test_config_from_dict_filters_unknown_keys(): + allowed = {f.name for f in fields(AdmissionControlConfig)} + raw = {"enabled": True, "max_inflight_requests": 42, "bogus_future_key": 1} + cfg = AdmissionControlConfig(**{k: v for k, v in raw.items() if k in allowed}) + assert cfg.enabled is True + assert cfg.max_inflight_requests == 42 diff --git a/tests/unit/server/services/test_config_service_bootstrap_keys_story_746.py b/tests/unit/server/services/test_config_service_bootstrap_keys_story_746.py index 38d82ba87..b9694a1ed 100644 --- a/tests/unit/server/services/test_config_service_bootstrap_keys_story_746.py +++ b/tests/unit/server/services/test_config_service_bootstrap_keys_story_746.py @@ -87,6 +87,7 @@ def test_all_server_config_fields_are_classified() -> None: "nfs_visibility_timeout_seconds", # Bug #1084 staging follow-up — runtime NFS visibility deadline "research_session_retention_days", # Bug #1085 — runtime research GC retention "query_embedding_cache_config", # Story #1105 — runtime Web UI cache config + "admission_control_config", # runtime; read from merged config at app wiring "search_event_log_retention_days", # Issue #1159 — runtime Web UI setting "export_retention_days", # Story #1160 — runtime Web UI export retention # Story #1197: four launch keys moved from BOOTSTRAP_KEYS to runtime