diff --git a/.gitignore b/.gitignore index 0cab644bab..c7ba921959 100644 --- a/.gitignore +++ b/.gitignore @@ -22,9 +22,13 @@ profiles.clj *.ruby-version .cljfmt.edn dev-system/local.edn -*pycache* .portal .snyk +*pycache* +*.pyc +*.egg-info/ +venv/ +.venv/ ############################### ### Test Files diff --git a/search-proxy/Dockerfile b/search-proxy/Dockerfile new file mode 100644 index 0000000000..26e5e5cd98 --- /dev/null +++ b/search-proxy/Dockerfile @@ -0,0 +1,17 @@ +FROM python:3.11-slim AS builder + +WORKDIR /build +COPY src/ src/ +COPY pyproject.toml . +RUN python -m venv /opt/venv && /opt/venv/bin/pip install --no-cache-dir . + +FROM python:3.11-slim + +WORKDIR /app +COPY --from=builder /opt/venv /opt/venv +COPY src/proxy/ proxy/ +COPY lanes.json /lanes.json + +ENV PATH="/opt/venv/bin:$PATH" + +EXPOSE 3013 diff --git a/search-proxy/README.md b/search-proxy/README.md new file mode 100644 index 0000000000..85cad1ae89 --- /dev/null +++ b/search-proxy/README.md @@ -0,0 +1,144 @@ +# CMR Search Proxy + +A traffic-shaping proxy that sits in front of CMR search. It classifies incoming requests into priority lanes, enforces concurrency limits via Redis-backed distributed semaphores, and caches responses to reduce backend load. + +## How it works + +Every request is classified into one of three lanes based on query complexity: + +| Lane | Permits | Cache TTL | Overflow | Retry-After | +|------|---------|-----------|----------|-------------| +| express | 200 | 10s | standard | 5s | +| standard | 150 | 15s | — | 5s | +| heavy | 50 | 30s | — | 10s | + +**Classification rules** (first match wins): + +- **Heavy**: `include_facets`, `online_only`, `cloud_cover`, temporal facet params (`temporal_facet[`), cycle/pass params (`cycle[`, `passes[`), `options[readable_granule_name][pattern]`, `options[granule_ur][pattern]`, `options[producer_granule_id][pattern]`, shapefile uploads, `polygon[]` (multi-polygon, always heavy), single `polygon` with >20 vertices, bounding boxes with area >5000 sq degrees, more than 2 bounding boxes (`bounding_box[]` with 3+ values) +- **Standard**: `temporal`, `updated_since`, `revision_date`, `orbit_number`, `point`, `point[]`, single `circle`, small polygon (≤20 vertices), small bounding box (≤5000 sq degrees) +- **Express**: `circle[]` (explicit fast path — always express regardless of other params), and everything not matched above + +**Concurrency**: each lane has a Redis sorted set (`lane:{name}:active`). When a request arrives, expired entries are pruned, the active count is checked against the permit limit, and if under the limit the request is added as a member scored by its expiry epoch. If the lane is full, the request either overflows to the configured overflow lane or is rejected with a 429. The entry is removed when the request completes. Entries whose score has passed are pruned automatically on the next acquire, so permits from crashed tasks recover without manual intervention. + +**Cache**: successful (2xx) responses are stored in Redis keyed on a SHA-256 hash of method, path, query string, hashed auth token, `Accept` header, `cmr-search-after` header, and POST body. Cache hits skip lane acquisition entirely. + +**Load shedding response**: +``` +HTTP 429 Too Many Requests +Retry-After: 10 + +{"errors": ["Service temporarily overloaded for heavy-tier queries"]} +``` + +## Configuration + +All settings are environment variables with the `CMR_PROXY_` prefix. + +| Variable | Default | Description | +|----------|---------|-------------| +| `CMR_PROXY_BACKEND_URL` | _none — required, startup fails if unset_ | CMR search base URL (no `/search` suffix) | +| `CMR_PROXY_REDIS_URL` | _none — required, startup fails if unset_ | Redis connection URL | +| `CMR_PROXY_LANES_CONFIG` | `lanes.json` | Path to lanes config file; used when `CMR_PROXY_LANES_JSON` is not set | +| `CMR_PROXY_LANES_JSON` | — | Lanes config as a JSON string; takes precedence over `CMR_PROXY_LANES_CONFIG` when set. Intended for deployments that inject the value from Parameter Store as an environment variable | +| `CMR_PROXY_LOG_LEVEL` | `INFO` | Log level (`DEBUG`, `INFO`, `WARNING`) | +| `CMR_PROXY_MAX_REQUEST_BODY_BYTES` | `52428800` | Max POST body size (50MB) | +| `CMR_PROXY_MAX_CACHE_RESPONSE_BYTES` | `1048576` | Max response size to cache (1MB) | +| `CMR_PROXY_BACKEND_TIMEOUT_SECONDS` | `300.0` | Backend request timeout | +| `CMR_PROXY_BACKEND_MAX_CONNECTIONS` | `500` | httpx connection pool size | +| `CMR_PROXY_BACKEND_MAX_KEEPALIVE` | `200` | httpx keepalive connection pool size | +| `CMR_PROXY_REDIS_MAX_CONNECTIONS` | auto | Redis pool size; defaults to total lane permits + 100 | +| `CMR_PROXY_REDIS_SOCKET_CONNECT_TIMEOUT` | `2.0` | Redis connection timeout in seconds | +| `CMR_PROXY_REDIS_SOCKET_TIMEOUT` | `2.0` | Redis read/write timeout in seconds | +| `CMR_PROXY_REDIS_HEALTH_CHECK_INTERVAL` | `30` | Seconds between Redis keepalive pings | + +### Feature toggles + +| Variable | Default | Description | +|----------|---------|-------------| +| `CMR_PROXY_BYPASS_ENABLED` | `false` | Skip classification, cache, and lanes — pure transparent proxy | +| `CMR_PROXY_CACHE_ENABLED` | `true` | Enable response caching | +| `CMR_PROXY_LOAD_SHEDDING_ENABLED` | `true` | Return 429 when lanes are full; when false, requests proceed over capacity but are still counted in the sorted set so pressure remains visible in `/health` | +| `CMR_PROXY_CLASSIFICATION_ENABLED` | `true` | Classify requests; when false, all traffic routes to the default lane | + +## Lanes configuration + +Lane definitions live in `lanes.json`. Each lane supports: + +```json +{ + "name": "express", + "permits": 200, + "overflow": "standard", + "cache_ttl": 10, + "retry_after": 5, + "default": true +} +``` + +- `permits` — maximum concurrent in-flight requests +- `overflow` — lane to try if this one is full (optional) +- `cache_ttl` — response cache TTL in seconds (0 disables caching) +- `retry_after` — value of the `Retry-After` header on 429 responses +- `default` — exactly one lane must be marked as the default + +## Health endpoints + +### `GET /health/shallow` + +Always returns HTTP 200. Used for ALB/ECS target group health checks so that Redis or backend failures do not trigger task replacement. + +### `GET /health` + +Informational health check, not cached. Nothing automated polls it — ALB/ECS use `/health/shallow`. Currently always returns HTTP 200: dependencies report their status but do not affect the top-level `ok?`. + +```json +{ + "ok?": true, + "dependencies": { + "redis": {"ok?": true}, + "search": {"ok?": true, "reachable": true}, + "lane-express": {"ok?": true, "active": 12, "permits": 200, "at_capacity": false}, + "lane-standard": {"ok?": true, "active": 3, "permits": 150, "at_capacity": false}, + "lane-heavy": {"ok?": true, "active": 0, "permits": 50, "at_capacity": false} + } +} +``` + +When a lane is at capacity, `at_capacity` is `true` but `ok?` remains `true`. Use this endpoint to monitor lane utilization rather than to drive automated remediation. + +## Running locally + +Requires Python 3.11+ (`pyproject.toml` sets `requires-python = ">=3.11"`). +Deploys run on `python:3.11-slim` and `ruff` targets `py311`, so develop on +3.11 to match — on macOS, `brew install python@3.11`. Use a virtualenv: + +```bash +python3.11 -m venv .venv +source .venv/bin/activate + +# Install dependencies +pip install -e ".[dev]" + +# Start Redis +docker run -d -p 6379:6379 redis + +# Run the proxy +CMR_PROXY_BACKEND_URL=http://localhost:3003 \ +CMR_PROXY_REDIS_URL=redis://localhost:6379 \ +uvicorn proxy.app:app --port 8080 +``` + +Requests to `http://localhost:8080/search/collections` are proxied to the backend at `http://localhost:3003/search/collections`. + +## Running tests + +```bash +pip install -e ".[dev]" +pytest +``` + +## Operational notes + +**Leaked permits**: A permit leaks when a task is killed before `_release` runs, or when Redis is briefly unavailable during release (the exception is swallowed so the ASGI handler can still return a response). Once a leaked entry's TTL score passes (defaulting to `backend_timeout_seconds`, 300 seconds), it stops affecting lane counts — the health endpoint's `ZCOUNT` filters on the current timestamp as a lower bound, and each acquire's `ZCARD` runs after `ZREMRANGEBYSCORE` prunes expired-score entries. Physical removal from Redis happens on the next acquire for that lane. Note: if Redis is unavailable during acquire, the fail-open path applies — no permit is stored and no release is attempted, so there is no leak in that case. To immediately reset a lane without waiting for TTL, delete its sorted set key from Redis: `lane:express:active`, `lane:standard:active`, `lane:heavy:active`. + +**Debugging**: Set `CMR_PROXY_LOG_LEVEL=DEBUG` to log backend response details including content encoding and actual byte counts. Remove when done — debug logging is verbose under load. diff --git a/search-proxy/lanes.json b/search-proxy/lanes.json new file mode 100644 index 0000000000..d1b51521a4 --- /dev/null +++ b/search-proxy/lanes.json @@ -0,0 +1,22 @@ +[ + { + "name": "express", + "permits": 200, + "overflow": "standard", + "cache_ttl": 10, + "retry_after": 5, + "default": true + }, + { + "name": "standard", + "permits": 150, + "cache_ttl": 15, + "retry_after": 5 + }, + { + "name": "heavy", + "permits": 50, + "cache_ttl": 30, + "retry_after": 10 + } +] diff --git a/search-proxy/pyproject.toml b/search-proxy/pyproject.toml new file mode 100644 index 0000000000..27d057c9af --- /dev/null +++ b/search-proxy/pyproject.toml @@ -0,0 +1,40 @@ +[build-system] +requires = ["setuptools>=68.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "search-proxy" +version = "0.1.0" +description = "Traffic lane proxy for CMR search" +requires-python = ">=3.11" +dependencies = [ + "fastapi>=0.115", + "uvicorn[standard]>=0.34", + "httpx>=0.28", + "redis>=5.0", + "pydantic-settings>=2.0", + "python-json-logger>=2.0", +] + +[project.optional-dependencies] +dev = [ + "pytest>=8.0", + "pytest-asyncio>=0.24", + "fakeredis>=2.0", + "ruff>=0.11", +] + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.ruff] +target-version = "py311" +line-length = 88 +src = ["src", "test"] + +[tool.ruff.lint] +select = ["E", "F", "W", "I"] + +[tool.pytest.ini_options] +asyncio_mode = "auto" +testpaths = ["test"] diff --git a/search-proxy/src/proxy/__init__.py b/search-proxy/src/proxy/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/search-proxy/src/proxy/app.py b/search-proxy/src/proxy/app.py new file mode 100644 index 0000000000..2568ed3343 --- /dev/null +++ b/search-proxy/src/proxy/app.py @@ -0,0 +1,549 @@ +import hashlib +import logging +import time +import uuid +from contextlib import asynccontextmanager +from urllib.parse import parse_qs + +import httpx +import redis.asyncio +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse, Response +from pythonjsonlogger import json as jsonlogger + +from proxy.cache import ResponseCache +from proxy.classifier import classify_request +from proxy.config import ProxySettings, load_lanes_config, parse_lanes_config +from proxy.lanes import LoadSheddingError, RequestLanes + +logger = logging.getLogger(__name__) + +# Hop-by-hop headers per RFC 2616 §13.5.1 +HOP_HEADERS = frozenset( + { + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailer", + "transfer-encoding", + "upgrade", + # httpx decompresses transparently; these headers reflect the + # compressed transport and must not be forwarded as-is + "content-encoding", + "content-length", + } +) + +DEFAULT_TOGGLES = { + "bypass_enabled": False, + "cache_enabled": True, + "load_shedding_enabled": True, + "classification_enabled": True, +} + + +class _PrefixedJsonFormatter(jsonlogger.JsonFormatter): + """Prepend a plain-text timestamp so awslogs-datetime-format can parse it.""" + + def format(self, record: logging.LogRecord) -> str: + return f"{self.formatTime(record)} {super().format(record)}" + + +def setup_logging(level: str = "INFO"): + handler = logging.StreamHandler() + handler.setFormatter( + _PrefixedJsonFormatter("%(asctime)s %(name)s %(levelname)s %(message)s") + ) + proxy_log = logging.getLogger("proxy") + proxy_log.setLevel(level.upper()) + proxy_log.addHandler(handler) + proxy_log.propagate = False + + + +@asynccontextmanager +async def lifespan(app: FastAPI): + """Initialize shared resources on startup, clean up on shutdown.""" + settings = ProxySettings() + setup_logging(settings.log_level) + if settings.lanes_json is not None: + lanes_config = parse_lanes_config(settings.lanes_json) + logger.info("lanes_config_source", extra={"source": "env"}) + else: + lanes_config = load_lanes_config(settings.lanes_config) + logger.info("lanes_config_source", extra={"source": settings.lanes_config}) + + app.state.settings = settings + app.state.lanes_config = lanes_config + app.state.toggles = { + "bypass_enabled": settings.bypass_enabled, + "cache_enabled": settings.cache_enabled, + "load_shedding_enabled": settings.load_shedding_enabled, + "classification_enabled": settings.classification_enabled, + } + + logger.info("toggles_loaded", extra={"toggles": app.state.toggles}) + logger.info( + "settings_loaded", + extra={ + "backend_url": settings.backend_url, + "lanes_config_source": "env" if settings.lanes_json is not None else settings.lanes_config, + "lanes": [{"name": l.name, "permits": l.permits, "overflow": l.overflow, "cache_ttl": l.cache_ttl} for l in lanes_config.lanes], + "cache_enabled": settings.cache_enabled, + "load_shedding_enabled": settings.load_shedding_enabled, + "bypass_enabled": settings.bypass_enabled, + "redis_max_connections": settings.redis_max_connections or (sum(l.permits for l in lanes_config.lanes) + 100), + "backend_timeout_seconds": settings.backend_timeout_seconds, + }, + ) + + # Connection-pooled httpx client for forwarding requests to the backend + app.state.backend = httpx.AsyncClient( + base_url=settings.backend_url, + timeout=settings.backend_timeout_seconds, + limits=httpx.Limits( + max_connections=settings.backend_max_connections, + max_keepalive_connections=settings.backend_max_keepalive, + ), + ) + # Redis connection for distributed lane semaphores and response cache. + # Pool size defaults to total lane permits + 100. Each acquire involves 4 + # sequential Redis ops (TIME, ZREMRANGEBYSCORE, ZCARD, ZADD) plus a ZREM + # on release and a GET/SET for cache. Connections are released between ops + # so peak concurrent demand tracks concurrent requests, not ops per request. + # The +100 provides headroom for health checks and reconnect storms. + # Override with CMR_PROXY_REDIS_MAX_CONNECTIONS if permits change without + # a redeploy or if additional headroom is needed. + redis_max_connections = settings.redis_max_connections or ( + sum(lane.permits for lane in lanes_config.lanes) + 100 + ) + app.state.redis = redis.asyncio.from_url( + settings.redis_url, + retry_on_timeout=True, + socket_connect_timeout=settings.redis_socket_connect_timeout, + socket_timeout=settings.redis_socket_timeout, + health_check_interval=settings.redis_health_check_interval, + max_connections=redis_max_connections, + ) + app.state.lanes = RequestLanes( + lanes_config, + app.state.redis, + permit_ttl=int(settings.backend_timeout_seconds), + ) + app.state.cache = ResponseCache(app.state.redis, settings.max_cache_response_bytes) + yield + await app.state.backend.aclose() + await app.state.redis.aclose() + + +app = FastAPI(title="CMR Search Priority Proxy", lifespan=lifespan) + + +def _extract_auth_token(request: Request) -> str: + """Extract and hash the auth token for cache key segmentation. + + Returns a SHA-256 hash so the plaintext token never appears in + cache keys, logs, or memory beyond this function.""" + token = request.headers.get("Echo-Token") or request.headers.get( + "Authorization", "" + ) + if token: + return hashlib.sha256(token.encode()).hexdigest() + return "guest" + + +def filter_hop_headers(headers: httpx.Headers) -> dict: + """Remove hop-by-hop headers that must not be forwarded.""" + return { + name: value + for name, value in headers.items() + if name.lower() not in HOP_HEADERS + } + + +async def forward_to_backend( + request: Request, + path: str, + request_id: str = "", +) -> httpx.Response: + """Forward request to the backend, preserving headers, body, and query + string verbatim. Injects cmr-request-id for log correlation.""" + backend: httpx.AsyncClient = request.app.state.backend + + # Forward all headers except host and content-length, which httpx + # sets from the base_url and body respectively + headers = { + name: value + for name, value in request.headers.items() + if name.lower() not in ("host", "content-length") + } + if request_id: + headers["cmr-request-id"] = request_id + headers["x-cmr-proxy-request"] = "1" + + # Append raw query string directly to avoid double-encoding + query = str(request.url.query) + url = f"/{path}?{query}" if query else f"/{path}" + + if request.method == "GET": + response = await backend.get(url, headers=headers) + else: + body = await request.body() + response = await backend.request( + request.method, + url, + headers=headers, + content=body, + ) + + logger.debug( + "backend_response_headers", + extra={ + "request_id": request_id, + "url": url, + "status_code": response.status_code, + "content_encoding": response.headers.get("content-encoding"), + "content_length_header": response.headers.get("content-length"), + "actual_content_bytes": len(response.content), + }, + ) + return response + + +@app.get("/health/shallow") +async def health_shallow(): + return JSONResponse(status_code=200, content={"ok?": True}) + + +@app.get("/health") +async def health(request: Request): + """Health check matching CMR's {:ok? bool :dependencies {...}} format. + + Informational only and not cached — nothing automated polls this + (ALB/ECS use /health/shallow). It exists so an operator can see + dependency and lane-utilization status in one place.""" + dependencies = {} + + # Redis + try: + await request.app.state.redis.ping() + dependencies["redis"] = {"ok?": True} + except Exception as exc: + dependencies["redis"] = {"ok?": True, "problem": str(exc)} # informational + + # Backend search service + try: + resp = await request.app.state.backend.get("/search/health") + backend_ok = resp.status_code < 500 + dependencies["search"] = { + "ok?": True, # informational + "reachable": backend_ok, + } + if not backend_ok: + dependencies["search"]["problem"] = f"status {resp.status_code}" + except Exception as exc: + dependencies["search"] = {"ok?": True, "problem": str(exc)} # informational + + # Lane utilization — count only non-expired entries (score > wall clock now) + lanes_config = request.app.state.lanes_config + redis_client = request.app.state.redis + wall_now = time.time() + for lane in lanes_config.lanes: + key = f"lane:{lane.name}:active" + try: + active = await redis_client.zcount(key, wall_now, "+inf") + at_capacity = active >= lane.permits + dep = { + "ok?": True, + "active": active, + "permits": lane.permits, + "at_capacity": at_capacity, + } + if at_capacity: + dep["note"] = "at capacity" + except Exception as exc: + dep = {"ok?": True, "problem": str(exc)} # informational + dependencies[f"lane-{lane.name}"] = dep + + ok = all(dep["ok?"] for dep in dependencies.values()) + status_code = 200 if ok else 503 + content = {"ok?": ok, "dependencies": dependencies} + + return JSONResponse(status_code=status_code, content=content) + + +def _extract_request_id(request: Request) -> str: + """Extract or generate a request ID.""" + return ( + request.headers.get("cmr-request-id") + or request.headers.get("x-request-id") + or str(uuid.uuid4()) + ) + + +@app.api_route("/{path:path}", methods=["GET", "POST"]) +async def proxy(request: Request, path: str): + """Main proxy handler: classify, cache check, acquire lane, forward.""" + t0 = time.monotonic() + full_path = f"/{path}" + toggles = request.app.state.toggles + + # Read POST body once; reused for size check, form param extraction, and cache key + content_type = request.headers.get("content-type", "") + body = b"" + body_hash = "" + if request.method == "POST": + content_length = request.headers.get("content-length") + try: + claimed_size = int(content_length) if content_length else 0 + except ValueError: + claimed_size = 0 + if claimed_size > request.app.state.settings.max_request_body_bytes: + return JSONResponse( + status_code=413, + content={"errors": ["Request body too large"]}, + headers={"CMR-Request-Id": _extract_request_id(request)}, + ) + body = await request.body() + if len(body) > request.app.state.settings.max_request_body_bytes: + return JSONResponse( + status_code=413, + content={"errors": ["Request body too large"]}, + headers={"CMR-Request-Id": _extract_request_id(request)}, + ) + body_hash = hashlib.sha256(body).hexdigest() + + params = dict(request.query_params) + if request.method == "POST" and "application/x-www-form-urlencoded" in content_type: + try: + body_params = parse_qs(body.decode(), keep_blank_values=True) + for param_name, param_values in body_params.items(): + # Query string params take precedence over body params + if param_name not in params: + params[param_name] = ( + param_values[0] if len(param_values) == 1 else param_values + ) + except UnicodeDecodeError: + logger.warning( + "Could not decode POST body as UTF-8, skipping body param extraction" + ) + + request_id = _extract_request_id(request) + auth_token = _extract_auth_token(request) + query_string = str(request.url.query) + search_after = request.headers.get("cmr-search-after", "") + accept = request.headers.get("accept", "") + + # Bypass: skip classification, cache, and lanes — pure transparent proxy + if toggles["bypass_enabled"]: + try: + backend_response = await forward_to_backend(request, path, request_id) + except httpx.TimeoutException: + logger.error( + "Backend timeout: %s %s [bypass]", request.method, full_path + ) + return JSONResponse( + status_code=504, + content={"errors": ["Backend timed out"]}, + headers={"CMR-Request-Id": request_id}, + ) + except httpx.ConnectError: + logger.error( + "Backend unavailable: %s %s [bypass]", request.method, full_path + ) + return JSONResponse( + status_code=502, + content={"errors": ["Backend unavailable"]}, + headers={"CMR-Request-Id": request_id}, + ) + resp_headers = filter_hop_headers(backend_response.headers) + resp_headers["CMR-Request-Id"] = request_id + logger.info( + "request_completed", + extra={ + "request_id": request_id, + "method": request.method, + "path": full_path, + "bypass": True, + "cache_hit": False, + "status_code": backend_response.status_code, + "response_bytes": len(backend_response.content), + "duration_ms": round((time.monotonic() - t0) * 1000), + "provider": params.get("provider"), + }, + ) + return Response( + content=backend_response.content, + status_code=backend_response.status_code, + headers=resp_headers, + ) + + # Classify the request into a traffic lane based on query parameters + lanes_config = request.app.state.lanes_config + if toggles["classification_enabled"]: + lane_name = classify_request(params, content_type) + else: + lane_name = lanes_config.default_lane + + logger.info( + "request_classified", + extra={ + "request_id": request_id, + "lane": lane_name, + "method": request.method, + "path": full_path, + "provider": params.get("provider"), + "has_spatial": any( + k in params + for k in ("polygon", "bounding_box", "circle[]", "point") + ), + "classification_enabled": toggles["classification_enabled"], + }, + ) + + lanes: RequestLanes = request.app.state.lanes + lane = lanes_config.get(lane_name) + cache: ResponseCache = request.app.state.cache + + # Check cache before acquiring a lane permit + if toggles["cache_enabled"] and lane.cache_ttl > 0: + try: + cached = await cache.get( + request.method, full_path, query_string, auth_token, search_after, accept, body_hash + ) + if cached: + logger.info( + "request_completed", + extra={ + "request_id": request_id, + "method": request.method, + "path": full_path, + "lane": lane_name, + "cache_hit": True, + "status_code": cached["status_code"], + "duration_ms": round((time.monotonic() - t0) * 1000), + "provider": params.get("provider"), + }, + ) + response = Response( + content=cached["body"], + status_code=cached["status_code"], + headers=cached.get("headers", {}), + ) + response.headers["CMR-Request-Id"] = request_id + return response + except Exception: + logger.warning("cache_read_failed", extra={"request_id": request_id}, exc_info=True) + + # Acquire a distributed semaphore permit for this lane, then forward + try: + async with lanes.acquire(lane_name, toggles["load_shedding_enabled"]) as actual_lane: + try: + backend_response = await forward_to_backend(request, path, request_id) + except httpx.TimeoutException: + logger.error( + "Backend timeout: %s %s tier=%s", + request.method, + full_path, + actual_lane, + ) + return JSONResponse( + status_code=504, + content={"errors": ["Backend timed out"]}, + headers={"CMR-Request-Id": request_id}, + ) + except httpx.ConnectError: + logger.error( + "Backend unavailable: %s %s tier=%s", + request.method, + full_path, + actual_lane, + ) + return JSONResponse( + status_code=502, + content={"errors": ["Backend unavailable"]}, + headers={"CMR-Request-Id": request_id}, + ) + + # Cache successful responses if this lane has a TTL + cache_stored = False + if toggles["cache_enabled"] and lane.cache_ttl > 0 and 200 <= backend_response.status_code < 300: + response_data = { + "status_code": backend_response.status_code, + "body": backend_response.text, + "headers": filter_hop_headers(backend_response.headers), + } + try: + await cache.set( + request.method, + full_path, + query_string, + auth_token, + response_data, + len(backend_response.content), + lane.cache_ttl, + search_after, + accept, + body_hash, + ) + cache_stored = True + except Exception: + logger.warning("cache_write_failed", extra={"request_id": request_id}, exc_info=True) + + # Strip hop-by-hop headers and attach the request ID + resp_headers = filter_hop_headers(backend_response.headers) + resp_headers["CMR-Request-Id"] = request_id + + logger.info( + "request_completed", + extra={ + "request_id": request_id, + "method": request.method, + "path": full_path, + "lane": lane_name, + "actual_lane": actual_lane, + "overflow": actual_lane != lane_name, + "cache_hit": False, + "cache_stored": cache_stored, + "status_code": backend_response.status_code, + "response_bytes": len(backend_response.content), + "duration_ms": round((time.monotonic() - t0) * 1000), + "provider": params.get("provider"), + }, + ) + + return Response( + content=backend_response.content, + status_code=backend_response.status_code, + headers=resp_headers, + ) + + # Lane is full — no permit available + except LoadSheddingError as shed_error: + logger.warning( + "load_shed", + extra={ + "request_id": request_id, + "requested_lane": lane_name, + "shed_lane": shed_error.lane_name, + "method": request.method, + "path": full_path, + "provider": params.get("provider"), + "retry_after": shed_error.retry_after, + }, + ) + return JSONResponse( + status_code=429, + content={ + "errors": [ + f"Service temporarily overloaded for " + f"{shed_error.lane_name}-tier queries" + ] + }, + headers={ + "Retry-After": str(shed_error.retry_after), + "CMR-Request-Id": request_id, + }, + ) diff --git a/search-proxy/src/proxy/cache.py b/search-proxy/src/proxy/cache.py new file mode 100644 index 0000000000..a24430db38 --- /dev/null +++ b/search-proxy/src/proxy/cache.py @@ -0,0 +1,64 @@ +import hashlib +import json +import logging +from typing import Optional + +import redis.asyncio + +logger = logging.getLogger(__name__) + + +class ResponseCache: + """Redis-backed response cache keyed on the full request signature.""" + + def __init__( + self, + redis_client: redis.asyncio.Redis, + max_response_bytes: int, + ): + self.redis = redis_client + self.max_response_bytes = max_response_bytes + + def _build_key(self, method: str, path: str, query: str, auth_token: str, search_after: str = "", accept: str = "", body_hash: str = "") -> str: + """Hash the full request signature into a Redis key.""" + raw = f"{method}|{path}|{query}|{auth_token}|{search_after}|{accept}|{body_hash}" + digest = hashlib.sha256(raw.encode()).hexdigest() + return f"cache:{digest}" + + async def get( + self, method: str, path: str, query: str, auth_token: str, search_after: str = "", accept: str = "", body_hash: str = "" + ) -> Optional[dict]: + """Look up a cached response. Returns None on miss.""" + key = self._build_key(method, path, query, auth_token, search_after, accept, body_hash) + cached = await self.redis.get(key) + if cached is not None: + return json.loads(cached) + return None + + async def set( + self, + method: str, + path: str, + query: str, + auth_token: str, + response_data: dict, + response_size: int, + ttl: int, + search_after: str = "", + accept: str = "", + body_hash: str = "", + ): + """Store a response with the given TTL. Skips oversized responses.""" + if response_size > self.max_response_bytes: + logger.debug( + "cache_skip_oversized", + extra={ + "size": response_size, + "limit": self.max_response_bytes, + "path": path, + }, + ) + return + + key = self._build_key(method, path, query, auth_token, search_after, accept, body_hash) + await self.redis.set(key, json.dumps(response_data), ex=ttl) diff --git a/search-proxy/src/proxy/classifier.py b/search-proxy/src/proxy/classifier.py new file mode 100644 index 0000000000..39b1ee3e13 --- /dev/null +++ b/search-proxy/src/proxy/classifier.py @@ -0,0 +1,165 @@ +from typing import Any, Dict, List, Optional + +EXPRESS = "express" +STANDARD = "standard" +HEAVY = "heavy" + + +# Non-spatial heavy signals + +HEAVY_PARAMS = { + "include_facets", + "online_only", + "cloud_cover", +} + +HEAVY_PREFIXES = ( + "temporal_facet[", + "cycle[", + "passes[", + "options[readable_granule_name][pattern]", + "options[granule_ur][pattern]", + "options[producer_granule_id][pattern]", +) + +# Non-spatial standard signals +STANDARD_PARAMS = { + "temporal", + "temporal[]", + "updated_since", + "revision_date", + "orbit_number", +} + +# Spatial geometry thresholds + +HEAVY_POLYGON_VERTEX_THRESHOLD = 20 +HEAVY_BBOX_AREA_THRESHOLD = 5000 # square degrees +HEAVY_MULTI_BBOX_COUNT_THRESHOLD = 2 + + +def _count_polygon_vertices(polygon_value: str) -> int: + """Count vertices in a polygon coordinate string. + Format: 'lon1,lat1,lon2,lat2,...,lonN,latN' (closing vertex may repeat first).""" + ordinates = polygon_value.strip().strip('"').split(",") + coords_per_vertex = 2 # lon, lat + return len(ordinates) // coords_per_vertex + + +def _compute_bbox_area(bbox_value: str) -> float: + """Approximate area in square degrees for a bounding box. + Format: 'west,south,east,north'. + + This is deliberately naive — no antimeridian handling, no spherical correction. + Edge cases (antimeridian crossing, duplicate closing points, backwards winding) + will overestimate area, pushing toward HEAVY. That's the conservative direction; + the semaphore timeout is the real safety net for misclassification. + + On parse failure, returns a value above the heavy threshold so malformed spatial + params are treated conservatively rather than getting the express lane.""" + try: + parts = bbox_value.strip().strip('"').split(",") + west, south, east, north = (float(p) for p in parts[:4]) + return abs(east - west) * abs(north - south) + except (ValueError, IndexError): + return HEAVY_BBOX_AREA_THRESHOLD + 1 + + +def _parse_multi_value(value: Any) -> List[str]: + """Parse a param value that may be a list or a single string.""" + if isinstance(value, list): + return value + return [str(value)] + + +def _is_shapefile_request(params: Dict[str, Any], content_type: str = "") -> bool: + """Detect shapefile uploads (multipart form with shapefile).""" + return "multipart/form-data" in content_type and any( + k in params for k in ("shapefile", "file") + ) + + +def _classify_spatial(params: Dict[str, Any]) -> Optional[str]: + """Classify spatial queries using computed geometric properties. + Returns None if no spatial params are present.""" + param_keys = set(params.keys()) + + has_spatial = False + + # polygon[] (multi-polygon): always heavy + if "polygon[]" in param_keys: + return HEAVY + + # Single polygon: check vertex count + if "polygon" in param_keys: + has_spatial = True + for poly_val in _parse_multi_value(params["polygon"]): + vertices = _count_polygon_vertices(poly_val) + if vertices > HEAVY_POLYGON_VERTEX_THRESHOLD: + return HEAVY + + # bounding_box[] or bounding_box: check area and count + for bbox_key in ("bounding_box[]", "bounding_box"): + if bbox_key in param_keys: + has_spatial = True + bbox_values = _parse_multi_value(params[bbox_key]) + if ( + bbox_key == "bounding_box[]" + and len(bbox_values) > HEAVY_MULTI_BBOX_COUNT_THRESHOLD + ): + return HEAVY + for bbox_val in bbox_values: + if _compute_bbox_area(bbox_val) > HEAVY_BBOX_AREA_THRESHOLD: + return HEAVY + + # circle[]: consistently cheap (avg 157ms, indexed filter, no script) + if "circle[]" in param_keys: + return EXPRESS + + # Other spatial params (point, point[], circle): standard + if any(k in param_keys for k in ("point", "point[]", "circle")): + has_spatial = True + + if has_spatial: + return STANDARD + + # No spatial params present + return None + + +def classify_request( + params: Dict[str, Any], + content_type: str = "", +) -> str: + """Classify a search request into a lane name. + + Uses three layers: + 1. Non-spatial heavy signals (deterministic from production data) + 2. Spatial geometry analysis (vertex count, area, multi-spatial) + 3. Non-spatial standard signals (temporal, date ranges) + + Anything not matching heavy or standard rules -> express. + """ + param_keys = set(params.keys()) + + # Layer 1: Non-spatial heavy signals (always win) + if param_keys & HEAVY_PARAMS: + return HEAVY + + if any(k.startswith(HEAVY_PREFIXES) for k in param_keys): + return HEAVY + + # Shapefile: always heavy + if _is_shapefile_request(params, content_type): + return HEAVY + + # Layer 2: Spatial classification (geometry-aware) + spatial_lane = _classify_spatial(params) + if spatial_lane is not None: + return spatial_lane + + # Layer 3: Non-spatial standard signals + if param_keys & STANDARD_PARAMS: + return STANDARD + + return EXPRESS diff --git a/search-proxy/src/proxy/config.py b/search-proxy/src/proxy/config.py new file mode 100644 index 0000000000..931b6dd288 --- /dev/null +++ b/search-proxy/src/proxy/config.py @@ -0,0 +1,152 @@ +import json +from pathlib import Path +from typing import List + +from pydantic import BaseModel, field_validator, model_validator +from pydantic_settings import BaseSettings + + +class ProxySettings(BaseSettings): + backend_url: str + redis_url: str + max_request_body_bytes: int = 52_428_800 + max_cache_response_bytes: int = 1_048_576 + backend_timeout_seconds: float = 300.0 + redis_socket_connect_timeout: float = 2.0 + redis_socket_timeout: float = 2.0 + redis_health_check_interval: int = 30 + redis_max_connections: int | None = None + backend_max_connections: int = 500 + backend_max_keepalive: int = 200 + + lanes_config: str = "lanes.json" + lanes_json: str | None = None + + log_level: str = "INFO" + bypass_enabled: bool = False + cache_enabled: bool = True + load_shedding_enabled: bool = True + classification_enabled: bool = True + + model_config = {"env_prefix": "CMR_PROXY_"} + + +class LaneConfig(BaseModel): + """Configuration for a single traffic lane. Defined in lanes.json.""" + + name: str + permits: int + overflow: str | None = None + cache_ttl: int = 0 + retry_after: int = 5 + default: bool = False + + @field_validator("name") + @classmethod + def name_must_not_be_blank(cls, v): + if not v.strip(): + raise ValueError("lane name must not be blank or whitespace-only") + return v + + @field_validator("permits") + @classmethod + def permits_must_be_positive(cls, v): + if v < 1: + raise ValueError(f"permits must be at least 1, got {v}") + return v + + @field_validator("cache_ttl") + @classmethod + def cache_ttl_must_be_non_negative(cls, v): + if v < 0: + raise ValueError(f"cache_ttl must be >= 0, got {v}") + return v + + @field_validator("retry_after") + @classmethod + def retry_after_must_be_positive(cls, v): + if v < 1: + raise ValueError(f"retry_after must be at least 1, got {v}") + return v + + +class LanesConfig(BaseModel): + """Validated collection of lane definitions loaded from lanes.json.""" + + lanes: List[LaneConfig] + + @model_validator(mode="after") + def validate_lanes(self): + all_names = [lane.name for lane in self.lanes] + names = set(all_names) + + # Lane names must be unique + if len(all_names) != len(names): + seen, dupes = set(), [] + for n in all_names: + if n in seen: + dupes.append(n) + seen.add(n) + raise ValueError(f"Duplicate lane names: {sorted(dupes)}") + + # Every overflow target must reference an existing lane + for lane in self.lanes: + if lane.overflow and lane.overflow not in names: + raise ValueError( + f"Lane '{lane.name}' overflows to '{lane.overflow}' " + f"which does not exist. Available: {sorted(names)}" + ) + + # Self-overflow and cycle detection + for lane in self.lanes: + if lane.overflow == lane.name: + raise ValueError(f"Lane '{lane.name}' overflows to itself") + + overflow_map = {lane.name: lane.overflow for lane in self.lanes} + for start in overflow_map: + visited, current = {start}, overflow_map.get(start) + while current: + if current in visited: + raise ValueError( + f"Overflow cycle detected involving lane '{current}'" + ) + visited.add(current) + current = overflow_map.get(current) + + # Exactly one lane must be marked as default + defaults = [lane for lane in self.lanes if lane.default] + if len(defaults) != 1: + raise ValueError( + f"Exactly one lane must have default=true, found {len(defaults)}" + ) + + return self + + @property + def default_lane(self) -> str: + """The name of the default lane.""" + return next(lane.name for lane in self.lanes if lane.default) + + def get(self, name: str) -> LaneConfig: + """Look up a lane by name. Returns the default lane if unknown.""" + for lane in self.lanes: + if lane.name == name: + return lane + return next(lane for lane in self.lanes if lane.default) + + +def parse_lanes_config(json_str: str) -> LanesConfig: + """Parse and validate lanes config from a JSON string.""" + return LanesConfig(lanes=json.loads(json_str)) + + +def load_lanes_config(path: str = "lanes.json") -> LanesConfig: + """Load and validate lane definitions from a JSON file.""" + config_path = Path(path) + if not config_path.is_absolute(): + config_path = Path(__file__).parent.parent.parent / config_path + + with open(config_path) as f: + raw = json.load(f) + + return LanesConfig(lanes=raw) diff --git a/search-proxy/src/proxy/lanes.py b/search-proxy/src/proxy/lanes.py new file mode 100644 index 0000000000..31bafb3c4c --- /dev/null +++ b/search-proxy/src/proxy/lanes.py @@ -0,0 +1,135 @@ +import logging +import uuid +from contextlib import asynccontextmanager + +import redis.asyncio + +from proxy.config import LanesConfig + +logger = logging.getLogger(__name__) + + +class LoadSheddingError(Exception): + def __init__(self, lane_name: str, retry_after: int): + self.lane_name = lane_name + self.retry_after = retry_after + + +class RequestLanes: + """Redis-backed distributed traffic lanes. + + Each lane's active permits are tracked in a Redis sorted set keyed + lane:{name}:active. Each in-flight request occupies one member (its UUID); + the score is the expiry epoch from Redis server time. Entries whose score + has passed are pruned on the next acquire, so leaked permits from crashed + tasks recover automatically without manual intervention. + + When Redis is unavailable, acquire fails open: requests are forwarded + without a permit rather than 429'd. This degrades gracefully under + infrastructure failure at the cost of temporary over-capacity.""" + + def __init__( + self, + config: LanesConfig, + redis_client: redis.asyncio.Redis, + permit_ttl: int = 300, + ): + self.config = config + self.redis = redis_client + self.permit_ttl = permit_ttl + + def _lane_key(self, lane_name: str) -> str: + return f"lane:{lane_name}:active" + + async def _redis_now(self) -> float: + """Return current time from the Redis server to avoid ECS task clock skew.""" + seconds, microseconds = await self.redis.time() + return seconds + microseconds / 1_000_000 + + async def _try_acquire(self, lane_name: str, permits: int, request_id: str, redis_now: float) -> bool: + """Prune expired entries, count active, and add this request if under limit. + + Redis exceptions propagate to the caller, which decides fail-open vs fail-closed.""" + key = self._lane_key(lane_name) + await self.redis.zremrangebyscore(key, "-inf", redis_now) + count = await self.redis.zcard(key) + if count >= permits: + return False + await self.redis.zadd(key, {request_id: redis_now + self.permit_ttl}) + return True + + async def _release(self, lane_name: str, request_id: str) -> None: + """Remove this request's permit entry from the sorted set.""" + key = self._lane_key(lane_name) + try: + await self.redis.zrem(key, request_id) + except Exception: + # Log and swallow so a Redis failure here doesn't propagate out of + # the finally block and crash the ASGI handler. The permit leaks + # but will expire naturally via the TTL score. + logger.error( + "permit_release_failed", + extra={"lane": lane_name}, + exc_info=True, + ) + + async def _acquire_permit( + self, lane_name: str, load_shedding_enabled: bool, request_id: str + ) -> tuple[str, bool]: + """Try to acquire a permit, returning (lane_name, permit_stored). + + permit_stored is False when Redis is unavailable and the request is + allowed through without a permit (fail-open). The caller must not + attempt a release in that case.""" + lane = self.config.get(lane_name) + + try: + redis_now = await self._redis_now() + + if await self._try_acquire(lane.name, lane.permits, request_id, redis_now): + return lane.name, True + + if lane.overflow: + overflow_lane = self.config.get(lane.overflow) + if await self._try_acquire(overflow_lane.name, overflow_lane.permits, request_id, redis_now): + return overflow_lane.name, True + + if not load_shedding_enabled: + await self.redis.zadd(self._lane_key(lane.name), {request_id: redis_now + self.permit_ttl}) + logger.warning( + "load_shed_suppressed", + extra={"lane": lane.name, "load_shedding_enabled": False}, + ) + return lane.name, True + + raise LoadSheddingError(lane.name, lane.retry_after) + + except LoadSheddingError: + raise + except Exception: + logger.error( + "permit_acquire_failed", + extra={"lane": lane.name}, + exc_info=True, + ) + logger.warning( + "permit_bypassed", + extra={"lane": lane.name, "reason": "redis_unavailable"}, + ) + return lane.name, False + + @asynccontextmanager + async def acquire(self, lane_name: str, load_shedding_enabled: bool = True): + """Acquire a distributed permit for the named lane. + + Yields the name of the lane that was actually acquired (may differ + from the requested lane if overflow occurred). The permit is always + released when the context exits, even on exception. If Redis was + unavailable during acquire (fail-open), no release is attempted.""" + request_id = str(uuid.uuid4()) + actual_name, permit_stored = await self._acquire_permit(lane_name, load_shedding_enabled, request_id) + try: + yield actual_name + finally: + if permit_stored: + await self._release(actual_name, request_id) diff --git a/search-proxy/test/__init__.py b/search-proxy/test/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/search-proxy/test/conftest.py b/search-proxy/test/conftest.py new file mode 100644 index 0000000000..b8c10c66d6 --- /dev/null +++ b/search-proxy/test/conftest.py @@ -0,0 +1,14 @@ +import pytest + +from proxy.app import DEFAULT_TOGGLES + + +@pytest.fixture(autouse=True) +def proxy_env_vars(monkeypatch): + """Set required env vars for ProxySettings in all tests. + + These have no defaults — the deployment must provide them. Tests set + them here so ProxySettings can be instantiated without hitting + validation errors.""" + monkeypatch.setenv("CMR_PROXY_BACKEND_URL", "http://localhost:3003") + monkeypatch.setenv("CMR_PROXY_REDIS_URL", "redis://localhost:6379") diff --git a/search-proxy/test/test_app.py b/search-proxy/test/test_app.py new file mode 100644 index 0000000000..6e91662c90 --- /dev/null +++ b/search-proxy/test/test_app.py @@ -0,0 +1,599 @@ +import gzip +import json +import time +from unittest.mock import AsyncMock + +import fakeredis.aioredis +import httpx +import pytest + +from proxy.app import DEFAULT_TOGGLES, app, filter_hop_headers +from proxy.cache import ResponseCache +from proxy.config import LaneConfig, LanesConfig, ProxySettings +from proxy.lanes import RequestLanes + +BACKEND_JSON = json.dumps({"items": []}) +BACKEND_HEADERS = {"content-type": "application/json", "cmr-hits": "5"} + + +def make_lanes_config(**overrides): + defaults = { + "express_permits": 200, + "standard_permits": 150, + "heavy_permits": 50, + "express_cache_ttl": 10, + } + defaults.update(overrides) + return LanesConfig( + lanes=[ + LaneConfig( + name="express", + permits=defaults["express_permits"], + overflow="standard", + cache_ttl=defaults["express_cache_ttl"], + default=True, + ), + LaneConfig( + name="standard", + permits=defaults["standard_permits"], + cache_ttl=15, + ), + LaneConfig( + name="heavy", + permits=defaults["heavy_permits"], + cache_ttl=30, + retry_after=10, + ), + ] + ) + + +def make_backend_response( + status_code=200, + body=BACKEND_JSON, + headers=None, +): + """Create a mock httpx.Response.""" + headers = headers or BACKEND_HEADERS + return httpx.Response( + status_code=status_code, + content=body.encode(), + headers=headers, + request=httpx.Request("GET", "http://backend/test"), + ) + + +@pytest.fixture +async def client(): + """Test client with fakeredis and mocked backend.""" + fake_redis = fakeredis.aioredis.FakeRedis() + config = make_lanes_config() + + settings = ProxySettings() + app.state.settings = settings + app.state.redis = fake_redis + app.state.lanes_config = config + app.state.lanes = RequestLanes(config, fake_redis) + app.state.cache = ResponseCache(fake_redis, settings.max_cache_response_bytes) + app.state.toggles = dict(DEFAULT_TOGGLES) + app.state.backend = AsyncMock(spec=httpx.AsyncClient) + app.state.backend.get = AsyncMock(return_value=make_backend_response()) + app.state.backend.request = AsyncMock(return_value=make_backend_response()) + + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), + base_url="http://test", + ) as c: + yield c + + await fake_redis.aclose() + + +# Routing + + +class TestRouting: + async def test_search_request_is_proxied(self, client): + resp = await client.get("/search/granules.json?provider=POCLOUD") + assert resp.status_code == 200 + assert resp.json() == {"items": []} + + async def test_health_returns_cmr_format(self, client): + """Health check matches CMR's ok?/dependencies format.""" + app.state.backend.get.return_value = make_backend_response() + resp = await client.get("/health") + assert resp.status_code == 200 + data = resp.json() + assert data["ok?"] is True + assert "dependencies" in data + assert data["dependencies"]["redis"]["ok?"] is True + assert data["dependencies"]["search"]["ok?"] is True + + async def test_health_includes_lane_status(self, client): + """Health check includes per-lane utilization.""" + app.state.backend.get.return_value = make_backend_response() + resp = await client.get("/health") + data = resp.json() + deps = data["dependencies"] + assert "lane-express" in deps + assert "lane-standard" in deps + assert "lane-heavy" in deps + assert deps["lane-heavy"]["permits"] == 50 + assert deps["lane-heavy"]["active"] == 0 + assert deps["lane-heavy"]["ok?"] is True + + async def test_health_reports_lane_at_capacity(self, client): + """Health reports at_capacity but remains ok? true — informational only.""" + app.state.backend.get.return_value = make_backend_response() + fake_redis = app.state.redis + future = time.time() + 300 + await fake_redis.zadd("lane:heavy:active", {f"req-{i}": future for i in range(50)}) + resp = await client.get("/health") + data = resp.json() + assert resp.status_code == 200 + assert data["ok?"] is True + assert data["dependencies"]["lane-heavy"]["ok?"] is True + assert data["dependencies"]["lane-heavy"]["at_capacity"] is True + await fake_redis.delete("lane:heavy:active") + + async def test_health_shallow_always_200(self, client): + resp = await client.get("/health/shallow") + assert resp.status_code == 200 + assert resp.json()["ok?"] is True + + async def test_health_not_cached(self, client): + """/health is not cached — every call re-checks the backend.""" + app.state.backend.get.return_value = make_backend_response() + await client.get("/health") + await client.get("/health") + health_calls = [ + c + for c in app.state.backend.get.call_args_list + if "/search/health" in str(c) + ] + assert len(health_calls) == 2 + + +# Response cache + + +class TestResponseCache: + async def test_second_identical_request_hits_cache(self, client): + """Polling the same query should hit cache on repeat.""" + await client.get("/search/granules.json?provider=POCLOUD") + await client.get("/search/granules.json?provider=POCLOUD") + assert app.state.backend.get.call_count == 1 + + async def test_different_params_miss_cache(self, client): + await client.get("/search/granules.json?provider=POCLOUD") + await client.get("/search/granules.json?provider=LPDAAC") + assert app.state.backend.get.call_count == 2 + + async def test_different_auth_tokens_miss_cache(self, client): + await client.get( + "/search/granules.json?provider=POCLOUD", + headers={"Echo-Token": "user-a"}, + ) + await client.get( + "/search/granules.json?provider=POCLOUD", + headers={"Echo-Token": "user-b"}, + ) + assert app.state.backend.get.call_count == 2 + + async def test_error_responses_not_cached(self, client): + app.state.backend.get.return_value = make_backend_response( + status_code=500, body='{"errors": ["oops"]}' + ) + await client.get("/search/granules.json?provider=POCLOUD") + app.state.backend.get.return_value = make_backend_response() + resp = await client.get("/search/granules.json?provider=POCLOUD") + assert resp.status_code == 200 + assert app.state.backend.get.call_count == 2 + + async def test_no_cache_when_ttl_is_zero(self, client): + """Lanes with cache_ttl=0 should not cache.""" + app.state.lanes_config = make_lanes_config(express_cache_ttl=0) + + await client.get("/search/granules.json?concept_id=C123") + await client.get("/search/granules.json?concept_id=C123") + assert app.state.backend.get.call_count == 2 + + +# Classification integration + + +class TestClassificationIntegration: + async def test_express_request(self, client): + resp = await client.get("/search/collections.json?concept_id=C123") + assert resp.status_code == 200 + + async def test_heavy_request(self, client): + resp = await client.get("/search/granules.json?include_facets=v2") + assert resp.status_code == 200 + + +# Load shedding + + +class TestLoadShedding: + async def test_429_on_load_shedding(self, client): + fake_redis = app.state.redis + config = make_lanes_config( + express_permits=1, + standard_permits=1, + heavy_permits=1, + ) + app.state.lanes = RequestLanes(config, fake_redis) + + # Fill the heavy lane via Redis + await fake_redis.zadd("lane:heavy:active", {"blocking-req": time.time() + 300}) + resp = await client.get("/search/granules.json?include_facets=v2") + assert resp.status_code == 429 + assert "Retry-After" in resp.headers + assert "overloaded" in resp.json()["errors"][0] + + await fake_redis.delete("lane:heavy:active") + + +# Backend errors + + +class TestBackendErrors: + async def test_backend_timeout_returns_504(self, client): + app.state.backend.get.side_effect = httpx.TimeoutException("timed out") + resp = await client.get("/search/granules.json?provider=POCLOUD") + assert resp.status_code == 504 + assert "timed out" in resp.json()["errors"][0].lower() + assert "cmr-request-id" in resp.headers + + async def test_backend_connect_error_returns_502(self, client): + app.state.backend.get.side_effect = httpx.ConnectError("refused") + resp = await client.get("/search/granules.json?provider=POCLOUD") + assert resp.status_code == 502 + assert "unavailable" in resp.json()["errors"][0].lower() + + async def test_malformed_content_length_does_not_crash(self, client): + app.state.backend.request.return_value = make_backend_response() + resp = await client.post( + "/search/granules.json", + content="provider=POCLOUD", + headers={ + "content-type": "application/x-www-form-urlencoded", + "content-length": "not-a-number", + }, + ) + assert resp.status_code == 200 + + +# Forwarding + + +class TestForwarding: + async def test_preserves_backend_status_code(self, client): + app.state.backend.get.return_value = make_backend_response(status_code=404) + resp = await client.get("/search/granules.json?concept_id=MISSING") + assert resp.status_code == 404 + + async def test_post_request_forwarded(self, client): + app.state.backend.request.return_value = make_backend_response() + resp = await client.post( + "/search/granules.json", + content="provider=POCLOUD", + headers={"content-type": "application/x-www-form-urlencoded"}, + ) + assert resp.status_code == 200 + app.state.backend.request.assert_called_once() + + async def test_request_id_injected_to_backend(self, client): + await client.get( + "/search/granules.json?provider=X", + headers={"cmr-request-id": "trace-123"}, + ) + call_args = app.state.backend.get.call_args + forwarded_headers = call_args.kwargs.get( + "headers", call_args[1].get("headers", {}) + ) + assert forwarded_headers["cmr-request-id"] == "trace-123" + + async def test_generated_request_id_in_response(self, client): + resp = await client.get("/search/granules.json?provider=X") + assert resp.headers.get("cmr-request-id") + + async def test_query_string_not_double_encoded(self, client): + await client.get("/search/granules.json?polygon=1,2,3,4&provider=POCLOUD") + call_args = app.state.backend.get.call_args + forwarded_url = ( + str(call_args.args[0]) + if call_args.args + else str(call_args.kwargs.get("url", "")) + ) + assert "polygon=1,2,3,4" in forwarded_url + assert "provider=POCLOUD" in forwarded_url + assert "%3D" not in forwarded_url + assert "%26" not in forwarded_url + + async def test_hop_headers_filtered(self, client): + app.state.backend.get.return_value = make_backend_response( + headers={ + "content-type": "application/json", + "transfer-encoding": "chunked", + "connection": "keep-alive", + "cmr-hits": "10", + } + ) + resp = await client.get("/search/granules.json?provider=X") + assert "transfer-encoding" not in resp.headers + assert "connection" not in resp.headers + + async def test_content_encoding_stripped(self, client): + compressed = gzip.compress(BACKEND_JSON.encode()) + app.state.backend.get.return_value = httpx.Response( + status_code=200, + content=compressed, + headers={"content-type": "application/json", "content-encoding": "gzip", "cmr-hits": "5"}, + request=httpx.Request("GET", "http://backend/test"), + ) + resp = await client.get("/search/granules.json?provider=X") + assert "content-encoding" not in resp.headers + + async def test_content_length_not_forwarded_from_backend(self, client): + # Backend claims 746 bytes (e.g. compressed size) but actual body is smaller. + # Proxy strips the backend content-length; framework sets the correct value. + app.state.backend.get.return_value = make_backend_response( + headers={"content-type": "application/json", "content-length": "746", "cmr-hits": "5"} + ) + resp = await client.get("/search/granules.json?provider=X") + assert resp.headers.get("content-length") != "746" + + async def test_proxy_marker_header_injected(self, client): + await client.get("/search/granules.json?provider=X") + call_args = app.state.backend.get.call_args + forwarded_headers = call_args.kwargs.get( + "headers", call_args[1].get("headers", {}) + ) + assert forwarded_headers.get("x-cmr-proxy-request") == "1" + + +# Hop header filtering + + +class TestFilterHopHeaders: + def test_removes_hop_headers(self): + headers = httpx.Headers( + { + "content-type": "application/json", + "transfer-encoding": "chunked", + "connection": "keep-alive", + } + ) + filtered = filter_hop_headers(headers) + assert "content-type" in filtered + assert "transfer-encoding" not in filtered + assert "connection" not in filtered + + def test_preserves_cmr_headers(self): + headers = httpx.Headers( + { + "cmr-hits": "100", + "cmr-took": "45", + "content-type": "application/json", + } + ) + filtered = filter_hop_headers(headers) + assert filtered["cmr-hits"] == "100" + assert filtered["cmr-took"] == "45" + + +# Request ID + + +class TestRequestId: + async def test_request_id_propagated(self, client): + resp = await client.get( + "/search/granules.json?provider=X", + headers={"cmr-request-id": "my-req-123"}, + ) + assert resp.headers.get("cmr-request-id") == "my-req-123" + + +# POST body classification + + +class TestPostBodyClassification: + async def test_post_json_body_not_parsed(self, client): + app.state.backend.request.return_value = make_backend_response() + resp = await client.post( + "/search/granules.json", + content='{"provider": "POCLOUD"}', + headers={"content-type": "application/json"}, + ) + assert resp.status_code == 200 + + async def test_oversized_content_length_returns_413(self, client): + app.state.settings.max_request_body_bytes = 100 + resp = await client.post( + "/search/granules.json", + content="x" * 50, + headers={ + "content-type": "application/x-www-form-urlencoded", + "content-length": "999999", + }, + ) + assert resp.status_code == 413 + assert "too large" in resp.json()["errors"][0] + + async def test_oversized_body_returns_413(self, client): + app.state.settings.max_request_body_bytes = 100 + resp = await client.post( + "/search/granules.json", + content="x" * 200, + headers={"content-type": "application/x-www-form-urlencoded"}, + ) + assert resp.status_code == 413 + + async def test_non_utf8_post_body_returns_200(self, client): + app.state.backend.request.return_value = make_backend_response() + resp = await client.post( + "/search/granules.json", + content=b"\xff\xfe\x00\x01", + headers={"content-type": "application/x-www-form-urlencoded"}, + ) + assert resp.status_code == 200 + + async def test_non_utf8_post_body_still_forwarded(self, client): + app.state.backend.request.return_value = make_backend_response() + await client.post( + "/search/granules.json", + content=b"\xff\xfe\x00\x01", + headers={"content-type": "application/x-www-form-urlencoded"}, + ) + app.state.backend.request.assert_called_once() + + +# Feature toggles + + +class TestBypassToggle: + async def test_bypass_forwards_directly(self, client): + app.state.toggles["bypass_enabled"] = True + try: + resp = await client.get("/search/granules.json?provider=POCLOUD") + assert resp.status_code == 200 + app.state.backend.get.assert_called_once() + finally: + app.state.toggles["bypass_enabled"] = False + + async def test_bypass_timeout_returns_504(self, client): + app.state.toggles["bypass_enabled"] = True + app.state.backend.get.side_effect = httpx.TimeoutException("timed out") + try: + resp = await client.get("/search/granules.json?provider=POCLOUD") + assert resp.status_code == 504 + finally: + app.state.backend.get.side_effect = None + app.state.toggles["bypass_enabled"] = False + + async def test_bypass_connect_error_returns_502(self, client): + app.state.toggles["bypass_enabled"] = True + app.state.backend.get.side_effect = httpx.ConnectError("refused") + try: + resp = await client.get("/search/granules.json?provider=POCLOUD") + assert resp.status_code == 502 + finally: + app.state.backend.get.side_effect = None + app.state.toggles["bypass_enabled"] = False + + +class TestCacheToggle: + async def test_cache_disabled_skips_caching(self, client): + app.state.toggles["cache_enabled"] = False + try: + await client.get("/search/granules.json?provider=POCLOUD") + await client.get("/search/granules.json?provider=POCLOUD") + assert app.state.backend.get.call_count == 2 + finally: + app.state.toggles["cache_enabled"] = True + + +class TestLoadSheddingToggle: + async def test_load_shedding_disabled_allows_over_capacity(self, client): + app.state.toggles["load_shedding_enabled"] = False + fake_redis = app.state.redis + config = make_lanes_config(heavy_permits=1) + app.state.lanes = RequestLanes(config, fake_redis) + await fake_redis.zadd("lane:heavy:active", {"blocking-req": time.time() + 300}) + try: + resp = await client.get("/search/granules.json?include_facets=v2") + assert resp.status_code == 200 + finally: + await fake_redis.delete("lane:heavy:active") + app.state.toggles["load_shedding_enabled"] = True + app.state.lanes = RequestLanes(make_lanes_config(), fake_redis) + + +class TestClassificationToggle: + async def test_classification_disabled_uses_default_lane(self, client): + """With classification off, a normally-heavy request routes to express.""" + app.state.toggles["classification_enabled"] = False + fake_redis = app.state.redis + await fake_redis.set("lane:heavy:active", 50) + try: + resp = await client.get("/search/granules.json?include_facets=v2") + assert resp.status_code == 200 + finally: + await fake_redis.delete("lane:heavy:active") + app.state.toggles["classification_enabled"] = True + + +# Cache key segmentation + + +class TestCacheKeySegmentation: + async def test_search_after_header_creates_separate_cache_entry(self, client): + await client.get("/search/granules.json?provider=POCLOUD") + await client.get( + "/search/granules.json?provider=POCLOUD", + headers={"cmr-search-after": "cursor-xyz"}, + ) + assert app.state.backend.get.call_count == 2 + + async def test_same_search_after_hits_cache(self, client): + await client.get( + "/search/granules.json?provider=POCLOUD", + headers={"cmr-search-after": "cursor-xyz"}, + ) + await client.get( + "/search/granules.json?provider=POCLOUD", + headers={"cmr-search-after": "cursor-xyz"}, + ) + assert app.state.backend.get.call_count == 1 + + async def test_different_accept_creates_separate_cache_entry(self, client): + await client.get( + "/search/granules.json?provider=POCLOUD", + headers={"accept": "application/json"}, + ) + await client.get( + "/search/granules.json?provider=POCLOUD", + headers={"accept": "application/xml"}, + ) + assert app.state.backend.get.call_count == 2 + + async def test_post_different_bodies_create_separate_cache_entries(self, client): + app.state.backend.request.return_value = make_backend_response() + await client.post( + "/search/granules.json", + content="provider=POCLOUD", + headers={"content-type": "application/x-www-form-urlencoded"}, + ) + await client.post( + "/search/granules.json", + content="provider=LPDAAC", + headers={"content-type": "application/x-www-form-urlencoded"}, + ) + assert app.state.backend.request.call_count == 2 + + async def test_post_same_body_hits_cache(self, client): + app.state.backend.request.return_value = make_backend_response() + await client.post( + "/search/granules.json", + content="provider=POCLOUD", + headers={"content-type": "application/x-www-form-urlencoded"}, + ) + await client.post( + "/search/granules.json", + content="provider=POCLOUD", + headers={"content-type": "application/x-www-form-urlencoded"}, + ) + assert app.state.backend.request.call_count == 1 + + async def test_same_accept_hits_cache(self, client): + await client.get( + "/search/granules.json?provider=POCLOUD", + headers={"accept": "application/json"}, + ) + await client.get( + "/search/granules.json?provider=POCLOUD", + headers={"accept": "application/json"}, + ) + assert app.state.backend.get.call_count == 1 diff --git a/search-proxy/test/test_cache.py b/search-proxy/test/test_cache.py new file mode 100644 index 0000000000..2a0ccde6e5 --- /dev/null +++ b/search-proxy/test/test_cache.py @@ -0,0 +1,149 @@ +import fakeredis.aioredis +import pytest + +from proxy.cache import ResponseCache + +SAMPLE_RESPONSE = { + "status_code": 200, + "body": '{"items": []}', + "headers": {"Content-Type": "application/json"}, +} + + +@pytest.fixture +async def cache(): + client = fakeredis.aioredis.FakeRedis() + c = ResponseCache(client, max_response_bytes=1_048_576) + yield c + await client.aclose() + + +class TestCacheKey: + async def test_same_request_same_key(self, cache): + await cache.set( + "GET", + "/search/granules.json", + "p=1", + "token", + SAMPLE_RESPONSE, + 100, + 30, + ) + result = await cache.get("GET", "/search/granules.json", "p=1", "token") + assert result == SAMPLE_RESPONSE + + async def test_different_method_misses(self, cache): + await cache.set( + "GET", + "/search/granules.json", + "p=1", + "token", + SAMPLE_RESPONSE, + 100, + 30, + ) + result = await cache.get("POST", "/search/granules.json", "p=1", "token") + assert result is None + + async def test_different_path_misses(self, cache): + await cache.set( + "GET", + "/search/granules.json", + "p=1", + "token", + SAMPLE_RESPONSE, + 100, + 30, + ) + result = await cache.get("GET", "/search/collections.json", "p=1", "token") + assert result is None + + async def test_different_auth_misses(self, cache): + await cache.set( + "GET", + "/search/granules.json", + "p=1", + "token-a", + SAMPLE_RESPONSE, + 100, + 30, + ) + result = await cache.get("GET", "/search/granules.json", "p=1", "token-b") + assert result is None + + async def test_oversized_response_not_cached(self, cache): + await cache.set( + "GET", + "/search/granules.json", + "p=1", + "token", + SAMPLE_RESPONSE, + 2_000_000, + 30, + ) + result = await cache.get("GET", "/search/granules.json", "p=1", "token") + assert result is None + + async def test_miss_returns_none(self, cache): + result = await cache.get("GET", "/search/granules.json", "p=1", "token") + assert result is None + + async def test_different_search_after_misses(self, cache): + await cache.set( + "GET", "/search/granules.json", "p=1", "token", + SAMPLE_RESPONSE, 100, 30, search_after="cursor-page1", + ) + result = await cache.get( + "GET", "/search/granules.json", "p=1", "token", search_after="cursor-page2" + ) + assert result is None + + async def test_same_search_after_hits(self, cache): + await cache.set( + "GET", "/search/granules.json", "p=1", "token", + SAMPLE_RESPONSE, 100, 30, search_after="cursor-abc", + ) + result = await cache.get( + "GET", "/search/granules.json", "p=1", "token", search_after="cursor-abc" + ) + assert result == SAMPLE_RESPONSE + + async def test_different_accept_misses(self, cache): + await cache.set( + "GET", "/search/granules.json", "p=1", "token", + SAMPLE_RESPONSE, 100, 30, accept="application/json", + ) + result = await cache.get( + "GET", "/search/granules.json", "p=1", "token", accept="application/xml" + ) + assert result is None + + async def test_same_accept_hits(self, cache): + await cache.set( + "GET", "/search/granules.json", "p=1", "token", + SAMPLE_RESPONSE, 100, 30, accept="application/json", + ) + result = await cache.get( + "GET", "/search/granules.json", "p=1", "token", accept="application/json" + ) + assert result == SAMPLE_RESPONSE + + async def test_different_body_hash_misses(self, cache): + await cache.set( + "POST", "/search/granules.json", "", "token", + SAMPLE_RESPONSE, 100, 30, body_hash="aaaa1111", + ) + result = await cache.get( + "POST", "/search/granules.json", "", "token", body_hash="bbbb2222" + ) + assert result is None + + async def test_same_body_hash_hits(self, cache): + await cache.set( + "POST", "/search/granules.json", "", "token", + SAMPLE_RESPONSE, 100, 30, body_hash="aaaa1111", + ) + result = await cache.get( + "POST", "/search/granules.json", "", "token", body_hash="aaaa1111" + ) + assert result == SAMPLE_RESPONSE diff --git a/search-proxy/test/test_classifier.py b/search-proxy/test/test_classifier.py new file mode 100644 index 0000000000..1a49e419a0 --- /dev/null +++ b/search-proxy/test/test_classifier.py @@ -0,0 +1,279 @@ +from proxy.classifier import ( + EXPRESS, + HEAVY, + STANDARD, + _compute_bbox_area, + _count_polygon_vertices, + classify_request, +) + +# Non-spatial heavy signals + + +class TestHeavyParamSignals: + def test_include_facets(self): + assert classify_request({"include_facets": "v2"}) == HEAVY + + def test_online_only(self): + assert classify_request({"online_only": "true"}) == HEAVY + + def test_cloud_cover(self): + assert classify_request({"cloud_cover": "0,50"}) == HEAVY + + +class TestHeavyPatternSignals: + def test_temporal_facet(self): + assert classify_request({"temporal_facet[0][year]": "2024"}) == HEAVY + + def test_cycle(self): + assert classify_request({"cycle[]": "1"}) == HEAVY + + def test_passes_pass(self): + assert classify_request({"passes[0][pass]": "100"}) == HEAVY + + def test_options_readable_granule_name_pattern(self): + assert ( + classify_request({"options[readable_granule_name][pattern]": "true"}) + == HEAVY + ) + + def test_options_granule_ur_pattern(self): + assert classify_request({"options[granule_ur][pattern]": "true"}) == HEAVY + + def test_options_producer_granule_id_pattern(self): + assert classify_request({"options[producer_granule_id][pattern]": "true"}) == HEAVY + + +# Shapefile detection + + +class TestShapefileDetection: + def test_shapefile_with_multipart_content_type(self): + assert ( + classify_request( + {"shapefile": "data"}, + content_type="multipart/form-data; boundary=----", + ) + == HEAVY + ) + + def test_file_param_with_multipart_content_type(self): + assert ( + classify_request( + {"file": "data"}, + content_type="multipart/form-data; boundary=----", + ) + == HEAVY + ) + + def test_shapefile_without_multipart_is_not_heavy(self): + # Without multipart content-type, shapefile param alone doesn't trigger heavy + assert classify_request({"shapefile": "data"}) == EXPRESS + + def test_multipart_without_shapefile_param_is_not_heavy(self): + assert ( + classify_request( + {"concept_id": "C123"}, + content_type="multipart/form-data; boundary=----", + ) + == EXPRESS + ) + + +# Non-spatial standard signals + + +class TestStandardParamSignals: + def test_temporal(self): + assert classify_request({"temporal": "2020-01-01T00:00:00Z,"}) == STANDARD + + def test_temporal_array(self): + assert classify_request({"temporal[]": "2020-01-01T00:00:00Z,"}) == STANDARD + + def test_updated_since(self): + assert classify_request({"updated_since": "2024-01-01"}) == STANDARD + + def test_revision_date(self): + assert classify_request({"revision_date": "2024-01-01,"}) == STANDARD + + def test_orbit_number(self): + assert classify_request({"orbit_number": "1000"}) == STANDARD + + +# Express default + + +class TestExpressDefault: + def test_no_signals_returns_express(self): + assert classify_request({"concept_id": "C1234"}) == EXPRESS + + def test_empty_params_returns_express(self): + assert classify_request({}) == EXPRESS + + def test_simple_filter_params(self): + assert ( + classify_request( + { + "provider": "POCLOUD", + "tag_key": "gov.nasa.eosdis", + "entry_id": "X", + } + ) + == EXPRESS + ) + + +# Spatial: polygon + + +class TestSpatialPolygon: + def test_polygon_array_always_heavy(self): + assert classify_request({"polygon[]": "1,2,3,4,5,6,1,2"}) == HEAVY + + def test_polygon_20_vertices_is_standard(self): + # 20 vertices = 40 coordinate values + coords = ",".join(str(i) for i in range(40)) + assert classify_request({"polygon": coords}) == STANDARD + + def test_polygon_21_vertices_is_heavy(self): + # 21 vertices = 42 coordinate values + coords = ",".join(str(i) for i in range(42)) + assert classify_request({"polygon": coords}) == HEAVY + + def test_polygon_small_is_standard(self): + # Simple triangle (3 vertices = 6 coords) + assert classify_request({"polygon": "0,0,1,0,0,1"}) == STANDARD + + +# Spatial: bounding box + + +class TestSpatialBoundingBox: + def test_bbox_area_under_threshold_is_standard(self): + # 99 * 50 = 4950 sq deg (just under 5000) + assert classify_request({"bounding_box": "-50,-25,49,25"}) == STANDARD + + def test_bbox_area_over_threshold_is_heavy(self): + # 101 * 50 = 5050 sq deg (just over 5000) + assert classify_request({"bounding_box": "-50,-25,51,25"}) == HEAVY + + def test_bbox_array_2_values_is_standard(self): + # 2 small bboxes -> standard (not exceeding count threshold) + assert ( + classify_request( + { + "bounding_box[]": ["0,0,1,1", "2,2,3,3"], + } + ) + == STANDARD + ) + + def test_bbox_array_3_values_is_heavy(self): + # 3 bboxes -> heavy (exceeds count threshold of 2) + assert ( + classify_request( + { + "bounding_box[]": ["0,0,1,1", "2,2,3,3", "4,4,5,5"], + } + ) + == HEAVY + ) + + def test_bbox_antimeridian_crossing_is_conservatively_heavy(self): + # Naive math: abs(-170 - 170) * abs(10 - -10) = 340 * 20 = 6800 sq deg. + # Real area is small, but overestimating toward HEAVY is the safe direction. + assert classify_request({"bounding_box": "170,-10,-170,10"}) == HEAVY + + +# Spatial: circle + + +class TestSpatialCircle: + def test_circle_array_is_express(self): + assert classify_request({"circle[]": "40,-90,100"}) == EXPRESS + + def test_single_circle_is_standard(self): + assert classify_request({"circle": "40,-90,100"}) == STANDARD + + +# Spatial: point + + +class TestSpatialPoint: + def test_point_is_standard(self): + assert classify_request({"point": "40,-90"}) == STANDARD + + def test_point_array_is_standard(self): + assert classify_request({"point[]": "40,-90"}) == STANDARD + + +# Combined signals + + +class TestCombinedSignals: + def test_heavy_nonspatial_overrides_spatial_standard(self): + # include_facets is a heavy signal; polygon with few vertices would be standard + assert ( + classify_request( + { + "polygon": "0,0,1,0,0,1", + "include_facets": "v2", + } + ) + == HEAVY + ) + + def test_no_spatial_no_signals_is_express(self): + assert classify_request({"provider": "POCLOUD"}) == EXPRESS + + def test_standard_temporal_with_simple_spatial(self): + # point (standard spatial) + temporal (standard non-spatial) -> standard + assert ( + classify_request( + { + "point": "40,-90", + "temporal": "2020-01-01,", + } + ) + == STANDARD + ) + + +# Helper unit tests + + +class TestCountPolygonVertices: + def test_triangle(self): + assert _count_polygon_vertices("0,0,1,0,0,1") == 3 + + def test_with_whitespace(self): + assert _count_polygon_vertices(" 0,0,1,0,0,1 ") == 3 + + def test_20_vertices(self): + coords = ",".join(str(i) for i in range(40)) + assert _count_polygon_vertices(coords) == 20 + + def test_21_vertices(self): + coords = ",".join(str(i) for i in range(42)) + assert _count_polygon_vertices(coords) == 21 + + +class TestComputeBboxArea: + def test_simple_box(self): + # 10 wide, 10 tall = 100 + assert _compute_bbox_area("0,0,10,10") == 100.0 + + def test_antimeridian_crossing_overestimates(self): + # Naive: abs(-170 - 170) * abs(10 - -10) = 340 * 20 = 6800 + # Overestimates, which is the conservative direction for classification. + assert _compute_bbox_area("170,-10,-170,10") == 6800.0 + + def test_invalid_value_returns_above_threshold(self): + # Malformed bbox should be treated conservatively (pushed toward heavy) + assert _compute_bbox_area("invalid") > 5000 + + def test_empty_string_returns_above_threshold(self): + assert _compute_bbox_area("") > 5000 + + def test_malformed_bbox_classifies_as_heavy(self): + assert classify_request({"bounding_box": "garbage"}) == HEAVY diff --git a/search-proxy/test/test_config.py b/search-proxy/test/test_config.py new file mode 100644 index 0000000000..bf736993dc --- /dev/null +++ b/search-proxy/test/test_config.py @@ -0,0 +1,388 @@ +import json +import tempfile + +import pytest +from pydantic import ValidationError + +from proxy.config import LaneConfig, LanesConfig, ProxySettings, load_lanes_config, parse_lanes_config + + +class TestProxySettingsDefaults: + def test_backend_url_required(self, monkeypatch): + """backend_url has no default — must come from CMR_PROXY_BACKEND_URL.""" + monkeypatch.delenv("CMR_PROXY_BACKEND_URL") + with pytest.raises(ValidationError, match="backend_url"): + ProxySettings() + + def test_redis_url_required(self, monkeypatch): + """redis_url has no default — must come from CMR_PROXY_REDIS_URL.""" + monkeypatch.delenv("CMR_PROXY_REDIS_URL") + with pytest.raises(ValidationError, match="redis_url"): + ProxySettings() + + def test_backend_url_from_env(self): + """backend_url is read from the env var set by conftest.""" + s = ProxySettings() + assert s.backend_url == "http://localhost:3003" + + def test_redis_url_from_env(self): + """redis_url is read from the env var set by conftest.""" + s = ProxySettings() + assert s.redis_url == "redis://localhost:6379" + + def test_max_request_body_bytes(self): + s = ProxySettings() + assert s.max_request_body_bytes == 52_428_800 + + def test_backend_timeout_seconds(self): + s = ProxySettings() + assert s.backend_timeout_seconds == 300.0 + + def test_redis_socket_connect_timeout(self): + s = ProxySettings() + assert s.redis_socket_connect_timeout == 2.0 + + def test_redis_socket_timeout(self): + s = ProxySettings() + assert s.redis_socket_timeout == 2.0 + + def test_redis_health_check_interval(self): + s = ProxySettings() + assert s.redis_health_check_interval == 30 + + def test_redis_max_connections_default_is_none(self): + s = ProxySettings() + assert s.redis_max_connections is None + + def test_backend_max_connections(self): + s = ProxySettings() + assert s.backend_max_connections == 500 + + def test_backend_max_keepalive(self): + s = ProxySettings() + assert s.backend_max_keepalive == 200 + + def test_lanes_config_default_path(self): + s = ProxySettings() + assert s.lanes_config == "lanes.json" + + +class TestEnvVarOverrides: + def test_proxy_backend_url_override(self, monkeypatch): + monkeypatch.setenv("CMR_PROXY_BACKEND_URL", "http://search:9999") + s = ProxySettings() + assert s.backend_url == "http://search:9999" + + def test_proxy_redis_url_override(self, monkeypatch): + monkeypatch.setenv("CMR_PROXY_REDIS_URL", "redis://redis-cluster:6380") + s = ProxySettings() + assert s.redis_url == "redis://redis-cluster:6380" + + def test_proxy_backend_timeout_override(self, monkeypatch): + monkeypatch.setenv("CMR_PROXY_BACKEND_TIMEOUT_SECONDS", "60.0") + s = ProxySettings() + assert s.backend_timeout_seconds == 60.0 + + def test_proxy_lanes_config_override(self, monkeypatch): + monkeypatch.setenv("CMR_PROXY_LANES_CONFIG", "/etc/cmr/lanes.json") + s = ProxySettings() + assert s.lanes_config == "/etc/cmr/lanes.json" + + def test_proxy_lanes_json_default_is_none(self): + s = ProxySettings() + assert s.lanes_json is None + + def test_proxy_lanes_json_override(self, monkeypatch): + monkeypatch.setenv("CMR_PROXY_LANES_JSON", '[{"name":"x","permits":1,"default":true}]') + s = ProxySettings() + assert s.lanes_json is not None + config = parse_lanes_config(s.lanes_json) + assert config.default_lane == "x" + assert config.get("x").permits == 1 + + +# Lane config model + + +class TestLaneConfigValidation: + def test_blank_name_raises(self): + with pytest.raises(ValidationError, match="blank"): + LaneConfig(name=" ", permits=10) + + def test_empty_name_raises(self): + with pytest.raises(ValidationError, match="blank"): + LaneConfig(name="", permits=10) + + def test_permits_zero_raises(self): + with pytest.raises(ValidationError, match="permits"): + LaneConfig(name="x", permits=0) + + def test_permits_negative_raises(self): + with pytest.raises(ValidationError, match="permits"): + LaneConfig(name="x", permits=-1) + + def test_cache_ttl_negative_raises(self): + with pytest.raises(ValidationError, match="cache_ttl"): + LaneConfig(name="x", permits=10, cache_ttl=-1) + + def test_retry_after_zero_raises(self): + with pytest.raises(ValidationError, match="retry_after"): + LaneConfig(name="x", permits=10, retry_after=0) + + def test_retry_after_negative_raises(self): + with pytest.raises(ValidationError, match="retry_after"): + LaneConfig(name="x", permits=10, retry_after=-5) + + def test_cache_ttl_zero_is_valid(self): + lane = LaneConfig(name="x", permits=10, cache_ttl=0) + assert lane.cache_ttl == 0 + + +class TestLaneConfigModel: + def test_minimal_lane(self): + lane = LaneConfig(name="test", permits=10) + assert lane.name == "test" + assert lane.permits == 10 + assert lane.overflow is None + assert lane.retry_after == 5 + assert lane.default is False + + def test_full_lane(self): + lane = LaneConfig( + name="express", + permits=200, + overflow="standard", + retry_after=3, + default=True, + ) + assert lane.overflow == "standard" + assert lane.retry_after == 3 + assert lane.default is True + + +# Lanes config validation + + +class TestLanesConfigValidation: + def test_valid_three_lane_config(self): + config = LanesConfig( + lanes=[ + LaneConfig( + name="express", permits=200, overflow="standard", default=True + ), + LaneConfig(name="standard", permits=150), + LaneConfig(name="heavy", permits=50), + ] + ) + assert config.default_lane == "express" + assert len(config.lanes) == 3 + + def test_duplicate_lane_names_raises(self): + with pytest.raises(ValidationError, match="Duplicate"): + LanesConfig( + lanes=[ + LaneConfig(name="express", permits=200, default=True), + LaneConfig(name="express", permits=100), + ] + ) + + def test_overflow_to_nonexistent_lane_fails(self): + with pytest.raises(ValidationError, match="does not exist"): + LanesConfig( + lanes=[ + LaneConfig( + name="fast", permits=10, overflow="missing", default=True + ), + ] + ) + + def test_no_default_fails(self): + with pytest.raises(ValidationError, match="default=true"): + LanesConfig( + lanes=[ + LaneConfig(name="a", permits=10), + LaneConfig(name="b", permits=10), + ] + ) + + def test_multiple_defaults_fails(self): + with pytest.raises(ValidationError, match="default=true"): + LanesConfig( + lanes=[ + LaneConfig(name="a", permits=10, default=True), + LaneConfig(name="b", permits=10, default=True), + ] + ) + + def test_get_existing_lane(self): + config = LanesConfig( + lanes=[ + LaneConfig(name="fast", permits=10, default=True), + LaneConfig(name="slow", permits=5), + ] + ) + lane = config.get("slow") + assert lane.name == "slow" + assert lane.permits == 5 + + def test_self_overflow_raises(self): + with pytest.raises(ValidationError, match="itself"): + LanesConfig( + lanes=[ + LaneConfig(name="express", permits=200, overflow="express", default=True), + ] + ) + + def test_overflow_cycle_raises(self): + with pytest.raises(ValidationError, match="cycle"): + LanesConfig( + lanes=[ + LaneConfig(name="a", permits=10, overflow="b", default=True), + LaneConfig(name="b", permits=10, overflow="a"), + ] + ) + + def test_three_node_cycle_raises(self): + with pytest.raises(ValidationError, match="cycle"): + LanesConfig( + lanes=[ + LaneConfig(name="a", permits=10, overflow="b", default=True), + LaneConfig(name="b", permits=10, overflow="c"), + LaneConfig(name="c", permits=10, overflow="a"), + ] + ) + + def test_rho_shape_cycle_raises(self): + """A→B→C→B: cycle not involving the start node.""" + with pytest.raises(ValidationError, match="cycle"): + LanesConfig( + lanes=[ + LaneConfig(name="a", permits=10, overflow="b", default=True), + LaneConfig(name="b", permits=10, overflow="c"), + LaneConfig(name="c", permits=10, overflow="b"), + ] + ) + + def test_get_unknown_lane_returns_default(self): + config = LanesConfig( + lanes=[ + LaneConfig(name="fast", permits=10, default=True), + LaneConfig(name="slow", permits=5), + ] + ) + lane = config.get("nonexistent") + assert lane.name == "fast" + + +# Loading from JSON file + + +class TestLoadLanesConfig: + def test_loads_default_lanes_json(self): + """The lanes.json in the project root should load successfully.""" + config = load_lanes_config("lanes.json") + assert config.default_lane == "express" + assert len(config.lanes) == 3 + + # Verify the values match what's in the file + express = config.get("express") + assert express.permits == 200 + assert express.overflow == "standard" + + heavy = config.get("heavy") + assert heavy.permits == 50 + assert heavy.retry_after == 10 + + def test_loads_from_absolute_path(self): + """Verify loading from an absolute path works.""" + lanes = [ + {"name": "only", "permits": 100, "default": True}, + ] + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: + json.dump(lanes, f) + f.flush() + config = load_lanes_config(f.name) + + assert config.default_lane == "only" + assert config.get("only").permits == 100 + + def test_four_lane_custom_config(self): + """Verify a custom 4-lane config loads and validates.""" + lanes = [ + { + "name": "fast", + "permits": 300, + "overflow": "normal", + "default": True, + }, + {"name": "normal", "permits": 200}, + {"name": "slow", "permits": 100}, + {"name": "bulk", "permits": 20, "retry_after": 30}, + ] + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: + json.dump(lanes, f) + f.flush() + config = load_lanes_config(f.name) + + assert len(config.lanes) == 4 + assert config.get("bulk").retry_after == 30 + assert config.get("slow").permits == 100 + + +# Parsing from a JSON string + + +VALID_LANES_JSON = json.dumps([ + {"name": "express", "permits": 200, "overflow": "standard", "cache_ttl": 10, "retry_after": 5, "default": True}, + {"name": "standard", "permits": 150, "cache_ttl": 15, "retry_after": 5}, + {"name": "heavy", "permits": 50, "cache_ttl": 30, "retry_after": 10}, +]) + + +class TestParseLanesConfig: + def test_parses_valid_json(self): + config = parse_lanes_config(VALID_LANES_JSON) + assert len(config.lanes) == 3 + assert config.default_lane == "express" + + def test_parsed_values_match_input(self): + config = parse_lanes_config(VALID_LANES_JSON) + assert config.get("heavy").permits == 50 + assert config.get("heavy").retry_after == 10 + + def test_invalid_json_raises(self): + with pytest.raises(json.JSONDecodeError): + parse_lanes_config("not-valid-json{{") + + def test_invalid_schema_raises(self): + with pytest.raises(ValidationError): + parse_lanes_config(json.dumps([{"name": "broken"}])) + + def test_empty_string_raises(self): + with pytest.raises(json.JSONDecodeError): + parse_lanes_config("") + + +class TestLanesConfigSourceSelection: + """Verify the is-not-None semantics that drive the lifespan branch selection.""" + + def test_lanes_json_set_is_not_none(self, monkeypatch): + """When CMR_PROXY_LANES_JSON is set, settings.lanes_json is not None and parse_lanes_config succeeds.""" + monkeypatch.setenv("CMR_PROXY_LANES_JSON", VALID_LANES_JSON) + settings = ProxySettings() + assert settings.lanes_json is not None + config = parse_lanes_config(settings.lanes_json) + assert config.default_lane == "express" + + def test_empty_lanes_json_is_not_none(self, monkeypatch): + """An empty string is not None — lifespan calls parse_lanes_config, which raises rather than silently falling back to the file.""" + monkeypatch.setenv("CMR_PROXY_LANES_JSON", "") + settings = ProxySettings() + assert settings.lanes_json is not None + with pytest.raises(json.JSONDecodeError): + parse_lanes_config(settings.lanes_json) + + def test_unset_lanes_json_is_none(self): + """When CMR_PROXY_LANES_JSON is not set, lanes_json is None and the file path is used.""" + settings = ProxySettings() + assert settings.lanes_json is None diff --git a/search-proxy/test/test_lanes.py b/search-proxy/test/test_lanes.py new file mode 100644 index 0000000000..ff8f81f072 --- /dev/null +++ b/search-proxy/test/test_lanes.py @@ -0,0 +1,322 @@ +import time +from unittest.mock import AsyncMock, patch + +import fakeredis.aioredis +import pytest +from redis.exceptions import MaxConnectionsError + +from proxy.config import LaneConfig, LanesConfig +from proxy.lanes import LoadSheddingError, RequestLanes + + +def make_config(**overrides): + """Build a LanesConfig with sensible defaults, overridable per lane.""" + defaults = { + "express_permits": 200, + "standard_permits": 150, + "heavy_permits": 50, + } + defaults.update(overrides) + return LanesConfig( + lanes=[ + LaneConfig( + name="express", + permits=defaults["express_permits"], + overflow="standard", + default=True, + ), + LaneConfig( + name="standard", + permits=defaults["standard_permits"], + ), + LaneConfig( + name="heavy", + permits=defaults["heavy_permits"], + ), + ] + ) + + +@pytest.fixture +async def redis_client(): + client = fakeredis.aioredis.FakeRedis() + yield client + await client.aclose() + + +@pytest.fixture +async def lanes(redis_client): + return RequestLanes(make_config(), redis_client) + + +@pytest.fixture +async def tight_lanes(redis_client): + """Lanes with 1 permit each for testing contention.""" + return RequestLanes( + make_config( + express_permits=1, + standard_permits=1, + heavy_permits=1, + ), + redis_client, + ) + + +class TestBasicAcquisition: + async def test_express_acquires_and_releases(self, lanes): + async with lanes.acquire("express") as actual: + assert actual == "express" + + async def test_standard_acquires_and_releases(self, lanes): + async with lanes.acquire("standard") as actual: + assert actual == "standard" + + async def test_heavy_acquires_and_releases(self, lanes): + async with lanes.acquire("heavy") as actual: + assert actual == "heavy" + + +class TestPermitRelease: + async def test_permit_released_after_normal_exit(self, tight_lanes): + async with tight_lanes.acquire("heavy"): + pass + async with tight_lanes.acquire("heavy") as actual: + assert actual == "heavy" + + async def test_permit_released_after_exception(self, tight_lanes): + with pytest.raises(ValueError): + async with tight_lanes.acquire("heavy"): + raise ValueError("boom") + async with tight_lanes.acquire("heavy") as actual: + assert actual == "heavy" + + +class TestDistributedCounting: + async def test_permits_shared_across_instances(self, redis_client): + """Two RequestLanes sharing the same Redis enforce a global limit.""" + config = make_config(heavy_permits=2) + lanes_a = RequestLanes(config, redis_client) + lanes_b = RequestLanes(config, redis_client) + + async with lanes_a.acquire("heavy") as a: + async with lanes_b.acquire("heavy") as b: + assert a == "heavy" + assert b == "heavy" + + with pytest.raises(LoadSheddingError): + async with lanes_a.acquire("heavy"): + pass + + async def test_active_count_is_zero_after_all_releases(self, redis_client): + """After all permits are released, the sorted set should be empty.""" + config = make_config(heavy_permits=5) + lanes = RequestLanes(config, redis_client) + + for _ in range(5): + async with lanes.acquire("heavy"): + pass + + count = await redis_client.zcard("lane:heavy:active") + assert count == 0 + + +class TestExpressOverflow: + async def test_express_overflows_to_standard(self, redis_client): + """When express is full, express requests overflow to standard.""" + config = make_config(express_permits=1, standard_permits=1) + lanes = RequestLanes(config, redis_client) + + async with lanes.acquire("express") as first: + assert first == "express" + async with lanes.acquire("express") as second: + assert second == "standard" + + async def test_express_sheds_when_both_full(self, redis_client): + """When both express and standard are full, express sheds.""" + config = make_config(express_permits=1, standard_permits=1) + lanes = RequestLanes(config, redis_client) + + async with lanes.acquire("express"): + async with lanes.acquire("standard"): + with pytest.raises(LoadSheddingError) as exc_info: + async with lanes.acquire("express"): + pass + assert exc_info.value.lane_name == "express" + + +class TestLoadShedding: + async def test_sheds_when_full(self, tight_lanes): + async with tight_lanes.acquire("heavy"): + with pytest.raises(LoadSheddingError) as exc_info: + async with tight_lanes.acquire("heavy"): + pass + assert exc_info.value.lane_name == "heavy" + + async def test_load_shedding_error_has_retry_after(self, tight_lanes): + async with tight_lanes.acquire("heavy"): + with pytest.raises(LoadSheddingError) as exc_info: + async with tight_lanes.acquire("heavy"): + pass + assert exc_info.value.retry_after == 5 + + +class TestLaneIsolation: + async def test_tiers_are_independent(self, tight_lanes): + """Filling one tier doesn't affect others.""" + async with tight_lanes.acquire("heavy"): + async with tight_lanes.acquire("standard") as actual: + assert actual == "standard" + + async def test_concurrent_acquisition(self, lanes): + """Multiple tiers can be held simultaneously.""" + async with lanes.acquire("express") as t1: + async with lanes.acquire("standard") as t2: + async with lanes.acquire("heavy") as t3: + assert t1 == "express" + assert t2 == "standard" + assert t3 == "heavy" + + +class TestUnknownLaneFallback: + async def test_unknown_lane_falls_back_to_default(self, lanes): + async with lanes.acquire("nonexistent") as actual: + assert actual == "express" + + +class TestRedisFailOpen: + async def test_acquire_succeeds_when_redis_unavailable(self, lanes): + """When Redis is down, acquire fails open rather than 429ing the request.""" + with patch.object(lanes.redis, "time", new=AsyncMock(side_effect=Exception("Redis down"))): + async with lanes.acquire("heavy") as actual: + assert actual == "heavy" + + async def test_no_release_attempted_on_fail_open(self, lanes): + """When Redis fails during acquire, no release is attempted (nothing was stored).""" + with patch.object(lanes.redis, "time", new=AsyncMock(side_effect=Exception("Redis down"))): + with patch.object(lanes.redis, "zrem", new=AsyncMock()) as mock_zrem: + async with lanes.acquire("heavy"): + pass + mock_zrem.assert_not_called() + + async def test_fail_open_does_not_propagate_exception(self, lanes): + """When Redis is unavailable during acquire, then no exception surfaces to the caller.""" + with patch.object(lanes.redis, "time", new=AsyncMock(side_effect=Exception("Redis down"))): + try: + async with lanes.acquire("heavy"): + pass + except Exception: + pytest.fail("Redis exception escaped from acquire") + + +class TestReleaseFailure: + async def test_release_redis_error_does_not_propagate(self, lanes): + """MaxConnectionsError in _release must not escape to the caller.""" + with patch.object(lanes.redis, "zrem", new=AsyncMock(side_effect=MaxConnectionsError("Too many connections"))): + try: + async with lanes.acquire("heavy"): + pass + except MaxConnectionsError: + pytest.fail("MaxConnectionsError escaped from _release") + + async def test_release_redis_error_is_logged(self, lanes, caplog): + import logging + with patch.object(lanes.redis, "zrem", new=AsyncMock(side_effect=MaxConnectionsError("Too many connections"))): + with caplog.at_level(logging.ERROR, logger="proxy.lanes"): + async with lanes.acquire("heavy"): + pass + assert "permit_release_failed" in caplog.text + + +class TestLoadSheddingDisabled: + async def test_force_acquires_when_lane_full(self, tight_lanes): + """With load shedding off, over-capacity requests are not rejected.""" + async with tight_lanes.acquire("heavy"): + async with tight_lanes.acquire("heavy", load_shedding_enabled=False) as actual: + assert actual == "heavy" + + async def test_does_not_raise_load_shedding_error(self, tight_lanes): + async with tight_lanes.acquire("heavy"): + try: + async with tight_lanes.acquire("heavy", load_shedding_enabled=False): + pass + except LoadSheddingError: + pytest.fail("LoadSheddingError raised with load_shedding_enabled=False") + + +class TestTTLExpiry: + async def test_expired_permit_is_pruned_on_next_acquire(self, redis_client): + """When a permit's TTL score has passed, then it is pruned on the next acquire and the slot is freed.""" + config = make_config(heavy_permits=1) + lanes = RequestLanes(config, redis_client) + + # Simulate a leaked permit: score is in the past so it is already expired + expired_score = time.time() - 10 + await redis_client.zadd("lane:heavy:active", {"leaked-request-id": expired_score}) + + # The expired entry should be pruned and the slot freed for a new request + async with lanes.acquire("heavy") as actual: + assert actual == "heavy" + + async def test_non_expired_permit_blocks_acquire(self, redis_client): + """When a permit's TTL score is in the future, then it is counted as active and blocks acquisition.""" + config = make_config(heavy_permits=1) + lanes = RequestLanes(config, redis_client) + + future_score = time.time() + 300 + await redis_client.zadd("lane:heavy:active", {"active-request-id": future_score}) + + with pytest.raises(LoadSheddingError): + async with lanes.acquire("heavy"): + pass + + async def test_multiple_expired_permits_all_pruned(self, redis_client): + """When multiple permits have expired scores, then all are pruned before counting capacity.""" + config = make_config(heavy_permits=2) + lanes = RequestLanes(config, redis_client) + + expired = time.time() - 10 + await redis_client.zadd("lane:heavy:active", { + "leaked-1": expired, + "leaked-2": expired, + "leaked-3": expired, + }) + + # All three expired entries pruned — both permits now free + async with lanes.acquire("heavy"): + async with lanes.acquire("heavy") as actual: + assert actual == "heavy" + + +class TestUniquePermitSlots: + async def test_identical_concurrent_requests_each_occupy_a_slot(self, redis_client): + """Two concurrent requests occupy two distinct slots in the sorted set.""" + config = make_config(express_permits=2) + lanes = RequestLanes(config, redis_client) + + async with lanes.acquire("express"): + async with lanes.acquire("express"): + count = await redis_client.zcard("lane:express:active") + assert count == 2 + + async def test_third_request_shed_when_two_permit_lane_full(self, redis_client): + """A third request is shed when a 2-permit no-overflow lane is fully occupied.""" + config = make_config(heavy_permits=2) + lanes = RequestLanes(config, redis_client) + + async with lanes.acquire("heavy"): + async with lanes.acquire("heavy"): + with pytest.raises(LoadSheddingError): + async with lanes.acquire("heavy"): + pass + + async def test_permit_slot_removed_on_release(self, redis_client): + """Each release removes exactly one entry from the sorted set.""" + config = make_config(heavy_permits=3) + lanes = RequestLanes(config, redis_client) + + async with lanes.acquire("heavy"): + async with lanes.acquire("heavy"): + count_during = await redis_client.zcard("lane:heavy:active") + assert count_during == 2 + count_after_one_release = await redis_client.zcard("lane:heavy:active") + assert count_after_one_release == 1