perf(proxy): reduce lock contention on stage-timing metrics

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

@@ -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)}'
                     )

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