akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(proxy/metrics): move the savings-ledger append off the event loop (#2439)

Browse source 
## Description

`PrometheusMetrics.record_request` appends one durable JSONL event per
compressed request. That append is synchronous: `open` + `fcntl.flock` +
`write`, plus a full-file rewrite once the ledger passes 1 MB. It runs
on the event loop, inside `self._lock`.

`export()` takes that same lock and holds it for the entire Prometheus
serialization, so a slow ledger write stops `/metrics` cold. In a repro
run of 200 compressed requests, `/metrics` completed zero scrapes and
the event loop never yielded once across 6.4 seconds.

The append now runs in a thread, outside the lock. `savings_ledger`
already takes its own `flock` across processes, so the metrics lock was
never what made the write safe.

Both halves are one change. Awaiting inside the lock would hold it for
the whole write rather than just the syscall, which is worse than what
is on main today.

The file already documents this hazard against itself.
`record_stage_timings` (`prometheus_metrics.py:867-874`) picks a plain
`threading.Lock` over `self._lock` specifically because "the async lock
is also held by `export()` during Prometheus scrapes." The ledger append
was the pattern that docstring warns about.

No filed issue for this one.

## Type of Change

- [ ] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [x] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- Move the `savings_ledger.record_savings_event` call in
`record_request` out of `async with self._lock` and run it through
`asyncio.to_thread`. The call site keeps its keyword arguments verbatim;
`to_thread` forwards `**kwargs`, so no `functools.partial` wrapper is
needed.
- Keep the `await`. Callers still see the event on disk when
`record_request` returns, which
`tests/test_savings_ledger_before_forwarded.py` asserts synchronously.
- Add `tests/test_savings_ledger_offload.py`: lock scope, event-loop
responsiveness, durability on return, and both arms of the `tokens_saved
> 0 and not stateless` gate.

`savings_ledger.py` is untouched. It stays synchronous so the MCP
`headroom_compress` caller in `ccr/mcp_server.py:789` does not have to
change.

Sizing the executor is left alone on purpose. `asyncio.to_thread` uses
the default pool, which is the documented tool for blocking I/O and
already the idiom here (`helpers.py:1297`, `server.py:1694`, `:3557`,
`:3613`, `:4244`). The compression pools are sized `max(1,
os.cpu_count())` for CPU-bound work, and `PrometheusMetrics` holds no
reference to `HeadroomProxy` anyway, so reaching them would mean a new
constructor parameter.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [x] Manual testing performed

### Test Output

```text
$ pytest tests/test_savings_ledger_offload.py tests/test_savings_ledger.py tests/test_savings_ledger_before_forwarded.py -q
======================== 26 passed, 1 warning in 4.08s =========================

$ ruff check . && ruff format --check headroom/proxy/prometheus_metrics.py tests/test_savings_ledger_offload.py
All checks passed!
2 files already formatted

$ mypy headroom/proxy/prometheus_metrics.py
Success: no issues found in 1 source file
```

Broader sweep across the blast radius, 145 test files matching savings /
metrics / outcome / stats / proxy / handler / server / ledger / cost /
prometheus, each run under a per-file wall-clock watchdog:

```text
138 files pass, 1470 tests passed
7 non-green:
  HANG tests/test_agent_savings.py
  HANG tests/test_ccr_mcp_server.py
  HANG tests/test_netcost_gate.py
  HANG tests/test_proxy_compress_endpoint.py
  HANG tests/test_proxy_mode_benchmark.py
  HANG tests/test_read_maturation_handler_nobust.py
  FAIL tests/test_proxy_copilot_auth_hooks.py::test_openai_passthrough_applies_copilot_auth

Same 7 files re-run with headroom/proxy/prometheus_metrics.py reverted to c400f908:
  identical set, identical failure. diff of the two non-green lists is empty.
```

The before and after non-green sets match exactly, so nothing here is a
regression from this PR. See Additional Notes for the hang.

## Real Behavior Proof

- Environment: macOS 15.4 (Darwin 25.4.0) arm64, Python 3.13.13,
uv-managed venv, git worktree at `upstream/main` `c400f908`.
- Exact command / steps: a standalone asyncio script, not the unit
tests. It builds a real `PrometheusMetrics` (no injected tracker, so it
self-constructs with `save_flush_every=PROXY_SAVINGS_FLUSH_EVERY`
exactly as the proxy does) against a real on-disk ledger pre-seeded to
3.00 MB so `_maybe_compact`'s full-file rewrite actually fires. It then
drives 200 `record_request` calls at concurrency 16 while a `/metrics`
scraper calls `export()` every 20 ms and a canary coroutine ticks every
5 ms. Ran twice from the same script: once with
`headroom/proxy/prometheus_metrics.py` reverted to `c400f908`, once with
this change. Seeding the ledger past 1 MB is the part that matters. On a
fresh ledger the write is microseconds, compaction never fires, and the
run shows no delta at all.
- Observed result: before, `/metrics` completed 0 scrapes and the canary
ticked once in 6419 ms. After, 206 scrapes at p50 0.1 ms and max 0.2 ms,
and 418 canary ticks with a 92.3 ms worst gap. Total wall clock barely
moved, 6419 ms to 6473 ms, which is the expected result and not a null
one: the same disk work still serializes on the ledger's own `flock`,
now in a thread instead of on the loop. Unit-test view of the same
behavior, with a 500 ms stub standing in for the write: before, `event
loop stalled 0.506s during a 0.500s ledger write` and the competing lock
holder waited `+0.502s`; after, both pass.
- Not tested: Windows, where `savings_ledger` already skips locking
because `fcntl` is unavailable. Multi-process contention on one ledger
file, which this change does not alter. The residual 92.3 ms loop gap
after the fix, which traces to `SavingsTracker._save_locked`'s
`os.fsync` (`savings_tracker.py:1445`) firing every 25th request from
inside the same lock, a separate path this PR leaves alone.

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I did **not** edit `CHANGELOG.md` — it is generated by
release-please from my Conventional Commit PR title (a CI guard enforces
this)

## Screenshots (if applicable)

N/A, no user-visible surface.

## Additional Notes

Docs checklist item is N/A. Nothing user-facing moves; `headroom
savings` reads the same ledger, with the same contents, written from a
thread.

Two related in-lock costs on this same code path are deliberately out of
scope, one logical change per commit:
`_current_savings_tracker_totals()` at `:798` rebuilds
`cost_tracker.stats()` per request, and `_resolve_litellm_model` in
`savings_tracker.py` is uncached across roughly seven calls per request.
Happy to follow up on either.

`make ci-precheck` was not run end to end. The uv-managed worktree venv
has no `pip`, so the `ci-precheck-python` hook's `pip install -e .` step
fails on this machine for reasons unrelated to the change. Ran `pytest`,
`ruff`, and `mypy` directly instead, output above. No Rust touched.

One heads up worth passing on, since it is why the numbers above are a
sweep and not a single full-suite line. Seven test files do not complete
on this macOS box: six hang and one fails. The hangs park the main
thread in `_dispatch_semaphore_wait_slow` with CPU time frozen and never
recover, and `pytest-timeout --timeout-method=signal` cannot break them
out, so the block is native, below the interpreter.
`tests/test_adversarial_grid.py::test_grid_shape_and_schema` is the
first one a full run reaches.

All seven reproduce identically on unmodified `c400f908` with this
change reverted, so they predate the PR. I ran the reverted comparison
specifically to rule out a thread-before-fork interaction from the new
`to_thread` call, which was the plausible way this change could have
caused it. It did not. Happy to open a separate issue with the sample
output if that is useful.


---

## Follow-up: a cancellation bug this change introduced

Self-review turned up a second problem in this change, so the fix rides
along here.

Moving the append into a thread added the first suspension point in
`record_request` that can tear state. The metrics lock at
`prometheus_metrics.py:701` suspends too, but only under contention, and
it sits ahead of every mutation, so a cancellation there recorded
nothing at all. The new await is different. It sits after the Prometheus
counters commit and before OTel and the funnel's effects 2/3/4, and it
suspends on every compressed request.

Four of the funnel's call sites are `finally:` blocks inside streaming
async generators (`streaming.py:1611`, `:1859`, `:2069`,
`openai.py:8614`). A client disconnect cancels that task. The
cancellation lands on the new await, so Prometheus counts the request
while the cost tracker, the request log, and the PERF line `headroom
perf` reads never see it. `emit_request_outcome` has one try/except and
it sits before `record_request`, so nothing catches this.

`_record_request_outcome` now wraps the funnel in `asyncio.shield`. One
line at a single choke point, covering all 28 call sites. The shield
leaves the cancellation itself alone: the await still raises
`CancelledError`, so generator teardown propagates as before. Only the
bookkeeping survives.

Real stack, uvicorn 0.40.0 + starlette 1.3.1, raw-socket disconnect
mid-stream:

| effect | before | after |
|---|---|---|
| Prometheus counters | committed | committed |
| ledger write | ran | ran |
| OTel | **skipped** | ran |
| cost tracker | **skipped** | ran |
| request log | **skipped** | ran |
| PERF line | **skipped** | ran |
| caller sees `CancelledError` | yes | yes |

The new test fails on its parent commit with a `TimeoutError`.

## Test changes

Dropped `test_event_loop_keeps_running_during_the_ledger_write`. It
detected a strict subset of what the lock test already detects:

| scenario | lock test | loop test |
|---|---|---|
| correct: outside lock + `to_thread` | PASS | PASS |
| regress: INSIDE lock + `to_thread` | FAIL | **PASS** |
| regress: outside lock + sync write | FAIL | FAIL |
| pre-fix: INSIDE lock + sync write | FAIL | FAIL |

Added a concurrency test in its place, which covers what the offload
actually introduces: before the move every proxy ledger write ran on the
one event-loop thread and was serialised for free, and now N in-flight
requests append from N worker threads.

One thing left open. That same intra-process concurrency reaches
`_maybe_compact`, which rewrites the file in place. On POSIX the
ledger's own `flock` serialises it. On Windows `_HAS_FCNTL` is false and
all locking is skipped, so a single Windows proxy can now interleave
writers where the loop thread used to serialise them. The cross-process
form of that is pre-existing and called out at `savings_ledger.py:38`.
Happy to take the intra-process guard here or in a follow-up.

Two notes on the sweep above, now that the diff is three files. The
blast radius re-run at this head is 136 of 145 files green, and the nine
non-green are identical with and without the change. Two of them
(`test_gemini_function_response_waste.py`,
`test_openai_responses_context_compaction.py`) are not in the seven
listed earlier; I re-ran both against a reverted `server.py` and they
hang the same way on both sides.
This commit is contained in:
inix 2026-07-21 02:01:34 +08:00 • committed by GitHub
parent a7dcb9e91c
commit 4aac068814
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 262 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

59
headroom/proxy/prometheus_metrics.py
View file

@ -810,29 +810,42 @@ class PrometheusMetrics:
output_tokens_saved=output_tokens_saved,
)
# Also append to the durable, multi-process savings ledger so
# `headroom savings` reflects proxy traffic alongside MCP-tool usage.
# The real upstream model means litellm prices it accurately. The
# client is the harness classified from the User-Agent / X-Client
# (claude-code, codex, cursor, ...); it falls back to "proxy" only
# when the harness is unidentified.
if tokens_saved > 0 and not self._stateless:
# `input_tokens` here is the optimized (post-compression) count
# that was actually forwarded — see emit_request_outcome, which
# passes `input_tokens=outcome.optimized_tokens`. The ledger's
# `before` is the pre-compression original and `after` is what we
# forwarded, and `headroom savings` derives the reduction percent
# as saved / before. Passing the forwarded count as `before`
# understated the original by `tokens_saved`, inflating that
# percentage (e.g. a real 40% reduction was reported as ~67%).
# Reconstruct the original as forwarded + saved.
savings_ledger.record_savings_event(
tokens_before=input_tokens + tokens_saved,
tokens_after=input_tokens,
model=model,
client=client or "proxy",
source="proxy",
)
# Also append to the durable, multi-process savings ledger so
# `headroom savings` reflects proxy traffic alongside MCP-tool usage.
# The real upstream model means litellm prices it accurately. The
# client is the harness classified from the User-Agent / X-Client
# (claude-code, codex, cursor, ...); it falls back to "proxy" only
# when the harness is unidentified.
#
# Deliberately outside `self._lock` and off the loop. The append does
# synchronous open + fcntl.flock + write, and rewrites the whole file
# once it passes 1 MB — and it fires on every compressed request. Under
# the lock that queued every other metrics caller behind the disk,
# including `export()`, which holds the same lock for the full
# Prometheus serialization. The ledger takes its own flock across
# processes, so the metrics lock was never what made it safe. Moving it
# out is not optional once it becomes an await: awaiting inside the lock
# would hold the lock for the whole write instead of just the syscall.
# ponytail: default thread pool, not a dedicated executor -- give it one
# if a profile ever shows writers parked on flock saturating the pool.
if tokens_saved > 0 and not self._stateless:
# `input_tokens` here is the optimized (post-compression) count
# that was actually forwarded — see emit_request_outcome, which
# passes `input_tokens=outcome.optimized_tokens`. The ledger's
# `before` is the pre-compression original and `after` is what we
# forwarded, and `headroom savings` derives the reduction percent
# as saved / before. Passing the forwarded count as `before`
# understated the original by `tokens_saved`, inflating that
# percentage (e.g. a real 40% reduction was reported as ~67%).
# Reconstruct the original as forwarded + saved.
await asyncio.to_thread(
savings_ledger.record_savings_event,
tokens_before=input_tokens + tokens_saved,
tokens_after=input_tokens,
model=model,
client=client or "proxy",
source="proxy",
)
self._get_otel_metrics().record_proxy_request(
provider=provider,

13
headroom/proxy/server.py
View file

@ -1936,7 +1936,18 @@ class HeadroomProxy(
"""
from headroom.proxy.outcome import emit_request_outcome
await emit_request_outcome(self, outcome)
# Shielded because four call sites are `finally:` blocks inside streaming
# async generators (streaming.py:1611, :1859, :2069, openai.py:8614). A
# client disconnect cancels that task, and the funnel now suspends partway
# through: `record_request` commits the Prometheus counters, then awaits
# the ledger append in a worker thread. A cancellation landing on that
# await leaves the request counted in Prometheus but missing from the cost
# tracker, the request log, and the PERF line `headroom perf` reads.
#
# The shield does not swallow the cancellation — the await below still
# raises CancelledError, so generator teardown propagates exactly as
# before. It only keeps the bookkeeping from being torn in half.
await asyncio.shield(emit_request_outcome(self, outcome))
async def _next_request_id(self) -> str:
"""Generate unique request ID."""

55
tests/test_request_outcome.py
View file

@ -11,6 +11,8 @@ silently regress the wire shape.
from __future__ import annotations
import asyncio
import contextlib
import logging
from dataclasses import FrozenInstanceError
from typing import Any
@ -326,6 +328,59 @@ async def test_funnel_skips_request_log_when_logger_absent() -> None:
h.metrics.record_request.assert_awaited_once() # still happens
@pytest.mark.asyncio
async def test_funnel_tail_survives_cancellation_inside_record_request() -> None:
"""A client disconnect must not tear per-request bookkeeping in half.
Four call sites are ``finally:`` blocks inside streaming async generators
(``streaming.py:1611``, ``:1859``, ``:2069``, ``openai.py:8614``), and
``record_request`` suspends partway through — it awaits the savings-ledger
append in a worker thread after the Prometheus counters have already been
committed. A cancellation landing on that await used to skip every effect
below it, leaving the request counted in Prometheus but absent from the cost
tracker, the request log, and the PERF line ``headroom perf`` reads.
Without the ``asyncio.shield`` in ``_record_request_outcome`` the release
below never resumes the funnel and this test times out on ``logged``.
"""
h = _FunnelHarness()
logged = asyncio.Event()
collect = h.logger.log
def log_and_signal(entry: Any) -> None:
collect(entry)
logged.set()
h.logger.log = log_and_signal # type: ignore[method-assign]
entered = asyncio.Event()
release = asyncio.Event()
async def suspending_record_request(**kwargs: Any) -> None:
# Stands in for the `await asyncio.to_thread(...)` ledger append: the
# counters are in, and the funnel is now parked on an await.
entered.set()
await release.wait()
h.metrics.record_request = suspending_record_request
task = asyncio.create_task(h._record_request_outcome(_outcome()))
await asyncio.wait_for(entered.wait(), timeout=5)
task.cancel()
# The shield deliberately does not swallow the cancellation — the caller
# still sees CancelledError, so generator teardown propagates unchanged.
with contextlib.suppress(asyncio.CancelledError):
await task
release.set()
await asyncio.wait_for(logged.wait(), timeout=5)
assert h.cost_tracker.record_tokens.called, "cost tracker was skipped by the cancellation"
assert len(h.logger.logs) == 1, "request log was skipped by the cancellation"
@pytest.mark.asyncio
async def test_funnel_emits_perf_log_with_canonical_shape(
caplog: pytest.LogCaptureFixture,

159
tests/test_savings_ledger_offload.py Normal file
View file

@ -0,0 +1,159 @@
"""The proxy savings-ledger append must not block the loop or hold the metrics lock.
``PrometheusMetrics.record_request`` appends one durable JSONL event per
compressed request. That append does synchronous ``open`` + ``fcntl.flock`` +
``write``, and rewrites the whole file once it passes 1 MB. Running it on the
event loop stalls every other request. Running it under ``self._lock`` also
queues every other metrics caller behind it, including the ``/metrics`` scrape
(``export`` holds that same lock for the full Prometheus serialization).
"""
from __future__ import annotations
import asyncio
import time
from pathlib import Path
from typing import Any
import pytest
from headroom import savings_ledger
from headroom.proxy import prometheus_metrics
# Long enough to dwarf scheduler noise, short enough to keep the suite quick.
_WRITE_SECONDS = 0.5
class _FakeSavingsTracker:
def snapshot(self) -> dict[str, dict[str, int | float]]:
return {"lifetime": {"total_input_tokens": 0, "total_input_cost_usd": 0.0}}
def record_request(self, **kwargs: Any) -> None:
pass
def record_lifetime_request(self, **kwargs: Any) -> None:
pass
class _FakeOtelMetrics:
def record_proxy_request(self, **kwargs: Any) -> None:
pass
def _metrics(**kwargs: Any) -> prometheus_metrics.PrometheusMetrics:
return prometheus_metrics.PrometheusMetrics(
savings_tracker=_FakeSavingsTracker(),
otel_metrics=_FakeOtelMetrics(),
**kwargs,
)
async def _record(
metrics: prometheus_metrics.PrometheusMetrics, *, tokens_saved: int = 400
) -> None:
await metrics.record_request(
provider="anthropic",
model="claude-opus-4-6",
input_tokens=600,
output_tokens=25,
tokens_saved=tokens_saved,
latency_ms=10.0,
client="claude-code",
)
@pytest.mark.asyncio
async def test_metrics_lock_is_free_while_the_ledger_write_runs(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""A concurrent ``self._lock`` holder proceeds mid-write, not after it.
``export()`` takes this same lock, so a write held under it blocks
``/metrics`` for the full duration of the disk write.
"""
window: dict[str, float] = {}
def slow_record(**kwargs: Any) -> None:
window["start"] = time.perf_counter()
time.sleep(_WRITE_SECONDS)
window["end"] = time.perf_counter()
monkeypatch.setattr(prometheus_metrics.savings_ledger, "record_savings_event", slow_record)
metrics = _metrics()
async def competitor() -> float:
async with metrics._lock:
return time.perf_counter()
_, acquired = await asyncio.gather(_record(metrics), competitor())
assert window, "the ledger write never ran"
# Only an upper bound. Once the write is offloaded, the competitor takes the
# free lock on the loop thread before the worker has even started, so
# `acquired` legitimately precedes `window["start"]`. What must not happen is
# the competitor queueing until the write is done.
assert acquired < window["end"] - _WRITE_SECONDS / 2, (
"the metrics lock was held across the ledger write: acquired "
f"{acquired - window['start']:+.3f}s relative to write start "
f"(write took {window['end'] - window['start']:.3f}s)"
)
@pytest.mark.asyncio
async def test_event_is_on_disk_once_record_request_returns(
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
"""Offloading must stay awaited: callers still see a durable write on return."""
monkeypatch.setenv("HEADROOM_SAVINGS_EVENTS_PATH", str(tmp_path / "savings_events.jsonl"))
await _record(_metrics())
report = savings_ledger.aggregate_savings()
assert report.lifetime["calls"] == 1
assert report.lifetime["tokens_saved"] == 400
assert report.lifetime["tokens_before"] == 1000
@pytest.mark.asyncio
@pytest.mark.parametrize(
("tokens_saved", "stateless"),
[(0, False), (400, True)],
)
async def test_no_ledger_write_when_gated_out(
monkeypatch: pytest.MonkeyPatch, tokens_saved: int, stateless: bool
) -> None:
"""The ``tokens_saved > 0 and not stateless`` gate survives the move."""
calls: list[dict[str, Any]] = []
monkeypatch.setattr(
prometheus_metrics.savings_ledger,
"record_savings_event",
lambda **kwargs: calls.append(kwargs),
)
await _record(_metrics(stateless=stateless), tokens_saved=tokens_saved)
assert calls == []
@pytest.mark.asyncio
async def test_concurrent_requests_all_land_their_events(
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
"""Offloading means N in-flight requests append from N worker threads.
Before the move every proxy ledger write ran on the one event-loop thread,
so they were serialised for free. Now they are not, and the ledger's own
``flock`` plus its past-1 MB full-file rewrite are what has to hold the line.
"""
monkeypatch.setenv("HEADROOM_SAVINGS_EVENTS_PATH", str(tmp_path / "savings_events.jsonl"))
metrics = _metrics()
await asyncio.gather(*(_record(metrics) for _ in range(24)))
report = savings_ledger.aggregate_savings()
assert report.lifetime["calls"] == 24, "a concurrent append was lost"
assert report.lifetime["tokens_saved"] == 24 * 400