diff --git a/headroom/proxy/savings_attribution.py b/headroom/proxy/savings_attribution.py index 56e2fe209..2ea8c3fb4 100644 --- a/headroom/proxy/savings_attribution.py +++ b/headroom/proxy/savings_attribution.py @@ -4,6 +4,7 @@ from __future__ import annotations import base64 import json +import math import re from collections.abc import MutableMapping from typing import Any @@ -33,6 +34,23 @@ MAX_STAGES = 16 # must not accumulate into the same series. STAGE_PREFIX = "ext:" +# NON-FINITE VALUES POISON EVERY CONSUMER DOWNSTREAM, and they do it long after +# the call that introduced them. Starlette's JSONResponse encodes with +# ``allow_nan=False``, so a single ``inf`` reaching ``/stats`` raises +# ``ValueError: Out of range float values are not JSON compliant`` -- and the +# value sits in the process-wide metrics totals, so the endpoint stays broken +# until restart. Prometheus is no better: Python renders ``inf``, the exposition +# format wants ``+Inf``, and the scrape fails to parse. +# +# The request itself still returns 200 throughout, which is the worst shape a +# bug can have: the extension looks healthy while the operator's dashboard and +# scrape are dead. +# +# One hour bounds a single stage inside one request -- unreachable in practice, +# and it makes overflow-to-infinity on accumulation structurally impossible +# (16 stages x 1h is nowhere near the float ceiling). +MAX_STAGE_MS = 3_600_000.0 + def _source_name(value: object) -> str: name = _NAME_RE.sub("_", str(value or "other").strip().lower()).strip("_.-") @@ -90,10 +108,12 @@ def record_scope_timing(scope: MutableMapping[str, Any], stage: object, ms: floa elapsed = float(ms) except (TypeError, ValueError): return - if not elapsed > 0.0: - # Non-positive is either a clock artifact or nothing happening. Either - # way it is not a measurement, and averaging it in would drag the mean - # toward zero exactly where the stage is cheapest to ignore. + # Non-positive is either a clock artifact or nothing happening; either way + # it is not a measurement, and averaging it in would drag the mean toward + # zero exactly where the stage is cheapest to ignore. Non-finite and + # absurdly large are not measurements either, and they break consumers + # rather than merely skewing them -- see MAX_STAGE_MS. + if not math.isfinite(elapsed) or not 0.0 < elapsed <= MAX_STAGE_MS: return state = scope.setdefault("state", {}) @@ -119,7 +139,16 @@ def timings_from_tags(tags: MutableMapping[str, Any] | None) -> dict[str, float] elapsed = float(value) except (TypeError, ValueError): continue - if elapsed > 0.0: + # Re-checked rather than trusted: the ledger is a plain dict reachable + # through ``tags``, so a handler can be handed one this module never + # wrote. The guarantee has to hold at the read, not only at the write. + # + # Finiteness ONLY. ``MAX_STAGE_MS`` bounds a single sample at the write, + # where it prevents overflow; applying it here would test it against an + # ACCUMULATED total and silently discard a stage that legitimately ran + # for longer across many samples -- throwing away real data to guard + # against a value this path cannot produce. + if math.isfinite(elapsed) and elapsed > 0.0: out[str(name)] = elapsed return out @@ -138,12 +167,25 @@ def record_savings( ledger = _ledger(tags) if len(ledger) >= MAX_SOURCES: return + # Same hazard as MAX_STAGE_MS, on the amounts rather than the durations: + # ``usd=inf`` reaches ``/stats`` and raises out of the JSON encoder, and + # ``int(inf)`` raises OverflowError right here, inside the handler, on a + # request that would otherwise have succeeded. Neither is a saving, so + # neither is recorded -- the alternative is a plugin's arithmetic bug + # taking down an endpoint it has nothing to do with. + try: + amount = float(usd or 0.0) + count = int(tokens or 0) + except (TypeError, ValueError, OverflowError): + return + if not math.isfinite(amount): + return item: dict[str, Any] = { "source": _source_name(source), "realized": bool(realized), "estimated": bool(estimated), - "tokens": max(0, int(tokens or 0)), - "usd": round(float(usd or 0.0), 12), + "tokens": max(0, count), + "usd": round(amount, 12), } if details: item["details"] = { diff --git a/tests/test_extension_attribution.py b/tests/test_extension_attribution.py index 58a227341..77deb8f61 100644 --- a/tests/test_extension_attribution.py +++ b/tests/test_extension_attribution.py @@ -9,9 +9,12 @@ half did not exist at all, so an extension's own latency was invisible — from __future__ import annotations +import math + import pytest from headroom.proxy.savings_attribution import ( + MAX_STAGE_MS, MAX_STAGES, SAVINGS_ATTRIBUTION_TAG, STAGE_PREFIX, @@ -106,10 +109,29 @@ def test_extension_stages_are_namespaced() -> None: assert list(timings_from_tags(tags)) == [f"{STAGE_PREFIX}deep_copy"] -@pytest.mark.parametrize("bad", [0, -1.0, None, "slow", float("nan")]) +@pytest.mark.parametrize( + "bad", + [ + 0, + -1.0, + None, + "slow", + float("nan"), + float("inf"), + float("-inf"), + 1e400, + MAX_STAGE_MS + 1, + ], +) def test_a_non_measurement_is_not_recorded(bad) -> None: """Zero and negative are clock artifacts, not observations; averaging them - in would drag the mean down exactly where the stage is cheapest to skip.""" + in would drag the mean down exactly where the stage is cheapest to skip. + + Non-finite is worse than skew. Starlette encodes ``/stats`` with + ``allow_nan=False``, so one ``inf`` raises out of the JSON encoder — and it + lands in process-wide metrics totals, so the endpoint stays broken until + restart while the request that caused it returns 200. + """ scope = _scope() record_scope_timing(scope, "ext", bad) @@ -118,6 +140,78 @@ def test_a_non_measurement_is_not_recorded(bad) -> None: assert timings_from_tags(tags) == {} +def test_a_poisoned_ledger_is_rejected_on_read_too() -> None: + """The ledger is a plain dict reachable through ``tags``, so a handler can + be handed one this module never wrote. The guarantee holds at the read.""" + assert timings_from_tags({STAGE_TIMING_TAG: {"ext:a": float("inf"), "ext:b": 2.0}}) == { + "ext:b": 2.0 + } + + +def test_accumulation_cannot_overflow_to_infinity() -> None: + """Two finite values can sum to ``inf``. Bounding each SAMPLE makes that + unreachable rather than merely unlikely.""" + scope = _scope() + for _ in range(4): + record_scope_timing(scope, "ext", MAX_STAGE_MS) + + tags: dict = {} + bind_scope(tags, scope) + (total,) = timings_from_tags(tags).values() + assert math.isfinite(total) + + +def test_an_accumulated_total_may_exceed_the_per_sample_bound() -> None: + """The bound is on one sample, not on the sum. Testing it against the + accumulated total would silently discard a stage that legitimately ran + longer across many samples — throwing away real data to guard a value the + write path cannot produce.""" + scope = _scope() + for _ in range(3): + record_scope_timing(scope, "ext", MAX_STAGE_MS) + + tags: dict = {} + bind_scope(tags, scope) + assert timings_from_tags(tags) == {f"{STAGE_PREFIX}ext": MAX_STAGE_MS * 3} + + +@pytest.mark.parametrize("bad", [float("inf"), float("-inf"), float("nan")]) +def test_a_non_finite_amount_is_not_a_saving(bad) -> None: + """Pre-existing, and the same crash: ``usd=inf`` reaches ``/stats`` and + raises out of the JSON encoder.""" + scope = _scope() + record_scope_savings(scope, "buggy", usd=bad) + + tags: dict = {} + bind_scope(tags, scope) + assert from_tags(tags) == [] + + +@pytest.mark.parametrize("bad", [float("inf"), float("nan")]) +def test_a_non_finite_token_count_does_not_raise_inside_the_handler(bad) -> None: + """``int(inf)`` is an OverflowError, raised on a request that would + otherwise have succeeded. A plugin's arithmetic bug must not become the + proxy's 500.""" + scope = _scope() + record_scope_savings(scope, "buggy", tokens=bad) + + tags: dict = {} + bind_scope(tags, scope) + assert from_tags(tags) == [] + + +def test_a_real_saving_still_records_after_the_guards() -> None: + """The direction that must not be lost while hardening the other one.""" + scope = _scope() + record_scope_savings(scope, "routemegood", tokens=10, usd=0.5) + record_scope_timing(scope, "routemegood", 3.0) + + tags: dict = {} + bind_scope(tags, scope) + assert from_tags(tags)[0]["usd"] == 0.5 + assert timings_from_tags(tags) == {f"{STAGE_PREFIX}routemegood": 3.0} + + def test_stage_cardinality_is_capped() -> None: """Stage names are extension-supplied, so they are bounded like every other client-influenced label in this proxy."""