diff --git a/headroom/proxy/prometheus_metrics.py b/headroom/proxy/prometheus_metrics.py index ef423c22a..4413feec5 100644 --- a/headroom/proxy/prometheus_metrics.py +++ b/headroom/proxy/prometheus_metrics.py @@ -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, diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index d7e698bfc..56a776d46 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -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.""" diff --git a/tests/test_request_outcome.py b/tests/test_request_outcome.py index 5f898b714..89df90a27 100644 --- a/tests/test_request_outcome.py +++ b/tests/test_request_outcome.py @@ -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, diff --git a/tests/test_savings_ledger_offload.py b/tests/test_savings_ledger_offload.py new file mode 100644 index 000000000..06f627a75 --- /dev/null +++ b/tests/test_savings_ledger_offload.py @@ -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