feat: Add load test from any-llm PR #1001, fix race condition

[brightsparc]

Description

Fixes a row-loss bug in BatchLogWriter during graceful shutdown, and adds the missing tests/load/ harness that was dropped when the gateway was split out of any-llm (mozilla-ai/any-llm#1001).

The bug

BatchLogWriter.stop() used to call task.cancel() on the background run loop. If the cancel landed while _flush() was mid-commit, the entire in-flight batch was lost: _collect_batch calls queue.task_done() the moment items are dequeued — before the flush commits — so a CancelledError interrupting await db.commit() rolls the session back and leaves the rows in neither the queue nor the database. _flush_all() in stop() only drains what's still in the queue, so those items are gone.

The load test caught it: k6 reported 8,518 iterations but usage_logs held 8,418 rows — exactly 100 missing, equal to max_batch. One in-flight batch dropped on shutdown. In production the same race silently drops up to max_batch rows on every clean SIGTERM.

The fix: graceful shutdown via asyncio.Event

Instead of cancelling the task from under the flush, stop() now signals an asyncio.Event, awaits the task with a bounded _STOP_TIMEOUT (10s), and only falls back to cancel() if the timeout expires (a genuinely wedged flush).

The run loop checks the event between batches, lets any in-flight _flush() complete, then drains any remaining queue contents via _flush_all() in a finally block. _collect_batch races the queue get against stop_event.wait() using asyncio.wait(FIRST_COMPLETED) so an idle collector wakes immediately on shutdown instead of waiting out the 1s flush interval.

asyncio.Event (rather than a plain bool) because _collect_batch needs to await the shutdown signal concurrently with the queue get — one primitive covers both the loop-top check and the collector race. asyncio.shield(_flush) was considered and rejected: the outer await still raises CancelledError, so stop() would return while an orphaned flush continued into a half-torn-down event loop.

Tradeoff: stop() can now take up to 10s in the wedged-flush case; previously it returned instantly but silently dropped the batch. Durability over teardown speed, timeout caps worst case. Idle-writer shutdown is actually faster than before because the asyncio.wait race wakes immediately instead of waiting for _collect_batch's wait_for to time out.

Full trace and rationale in bug.md.

Missing load-test harness

mozilla-ai/any-llm#1001 split the gateway to this repo but missed tests/load/** (6 files) from introspection-org:julian/async-asyncpg. Restored here, with one small fix to run_load_test.sh: dropped uv run --extra gateway in two places — that extra was an any-llm concept; in this repo the gateway is the project itself, so plain uv run works.

Benchmark results

Same k6 setup as the prior reference run (GATEWAY_PORT=4001 BUDGET_STRATEGY=disabled LOG_WRITER_STRATEGY=batch, 100 VUs × 30s per scenario + 5s warmup, noop fake provider):

Metric	Before fix	After fix
k6 iterations	8,518	8,577
Rows in usage_logs	8,418	8,577
Coverage	98% (100 dropped)	100%
distinct_users rps	54.7	55.0
same_user rps	53.2	53.7
distinct p50 / p95 / p99	778 / 1243 / 1649 ms	750 / 1259 / 1688 ms
same p50 / p95 / p99	768 / 1281 / 1598 ms	763 / 1240 / 1702 ms
Failures	0	0
CPU avg / max	89.0% / 101%	90.0% / 101%
RSS max	430 MB	394 MB

No throughput or latency regression — all changes are within noise. The meaningful delta is 100% row coverage, zero drops.

For comparison against the original historical reference run from any-llm#1001 (different machine, but same k6 config): 97.4 total rps / 7,637 requests / 100%. This branch: ~108 total rps / 8,577 requests / 100%. Faster hardware this time, but the durability guarantee matches.

Changes

src/gateway/services/log_writer.py

• BatchLogWriter.__init__: adds self._stop_event = asyncio.Event() and _STOP_TIMEOUT = 10.0.
• BatchLogWriter.stop: sets the event, await asyncio.wait_for(self._task, _STOP_TIMEOUT), cancels only on timeout.
• BatchLogWriter._run: while not self._stop_event.is_set() loop; try/finally ensures _flush_all() always runs on exit.
• BatchLogWriter._collect_batch: asyncio.wait({queue.get(), stop_event.wait()}, FIRST_COMPLETED) replaces the old asyncio.wait_for(queue.get()).
• BatchLogWriter._flush_all: simplified single-loop drain.
• SingleLogWriter, NoopLogWriter, LogWriter Protocol all unchanged. Lifespan hook in main.py unchanged — stop() signature is still parameter-free.

tests/unit/test_log_writer.py — 4 new tests around the batch writer shutdown contract:

• test_batch_writer_flushes_queued_items_on_stop — baseline.
• test_batch_writer_does_not_drop_in_flight_batch_on_stop — the regression test. Uses a gated fake _flush to block a batch mid-commit, fires stop() while blocked, releases the gate, asserts every item is flushed (including ones queued after stop_event was set).
• test_batch_writer_stop_times_out_and_cancels — verifies _STOP_TIMEOUT fallback fires on a wedged flush.
• test_batch_writer_stop_is_idempotent_when_not_started — stop() on a never-started writer is a no-op.

tests/load/ — restored from introspection-org:julian/async-asyncpg (6 files: README.md, fake_provider.py, gateway-config.yml, load_test.js, run_load_test.sh, results/results.md). Minor fix to run_load_test.sh to drop --extra gateway for this repo's layout.

PR Type

• 🐛 Bug Fix (in-flight batch loss on BatchLogWriter shutdown)
• 🚦 Infrastructure (restore missing tests/load/ harness)

Relevant issues

Checklist

• I understand the code I am submitting.
• I have added unit tests that prove my fix/feature works
• I have run this code locally and verified it fixes the issue.
• New and existing tests pass locally
• Documentation was updated where necessary
• I have read and followed the contribution guidelines
• AI Usage:
  • No AI was used.
  • AI was used for drafting/refactoring.
  • This is fully AI-generated.

AI Usage Information

• AI Model used: Claude Opus 4.6
• AI Developer Tool used: Claude Code
• Any other info you'd like to share: AI diagnosed the race from the load-test coverage mismatch (100 missing rows == max_batch), traced it to the cancel-mid-_flush path, designed the asyncio.Event-based graceful shutdown, wrote the regression tests, and produced the comparison benchmarks. All results are reproducible via GATEWAY_PORT=4001 BUDGET_STRATEGY=disabled LOG_WRITER_STRATEGY=batch ./tests/load/run_load_test.sh async-batch-disabled.

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 :)

• I am an AI Agent filling out this form (check box if true)

[tbille]

@brightsparc could you review the PR. There are few extra markdown files that might not be required in here and that contain mentions of branches etc.

Could you also add a few details about the version of k6 you are using. I will play around with it and add a github workflow as a follow up PR.

Also can you rebase on main?

[njbrake]

Hello, closing out do to staleness, feel free to re-open if the issue continues!

[brightsparc]

Sorry, missed this one, can revisit - it was helpful in finding the asyncio bottleneck issues, i'll create a new PR when I have a moment.

[njbrake]

@brightsparc no worries, thanks for reaching out. Yes we're getting back into showing otari some more love and happy to help collaborate if there's any way we can help out here 🙏

[brightsparc]

@brightsparc no worries, thanks for reaching out. Yes we're getting back into showing otari some more love and happy to help collaborate if there's any way we can help out here 🙏

One thing that would be nice is if we could support a google interactions api such that we could preserve the thought / reasoning effort, I see this is in the base provider at the moment
https://github.com/mozilla-ai/any-llm/blob/main/src/any_llm/providers/gemini/base.py#L137-L142

Wonder how this works with the new api: https://ai.google.dev/api/interactions-api ?