Skip to content

perf: keep hot-path DynamoDB calls off the event loop#108

Open
warren830 wants to merge 2 commits into
aws-samples:mainfrom
warren830:feat/hot-path-ddb-optimization
Open

perf: keep hot-path DynamoDB calls off the event loop#108
warren830 wants to merge 2 commits into
aws-samples:mainfrom
warren830:feat/hot-path-ddb-optimization

Conversation

@warren830

Copy link
Copy Markdown
Contributor

Summary

Every request currently makes 3+ synchronous boto3 DynamoDB calls directly on the asyncio event loop: API key validation (auth middleware), model mapping lookup, and usage recording (including inside streaming paths). Each call freezes the event loop for the full round trip (2–10ms), which:

  • caps per-worker throughput at roughly 1s / total_blocking_time regardless of Bedrock latency or connection count
  • stalls chunk forwarding for all concurrent SSE streams in the worker every time any request does an auth lookup or usage write
  • costs one DynamoDB read per request for data that changes rarely (keys, mappings)

Performance

  • API key validation results are cached in-process (API_KEY_CACHE_TTL_SECONDS, default 60s; invalid keys only 5s so brute-force spam can't pin entries and new keys work almost immediately). Cache misses run on a dedicated 4-thread pool — not the default executor, which this codebase already shares with multi-second Tavily/web-fetch/docker work — with single-flight coalescing so a hot key's TTL expiry triggers exactly one DynamoDB read.
  • Usage recording moved to a single background writer thread via UsageTracker.record_usage_nowait() (all 13 call sites). Streaming responses no longer stall on put_item.
  • Model mapping lookups cached (MODEL_MAPPING_CACHE_TTL_SECONDS, default 300s), including negative results — clients passing Bedrock ARNs directly (pass-through) previously paid a wasted DynamoDB read per request.

Correctness / hardening (from a 5-perspective adversarial review of the first commit; see below)

  • APIKeyManager.validate_api_key now re-raises ClientError instead of returning None: previously DynamoDB throttling was indistinguishable from "invalid key", which (with negative caching) would have locked valid keys out in rolling 5s windows during a throttling event. Lookup failures now return 401 without caching.
  • TTLCache overflow drops the soonest-expiring entries instead of clearing everything, so attacker-inserted short-TTL negative entries evict themselves rather than hot positive entries.
  • Auth cache is keyed by SHA-256 digest (no plaintext keys as dict keys, constant per-entry key size); values are deep-copied both ways so a handler mutating nested api_key_info fields can't poison the cache.
  • Usage writer: bounded backlog (drops + logs + usage_writes_dropped_total counter instead of unbounded memory growth if DynamoDB hangs), submit-after-shutdown guard (a usage write inside except/finally can never mask the original exception), kwargs validated at the call site (a typo raises TypeError instead of becoming silently dropped billing data), timestamps stamped at submit time with millisecond precision (same-second rows for one key previously overwrote each other via put_item; a draining backlog made that worse), and a drain_usage_writes() flush hook wired into app lifespan shutdown.

Review

The first commit was reviewed by five independent passes (security, performance, testing, maintainability, adversarial). All CRITICAL and high-confidence findings are fixed in the second commit. Deferred items, documented deliberately:

  • get_mapping cache misses still run synchronously on the loop in the converter path (anthropic_to_bedrock._convert_model_id). The cache makes misses rare for legitimate traffic; a full fix requires async-ifying the converter chain and is left as a follow-up.
  • Cross-process key revocation (admin portal) takes effect within the cache TTL (≤60s by default, documented in config.py and env.example). There is no cross-process invalidation channel; operators needing faster revocation can lower or disable the TTL.
  • The SIGTERM handler in app/main.py calls sys.exit(0) directly, bypassing uvicorn's graceful shutdown — pre-existing behavior, out of scope here, but it limits how much the new shutdown drain hook can help on deploys.
  • boto3 Table resources are shared across threads. Table operations delegate to the underlying thread-safe client and writes are serialized on one thread; noted because boto3 documents resources as not thread-safe.

Test coverage

  • 41 new unit tests across tests/unit/test_ttl_cache.py, test_auth_cache.py, test_usage_nowait.py, test_model_mapping_cache.py, covering TTL boundaries (exactly-at-expiry), negative caching, eviction under pressure, single-flight coalescing, dedicated-executor placement, off-loop execution, error paths (ClientError, shutdown executor, full backlog), and cache invalidation on set/delete.
  • New shared tests/conftest.py clears the class-level model-mapping cache between tests (the review caught that the shared cache made existing moto-based tests order-dependent).
  • Full suite: 355 passed. tests/unit/test_cdk_deploy_env_validation.py::test_deploy_requires_bedrock_api_key_when_openai_compat_enabled fails identically on a clean main checkout (environment-dependent, unrelated to this change).

Config

Two new settings (documented in env.example):

API_KEY_CACHE_TTL_SECONDS=60        # 0 disables
MODEL_MAPPING_CACHE_TTL_SECONDS=300 # 0 disables

Both default to safe values; setting them to 0 restores the previous per-request read behavior exactly.

🤖 Generated with Claude Code

warren830 added 2 commits July 5, 2026 23:13
- Cache API key validation in-process (TTL 60s, negative 5s) and run
  DynamoDB misses via asyncio.to_thread so auth never blocks the loop
- Record usage on a single background writer thread (record_usage_nowait)
  so streaming responses no longer stall on put_item
- Cache model mapping lookups (TTL 300s) including pass-through misses
- Add shared TTLCache utility (bounded, thread-safe); new config flags
  API_KEY_CACHE_TTL_SECONDS / MODEL_MAPPING_CACHE_TTL_SECONDS
- validate_api_key re-raises ClientError so DynamoDB throttling is a
  lookup failure (401, uncached), never a negative-cached 'invalid key'
  that locks valid keys out in rolling 5s windows
- TTLCache overflow now drops soonest-expiring entries instead of
  clearing everything: negative-cache spam evicts itself, not hot keys
- Auth cache keyed by SHA-256 digest (no plaintext keys as dict keys,
  bounded per-entry size), values deep-copied both ways, lookups run on
  a dedicated 4-thread pool with single-flight coalescing per key
- Usage writer: bounded backlog (drop + warn + prometheus counter),
  submit-after-shutdown guard, kwargs validated at the call site,
  timestamps stamped at submit with ms precision, drain hook wired
  into app lifespan shutdown
- Shared tests/conftest.py clears the model-mapping cache between tests
  (class-level cache made existing tests order-dependent)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant