From b9d7dcc3da3ec67a968d0ac1e35bffba9b7cf1c2 Mon Sep 17 00:00:00 2001 From: inix <62450194+inix-x@users.noreply.github.com> Date: Mon, 24 Aug 2026 13:43:58 +0800 Subject: [PATCH] fix(proxy): make output-savings flush atomic and keep it off the event loop (#3231) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description The output-shaper's periodic savings-ledger flush ran synchronously on the asyncio event loop: every 25th shaped request, `emit_request_outcome` performed a full ledger reload (file read + `json.loads`) followed by a `json.dumps` + in-place `write_text`, with no await or executor. The write was also non-atomic, so a crash mid-write truncated the existing ledger, and `SavingsLedger.load()` silently swallowed the resulting decode error — corrupted history was indistinguishable from no history yet. ## Type of Change - [x] 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 - `emit_request_outcome` now runs `record_from_labels` + `estimate_request_savings` together on a worker thread via one `asyncio.to_thread` call — both take the recorder lock, and the periodic flush holds that lock across disk I/O, so nothing touches it from the event loop anymore. - `SavingsLedger.save()` writes through the existing `headroom.fsutil.write_text` helper (temp file in the target directory, fsync, atomic `os.replace`, temp cleanup on failure) instead of a truncating in-place write. - `SavingsLedger.load()` logs a warning naming the unreadable ledger file and still fails open with an empty ledger. - Added `TestFlushDurability` to `tests/test_output_savings.py`: failed-save intactness (+ no temp residue), corrupt-file warning, and off-loop-thread assertions. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text $ python -m pytest tests/test_output_savings.py tests/test_output_shaping_rollup.py tests/test_output_savings_cli.py -q ============================== 49 passed in 1.38s ============================== $ ruff check headroom/proxy/output_savings.py headroom/proxy/outcome.py tests/test_output_savings.py All checks passed! $ ruff format --check headroom/proxy/output_savings.py headroom/proxy/outcome.py tests/test_output_savings.py 3 files already formatted $ mypy headroom/proxy/output_savings.py headroom/proxy/outcome.py Success: no issues found in 2 source files Fail-before evidence (the three files checked out at upstream/main, fix reverted): $ python -m pytest tests/test_output_savings.py::TestFlushDurability -q FAILED tests/test_output_savings.py::TestFlushDurability::test_crash_mid_write_leaves_previous_ledger_intact - KeyError: 'opus|code|m|tools' FAILED tests/test_output_savings.py::TestFlushDurability::test_corrupt_ledger_warns_and_starts_empty - AssertionError: corrupt ledger was swallowed silently FAILED tests/test_output_savings.py::TestFlushDurability::test_emit_request_outcome_flushes_off_the_loop_thread - assert False 3 failed in 0.49s ``` ## Real Behavior Proof - Environment: macOS 15 (arm64), CPython 3.13.13, project venv; branch `fix/output-savings-atomic-offload` = upstream/main `7784bb18` + the single fix commit. - Exact command / steps: With the three touched files reverted to upstream/main: `python -m pytest tests/test_output_savings.py::TestFlushDurability -q` → all 3 new tests fail (torn write destroys the prior ledger; no warning on a corrupt file; flush observed on the loop thread). Re-applied the commit and re-ran the same command plus ruff/format/mypy as pasted under Test Output. - Observed result: All 3 fail-before cases now pass — a save failure before rename leaves the previous ledger loadable with no `*.tmp` residue, a corrupt ledger logs a warning and still fails open empty, and the flush triggered through `emit_request_outcome` runs on a worker thread distinct from the event-loop thread; 49 recorder/rollup/CLI tests pass. - Not tested: Windows behavior of the atomic rename (covered by `fsutil.write_text`, exercised only on POSIX here), a full local suite run (unrelated pre-existing native hangs on macOS), and the dashboard rendering of the ledger. ## Runtime Rollout Safety - Rollout-managed feature(s): None changed. The shaper itself is opt-in; this PR only changes how and where its ledger persistence happens. - Minimum rollout channel: Stable. - Stable/default behavior changed: No. With output shaping disabled the funnel never reaches this code path; when enabled, identical data is persisted — written atomically instead of truncating, and off the event loop. - Kill switch / disable path: Unset `HEADROOM_OUTPUT_SHAPER` (or disable the `proxy_output_shaper` rollout flag); the recorder then neither records nor flushes. - Unsafe override required: No. - Qualification impact: None. - Rollback path: Revert this single commit; the on-disk ledger format is unchanged, so no data migration is involved either way. ## 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) ## Additional Notes - Two failures seen while running neighbouring suites locally (`tests/test_stateless_writers.py::test_memory_disabled_under_stateless`, `tests/test_5xx_accounting_all_providers.py::test_gemini_count_tokens_handler_threads_real_529_onto_outcome`) reproduce on pristine upstream/main without this patch — pre-existing, not introduced here. - Known adjacent shapes deliberately left out of scope: `get_recorder().estimate()` in the `/stats` payload also reads the ledger file inline (already exception-guarded there), and `SavingsRecorder.flush()` is not yet wired into graceful shutdown. --- headroom/proxy/outcome.py | 16 +++-- headroom/proxy/output_savings.py | 15 ++++- tests/test_output_savings.py | 103 +++++++++++++++++++++++++++++-- 3 files changed, 122 insertions(+), 12 deletions(-) diff --git a/headroom/proxy/outcome.py b/headroom/proxy/outcome.py index c2b78b1dc..ba94feee7 100644 --- a/headroom/proxy/outcome.py +++ b/headroom/proxy/outcome.py @@ -25,6 +25,7 @@ actually reports. from __future__ import annotations +import asyncio import logging from dataclasses import dataclass, field from datetime import datetime, timezone @@ -455,10 +456,17 @@ async def emit_request_outcome(handler: Any, outcome: RequestOutcome) -> None: from headroom.proxy.output_savings import get_recorder _rec = get_recorder() - _rec.record_from_labels(outcome.transforms_applied, outcome.output_tokens) - output_tokens_saved_est = _rec.estimate_request_savings( - outcome.transforms_applied, outcome.output_tokens - ) + + def _record_and_estimate() -> int: + _rec.record_from_labels(outcome.transforms_applied, outcome.output_tokens) + return _rec.estimate_request_savings( + outcome.transforms_applied, outcome.output_tokens + ) + + # Both calls take the recorder lock, and the every-Nth record also + # does a full read-modify-write of the ledger file — run them + # together off the event loop (#18) so a slow flush can't stall it. + output_tokens_saved_est = await asyncio.to_thread(_record_and_estimate) except Exception: # pragma: no cover - defensive pass diff --git a/headroom/proxy/output_savings.py b/headroom/proxy/output_savings.py index b0ee6baf3..5fe317cb0 100644 --- a/headroom/proxy/output_savings.py +++ b/headroom/proxy/output_savings.py @@ -39,6 +39,7 @@ Pure module: no I/O except explicit ``load``/``save``. from __future__ import annotations import json +import logging import math from dataclasses import asdict, dataclass, field from typing import Any @@ -68,6 +69,8 @@ from .output_savings_policy import ( stratum_label as stratum_label, ) +logger = logging.getLogger(__name__) + @dataclass class _Accum: @@ -328,9 +331,13 @@ class SavingsLedger: def save(self, path: Any) -> None: from pathlib import Path + from headroom import fsutil + p = Path(path) p.parent.mkdir(parents=True, exist_ok=True) - p.write_text(json.dumps(self.to_dict(), separators=(",", ":"))) + # fsutil.write_text is atomic (temp file + os.replace), so a crash + # mid-write cannot truncate the ledger already on disk (#18). + fsutil.write_text(p, json.dumps(self.to_dict(), separators=(",", ":"))) @classmethod def load(cls, path: Any) -> SavingsLedger: @@ -341,7 +348,11 @@ class SavingsLedger: return cls() try: return cls.from_dict(json.loads(p.read_text())) - except (json.JSONDecodeError, ValueError, OSError): + except (json.JSONDecodeError, ValueError, OSError) as exc: + # Fail open (empty ledger), but surface the loss — silently + # swallowing a corrupt file made lost history indistinguishable + # from no history yet (#18). + logger.warning("output-savings ledger %s unreadable, starting empty: %s", p, exc) return cls() diff --git a/tests/test_output_savings.py b/tests/test_output_savings.py index e9d67a826..f62bc3e37 100644 --- a/tests/test_output_savings.py +++ b/tests/test_output_savings.py @@ -389,12 +389,7 @@ class TestRecorderBaselineReload: @staticmethod def _key() -> str: - return stratum_key( - turn_kind="code", - input_tokens=8000, - model="claude-opus-4-8", - has_tools=True, - ) + return SAMPLE_KEY def test_adopts_baseline_learned_after_start(self, tmp_path): path = str(tmp_path / "output_savings.json") @@ -482,3 +477,99 @@ class TestRecorderBaselineReload: relearned.save(path) assert recorder.estimate().baseline_tokens > baseline_tokens_v1 + + +# --------------------------------------------------------------------------- +# flush durability + event-loop safety +# --------------------------------------------------------------------------- + +# Deterministic stratum key shared by the recorder tests below. +SAMPLE_KEY = stratum_key( + turn_kind="code", + input_tokens=8000, + model="claude-opus-4-8", + has_tools=True, +) + + +class TestFlushDurability: + def test_crash_mid_write_leaves_previous_ledger_intact(self, tmp_path, monkeypatch): + import headroom.fsutil + + path = str(tmp_path / "output_savings.json") + key = SAMPLE_KEY + + recorder = SavingsRecorder(path, flush_every=1) + recorder.record_from_labels([stratum_label("treatment", key)], 200) + recorder.flush() + assert SavingsLedger.load(path).treatment[key].n == 1 + + def _die_before_rename(*args, **kwargs): + raise OSError(5, "simulated crash before rename") + + monkeypatch.setattr(headroom.fsutil.os, "replace", _die_before_rename) + recorder.record_from_labels([stratum_label("treatment", key)], 210) + recorder.flush() # OSError swallowed by the recorder — fail-open by design + + # The pre-crash sample must survive and no temp residue may be left + # behind: a failed save may not corrupt or clutter the ledger. + assert SavingsLedger.load(path).treatment[key].n == 1 + assert not list(tmp_path.glob("*.tmp")) + + def test_corrupt_ledger_warns_and_starts_empty(self, tmp_path, caplog): + import logging + + path = tmp_path / "output_savings.json" + path.write_text("{not json") + + with caplog.at_level(logging.WARNING): + SavingsRecorder(str(path)) + + assert caplog.records, "corrupt ledger was swallowed silently" + + def test_emit_request_outcome_flushes_off_the_loop_thread(self, tmp_path, monkeypatch): + import asyncio + import threading + + from headroom.proxy.outcome import RequestOutcome, emit_request_outcome + + path = str(tmp_path / "output_savings.json") + recorder = SavingsRecorder(path, flush_every=1) + monkeypatch.setattr("headroom.proxy.output_savings.get_recorder", lambda: recorder) + + saved_on_threads = [] + real_save = SavingsLedger.save + + def _spy_save(self, save_path): + saved_on_threads.append(threading.get_ident()) + real_save(self, save_path) + + monkeypatch.setattr(SavingsLedger, "save", _spy_save) + + class _Metrics: + async def record_request(self, **kwargs): + pass + + class _Handler: + def __init__(self): + self.metrics = _Metrics() + self.cost_tracker = None + self.logger = None + + outcome = RequestOutcome( + request_id="req-shaper", + provider="openai", + model="gpt-5", + status_code=200, + original_tokens=100, + optimized_tokens=80, + output_tokens=50, + tokens_saved=20, + attempted_input_tokens=100, + transforms_applied=(stratum_label("treatment", SAMPLE_KEY),), + ) + asyncio.run(emit_request_outcome(_Handler(), outcome)) + + loop_thread = threading.get_ident() + assert saved_on_threads, "flush never ran" + assert all(t != loop_thread for t in saved_on_threads)