diff --git a/headroom/proxy/prometheus_metrics.py b/headroom/proxy/prometheus_metrics.py index d4b373a11..c797fee28 100644 --- a/headroom/proxy/prometheus_metrics.py +++ b/headroom/proxy/prometheus_metrics.py @@ -11,6 +11,7 @@ from __future__ import annotations import asyncio import logging +import threading from collections import defaultdict from datetime import datetime from typing import TYPE_CHECKING @@ -170,6 +171,13 @@ class PrometheusMetrics: ) self._lock = asyncio.Lock() + # Tiny synchronous critical section for stage-timing triple updates + # (sum + count + max must move together for a consistent scrape). + # threading.Lock is cheaper than asyncio.Lock and does NOT contend + # with the async ``export()`` path — scrapes snapshot these dicts + # under this lock in a microsecond block, then build the metrics + # string without holding anything. + self._stage_timing_lock = threading.Lock() self._otel_metrics = otel_metrics def _get_otel_metrics(self) -> HeadroomOtelMetrics: @@ -348,10 +356,19 @@ class PrometheusMetrics: maps stage names to millisecond durations. Mirrors the ``transform_timing_*`` aggregation pattern so the ``/metrics`` endpoint exposes sum/count/max series per ``(path, stage)``. + + Uses a tiny synchronous ``threading.Lock`` around the triple + update (sum + count + max) rather than the async + ``self._lock``: (1) the updates have no awaits, so there is no + async contention benefit, and (2) the async lock is also held + by ``export()`` during Prometheus scrapes — which does + string-building while holding it. Under N concurrent request + finalizations + an active scrape, callers would queue behind + the scrape's string-building. """ if not timings: return - async with self._lock: + with self._stage_timing_lock: for stage, ms in timings.items(): try: ms_val = float(ms) @@ -423,6 +440,15 @@ class PrometheusMetrics: async def export(self) -> str: """Export metrics in Prometheus format.""" + # Snapshot stage-timing dicts under the tiny synchronous lock so + # we don't race a concurrent ``record_stage_timings`` and observe + # an inconsistent (sum, count, max) triple. Freeze into plain + # dicts so the scrape's string-building below doesn't hold the + # stage-timing lock during I/O-ish work. + with self._stage_timing_lock: + stage_timing_sum_snapshot = dict(self.stage_timing_sum) + stage_timing_count_snapshot = dict(self.stage_timing_count) + stage_timing_max_snapshot = dict(self.stage_timing_max) async with self._lock: lines: list[str] = [] _append_metric( @@ -628,14 +654,14 @@ class PrometheusMetrics: ) lines.append("") - if self.stage_timing_sum: + if stage_timing_sum_snapshot: lines.extend( [ "# HELP headroom_stage_timing_ms_sum Sum of per-stage handler timings in milliseconds", "# TYPE headroom_stage_timing_ms_sum counter", ] ) - for (path_label, stage), total in self.stage_timing_sum.items(): + for (path_label, stage), total in stage_timing_sum_snapshot.items(): lines.append( f'headroom_stage_timing_ms_sum{{path="{_escape_label_value(path_label)}",stage="{_escape_label_value(stage)}"}} {round(total, 2)}' ) @@ -646,7 +672,7 @@ class PrometheusMetrics: "# TYPE headroom_stage_timing_ms_count counter", ] ) - for (path_label, stage), count in self.stage_timing_count.items(): + for (path_label, stage), count in stage_timing_count_snapshot.items(): lines.append( f'headroom_stage_timing_ms_count{{path="{_escape_label_value(path_label)}",stage="{_escape_label_value(stage)}"}} {count}' ) @@ -657,7 +683,7 @@ class PrometheusMetrics: "# TYPE headroom_stage_timing_ms_max gauge", ] ) - for (path_label, stage), max_value in self.stage_timing_max.items(): + for (path_label, stage), max_value in stage_timing_max_snapshot.items(): lines.append( f'headroom_stage_timing_ms_max{{path="{_escape_label_value(path_label)}",stage="{_escape_label_value(stage)}"}} {round(max_value, 2)}' ) diff --git a/tests/test_prometheus_stage_timing_concurrency.py b/tests/test_prometheus_stage_timing_concurrency.py new file mode 100644 index 000000000..34c550d3f --- /dev/null +++ b/tests/test_prometheus_stage_timing_concurrency.py @@ -0,0 +1,94 @@ +"""Concurrency test for PrometheusMetrics stage-timing lock-free path. + +Covers the P2 fix that moves ``record_stage_timings`` off the global +``asyncio.Lock`` and onto a tiny synchronous ``threading.Lock`` so +request finalizations don't queue behind ``export()``'s string-building. + +The invariant we assert: sum / count / max agree after many concurrent +writes + concurrent scrapes. +""" + +from __future__ import annotations + +import asyncio + +import pytest + +from headroom.proxy.prometheus_metrics import PrometheusMetrics + + +@pytest.mark.asyncio +async def test_record_stage_timings_agrees_under_concurrent_scrapes(): + metrics = PrometheusMetrics() + + writes_per_task = 25 + writer_count = 20 + scrape_count = 5 + + async def _writer(stage_id: int) -> None: + for i in range(writes_per_task): + await metrics.record_stage_timings( + "some_path", + {f"stage-{stage_id}": float(i + 1)}, + ) + + async def _scraper() -> None: + for _ in range(10): + await metrics.export() + await asyncio.sleep(0) # yield + + writers = [asyncio.create_task(_writer(i)) for i in range(writer_count)] + scrapers = [asyncio.create_task(_scraper()) for _ in range(scrape_count)] + await asyncio.gather(*writers, *scrapers) + + # Invariant: for each (path, stage) key, count equals writes_per_task + # and sum equals 1+2+...+writes_per_task. + expected_count = writes_per_task + expected_sum = sum(range(1, writes_per_task + 1)) + expected_max = float(writes_per_task) + + for stage_id in range(writer_count): + key = ("some_path", f"stage-{stage_id}") + assert metrics.stage_timing_count[key] == expected_count, ( + f"stage {stage_id}: count mismatch " + f"got={metrics.stage_timing_count[key]}, want={expected_count}" + ) + assert metrics.stage_timing_sum[key] == pytest.approx(expected_sum), ( + f"stage {stage_id}: sum mismatch " + f"got={metrics.stage_timing_sum[key]}, want={expected_sum}" + ) + assert metrics.stage_timing_max[key] == pytest.approx(expected_max), ( + f"stage {stage_id}: max mismatch " + f"got={metrics.stage_timing_max[key]}, want={expected_max}" + ) + + +@pytest.mark.asyncio +async def test_record_stage_timings_does_not_hold_async_lock(): + """Sanity: ``record_stage_timings`` must not contend with the async + lock. If it did, a long-running holder of the async lock would block + stage-timing writes.""" + metrics = PrometheusMetrics() + + # Hold the async lock from a background task and ensure + # record_stage_timings completes anyway. + lock_acquired = asyncio.Event() + release_lock = asyncio.Event() + + async def _hold_async_lock() -> None: + async with metrics._lock: + lock_acquired.set() + await release_lock.wait() + + holder = asyncio.create_task(_hold_async_lock()) + await lock_acquired.wait() + + # This should complete even though _lock is held. + await asyncio.wait_for( + metrics.record_stage_timings("p", {"s": 1.0}), + timeout=1.0, + ) + assert metrics.stage_timing_count[("p", "s")] == 1 + + release_lock.set() + await holder