From a4345f1cce7418b8719fcfb8ce328808724cf670 Mon Sep 17 00:00:00 2001 From: Adryan Eka Vandra Date: Sat, 18 Apr 2026 01:59:34 +0700 Subject: [PATCH] perf(proxy): reduce lock contention on stage-timing metrics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit record_stage_timings() acquired the global asyncio.Lock on every request finalization — the same lock held by export() during Prometheus scrapes, which does string-building while holding it. Under N concurrent request finalizations + an active scrape, all N queued behind the scrape. Move the stage-timing triple update (sum + count + max) onto a tiny synchronous threading.Lock. This keeps the multi-field update consistent for scrapers without contending on the async lock. export() now snapshots the three dicts under the sync lock, then builds the metrics string from the snapshots outside any lock scope. --- headroom/proxy/prometheus_metrics.py | 36 ++++++- ...est_prometheus_stage_timing_concurrency.py | 94 +++++++++++++++++++ 2 files changed, 125 insertions(+), 5 deletions(-) create mode 100644 tests/test_prometheus_stage_timing_concurrency.py 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