feat(asyncio): Make all calls to postgress using sqlalchemy[asyncio] - #1001
feat(asyncio): Make all calls to postgress using sqlalchemy[asyncio]#1001brightsparc wants to merge 8 commits into
Conversation
|
@brightsparc sorry for the slow reply! Working on it |
njbrake
left a comment
njbrake
left a comment
There was a problem hiding this comment.
Review by @njbrake (assisted by Claude)
Hey @brightsparc, thanks for this. The diagnosis is spot-on and the benchmarks speak for themselves (391 -> 7,637 requests is compelling). The async migration across ~30 files is mechanically solid and the pluggable strategy design is clean.
That said, there are several issues that need addressing before we can merge. Splitting into must-fix and should-fix:
Must Fix
1. Test suite is broken: log_usage() calls missing log_writer parameter
log_usage now requires log_writer: LogWriter as its second positional argument, but 11 call sites across 6 test files were not updated:
tests/gateway/test_log_usage_commit_scope.py(4 calls)tests/gateway/test_atomic_spend_update.py(2 calls)tests/gateway/test_transaction_safety.py(2 calls)tests/gateway/test_pricing_config.py(2 calls)tests/gateway/test_timezone_consistency.py(1 call)
These will all raise TypeError at runtime.
2. Pool config values are dead code
GatewayConfig defines db_pool_size, db_max_overflow, db_pool_timeout, db_pool_recycle but init_db(database_url, auto_migrate) never receives them. They're never passed to create_async_engine. Either wire them through or remove them from this PR until they are.
3. CAS race condition: stale spend after losing reset race
In _validate_cas, after _claim_reset_cas returns False (another request won the race), the local user object still has the pre-reset spend value (expire_on_commit=False). _enforce_budget_limit then checks the stale value and can false-reject.
Scenario: user at spend=9.50 / max_budget=10.00, reset is overdue. Request A wins the CAS (spend -> 0.0). Request B loses but still sees spend=9.50, gets a 403 even though the budget was just reset.
Fix: await db.refresh(user) after losing the CAS race.
Relatedly, previous_spend captured for BudgetResetLog comes from the same stale read, so audit records may be inaccurate under concurrent load.
4. Unrelated changes need to be split out
The branch includes changes unrelated to the async migration:
- Removal of dict-based
response_formathandling from gemini and xai providers + their tests - Removal of the "Mypy and Provider SDKs" section from AGENTS.md
- Pre-commit config downgrades (ruff v0.15.8 -> v0.15.4, codespell v2.4.2 -> v2.4.1)
docs/package-lock.jsonversion downgrade
These look like they came from the branch being forked from an older main. Please rebase onto current main or revert these changes so the PR only contains the async migration.
5. Docs reference nonexistent split strategy
In budget-management.md:
If you need user-blocked enforcement without budget checks, use
splitorcas...
There is no split strategy. Only for_update, cas, and disabled exist in code.
6. Large benchmark artifacts should not be committed
tests/load/results/ contains ~4,600 lines of CSV, JSON, and TXT files. These are point-in-time snapshots that will bloat the repo and never be usefully diffed. Please keep results/results.md (the human summary is great) but gitignore the raw data files.
Should Fix
7. health_readiness manual async generator management
The readiness endpoint manually advances get_db() with double anext calls and nested try blocks instead of using Depends(get_db) like every other endpoint. This works but is fragile. Consider using dependency injection here too, or at minimum await db_gen.aclose() for cleanup.
8. Strategy config values not validated
budget_strategy="foobar" silently falls through to for_update. Same for log_writer_strategy. A log warning on unrecognized values would prevent invisible misconfigurations.
Overall the architecture is the right call. Looking forward to landing this once the above is addressed.
cf2991c to
ad85bc3
Compare
|
I’ve updated this from the base. Let me know if there are any more changes you would like to see before we can land @njbrake |
GitHub Actions
Error
GitHub Actions
Inactive
3 participants
Description
Fixes the gateway lockup diagnosed in #1000 by removing the sync-psycopg2-in-async-handler anti-pattern at its root, and adds two pluggable strategies for common bottlenecks on the request path.
Four stacked changes, each behind a config flag:
psycopg2+sqlalchemy.orm.Sessionforasyncpg+AsyncSession. Contended DB calls now yield to the event loop instead of blocking it.budget_strategyconfig —for_update(default, legacy) |cas(lock-free, recommended) |disabled(escape hatch)log_writer_strategyconfig —single(default, inline) |batch(recommended for high-throughput)db_pool_size/db_max_overflow/db_pool_timeout/db_pool_recycleexposed to config, defaults bumped for async workloads.Plus a k6-based load-test harness under
tests/load/with recorded comparison results.Why
The gateway uses
async defroute handlers but (until this PR) every DB call underneath went through synchronouspsycopg2. When a query had to wait on anything — aFOR UPDATErow lock, a saturated pool, Postgres itself — the sync call blocked the entire event loop, freezing every in-flight request including/health. The load test onmaindemonstrates this clearly: under 100 VUs the gateway stalls at 228 requests completed (all on the warmup), the 15-connection pool saturates, and the process needs to be killed after piling up 30s-timeout errors.After this change, a contended DB call
awaits and other coroutines make progress. Fixes the architectural root cause rather than papering over it.Changes
Dependencies (
pyproject.toml)asyncpg>=0.29.0(async Postgres driver)aiosqlite>=0.19.0(async SQLite for local dev/tests)sqlalchemy[asyncio]extrapsycopg2-binaryfor Alembic (which remains sync)Database core (
core/database.py)create_engine→create_async_enginewith configurable pool (pool_size, max_overflow, pool_timeout, pool_recycle)sessionmaker→async_sessionmaker(expire_on_commit=False)get_dbis now anasyncgenerator yieldingAsyncSessionpostgresql://…/sqlite:///…and wires the async driver automatically; Alembic gets the sync formcreate_session()helper for startup-time use outside request scopeConfig (
core/config.py)budget_strategy(default"for_update") — how per-user budget validation serializes concurrent resetslog_writer_strategy(default"single") — how usage log rows are persisteddb_pool_size(default 10),db_max_overflow(default 20),db_pool_timeout(default 30.0),db_pool_recycle(default -1)Budget service (
services/budget_service.py) — dispatch viamatch/case:for_update: legacy behavior (FOR UPDATE across entire request)cas: lock-free — atomic conditional UPDATE (WHERE next_budget_reset_at < now) claims the reset. Compare-and-swap semantics protect theBudgetResetLoginsert from duplicates.disabled: skipvalidate_user_budgetentirely, return NoneLog writers (
services/log_writer.py) — new module with Protocol + 2 implementations:SingleLogWriter: write each event inline, one txn per eventBatchLogWriter: queue events ontoasyncio.Queue, background task flushes batches of up to 100 rows / 1s.queue.join()drains on shutdown via the FastAPI lifespan.start()/stop()and are driven from the lifespan. Both treat flush failures as best-effort (rollback + log + drop), consistent with the historical contract.Startup (
main.py)lifespanasync context managerAuth / deps (
api/deps.py)_verify_and_update_api_key,verify_api_key,verify_api_key_or_master_keyasync end-to-endget_log_writerFastAPI dependencyRepositories + services + routes (~20 files)
db.query(X).filter(...).first()→(await db.execute(select(X).where(...))).scalar_one_or_none(),db.commit/rollback/refresh/delete→await db.commit/rollback/refresh/delete, atomic spend update →await db.execute(update(User)...values(spend=User.spend + cost))log_usageinchat.pyrefactored to build aUsageLogand enqueue viaLogWriter.put()— no longer commits inlineMetrics (
metrics.py)gateway_usage_log_queue_depth(gauge)gateway_usage_log_batch_size(histogram, labeled by writer)gateway_usage_log_flush_duration_seconds(histogram, labeled by writer + result)gateway_usage_log_rows(counter, labeled by writer + result=written|dropped)Tests (
tests/gateway/conftest.py+ ~10 test files)test_db,db_session,client.override_get_dball yieldAsyncSessioncasanddisabledbudget strategies (no-FOR-UPDATE assertions viaget_active_userspy)Load test harness (
tests/load/)fake_provider.py— click CLI, seedable log-normal jittered delay, zero-delay noop modeload_test.js— k6 scenarios: warmup, distinct_users, same_user withsetup()/teardown()lifecyclerun_load_test.sh— orchestrator (Postgres container, fake provider, gateway, k6) + post-run DB verification that countsusage_logsrows to validate 100% coverageresults/— recorded sync-vs-async comparison across all 4 budget strategies + the batch log writerBenchmark results
100 VUs × 30s × noop fake provider, single uvicorn worker:
Key findings:
same_usercontention cliff (34.8 → 40.8 req/sec)same_userbeatsdistinct_usersin the batch run because batch groups spend UPDATEs per userFull results and per-transition analysis:
tests/load/results/results.md.Not in scope
postgresql://URL driven bypsycopg2-binary.for_updatefor backwards compatibility. Users opt intocasvia config.singlefor backwards compatibility. Users opt intobatchvia config.Possible follow-ups
psycopg2-binaryby switching Alembic topsycopg[binary](psycopg3 sync mode). Consolidates to one driver family.--workersto match production deployment (currently 1 for benchmarking signal isolation).find_model_pricingis a PK lookup per request.PR Type
validate_user_budgetholdsSELECT FOR UPDATEacross the entire LLM call #1000)Relevant issues
Fixes #1000.
Checklist
AI Usage Information
casstrategy, built the log writer abstraction, and wrote the k6 load test harness. The lockup diagnosis was done live in a running container by inspectingpg_stat_activity/pg_blocking_pids. All benchmarks and results are reproducible via./tests/load/run_load_test.sh.When answering questions by the reviewer, please respond yourself, do not copy/paste the reviewer comments into an AI system and paste back its answer. We want to discuss with you, not your AI :)