Single-flight remote gateway health probe to stop cold-cache stampede

_probe_remote_gateway held _remote_probe_lock only around the cache
read and cache write; the network probe itself (_http_probe per path,
each up to ~2s) ran unlocked. On a cold cache a dashboard that fans out
to N panels fired N concurrent probe sets at the (possibly dead)
gateway, each blocking up to ~6s (#5455, #2476).

Wrap the existing lock in a Condition and add an in-flight set: the
first 'leader' thread marks the base_url in-flight and runs the probe
(via extracted _run_remote_probe), while latecomers wait on the
Condition for the leader's cached result instead of probing themselves.
The leader clears the marker and notify_all()s in a finally block, so
waiters are woken even if the probe raises (no deadlock); a bounded wait
budget lets a waiter fall back to probing itself if a leader hangs.
_reset_remote_probe_cache_for_tests also clears the in-flight set.

Expire the cache from the probe's COMPLETION time, not the leader's
entry time: walking every path of a hung gateway takes
len(paths) * timeout (~6s), which exceeds the 5s TTL, so
`entry_time + TTL` would store an already-expired cache line. Waiters
woken right after would then miss the cache and each re-probe the dead
gateway — collapsing single-flight and regressing latency to worse-than-
serial. A regression test drives a walk longer than the TTL and asserts
still exactly one probe set (the shipped happy-path test kept the walk
well under the TTL and so never exercised this).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

api/agent_health.py

@@ -314,7 +314,24 @@ _REMOTE_PROBE_PATHS: tuple[str, ...] = ("/health/detailed", "/health", "/v1/heal
 _REMOTE_PROBE_BODY_LIMIT_BYTES: int = 64 * 1024
 
 _remote_probe_lock = threading.Lock()
+# Condition wraps the same lock so cache reads/writes and single-flight waits
+# share one mutex (mirrors the Condition(_lock) idiom in api/session_lifecycle).
+_remote_probe_cond = threading.Condition(_remote_probe_lock)
 _remote_probe_cache: dict[str, Any] = {"url": None, "expires_at": 0.0, "result": None}
+# base_urls currently being probed by a "leader" thread. Latecomers wait on the
+# Condition for the leader's result instead of stampeding the (possibly dead)
+# gateway themselves (#5455 dashboard fan-out, #2476).
+_remote_probe_inflight: set[str] = set()
+def _remote_probe_wait_budget_s() -> float:
+    """How long a latecomer waits for the leader before giving up and self-probing.
+
+    The leader can walk every path (each up to ``_REMOTE_PROBE_TIMEOUT_S``), so
+    budget for the full walk plus a small margin; timing out is a safety valve
+    against a hung leader, never the normal path. Computed on each call (not a
+    module-level constant) so a test that monkeypatches the timeout or the path
+    list gets a budget consistent with those values instead of a stale one.
+    """
+    return _REMOTE_PROBE_TIMEOUT_S * len(_REMOTE_PROBE_PATHS) + 1.0
 
 
 def _remote_gateway_base_url() -> str | None:
@@ -395,23 +412,27 @@ def _http_probe(
         return (False, None, type(exc).__name__, None)
 
 
-def _probe_remote_gateway(base_url: str, *, now: float | None = None) -> dict[str, Any]:
-    """Return an agent-health payload dict for a remote gateway base URL.
+def _cached_remote_result_locked(base_url: str, current: float) -> dict[str, Any] | None:
+    """Return a fresh cached probe result for *base_url*, or None if stale/absent.
 
-    Result is cached for ``_REMOTE_PROBE_CACHE_TTL_S`` seconds per base_url.
+    Must be called while holding ``_remote_probe_cond``.
     """
-    current = time.monotonic() if now is None else now
-    with _remote_probe_lock:
-        if (
-            _remote_probe_cache.get("url") == base_url
-            and _remote_probe_cache.get("expires_at", 0.0) > current
-            and _remote_probe_cache.get("result") is not None
-        ):
-            cached = _remote_probe_cache["result"]
-            # Refresh checked_at so the UI shows a current timestamp without
-            # actually re-hitting the gateway.
-            return {**cached, "checked_at": _checked_at()}
+    if (
+        _remote_probe_cache.get("url") == base_url
+        and _remote_probe_cache.get("expires_at", 0.0) > current
+        and _remote_probe_cache.get("result") is not None
+    ):
+        return _remote_probe_cache["result"]
+    return None
 
+
+def _run_remote_probe(base_url: str) -> dict[str, Any]:
+    """Walk the remote gateway health paths and build a payload (no caching/locking).
+
+    This is the expensive, network-bound step: each path can block up to
+    ``_REMOTE_PROBE_TIMEOUT_S``. It runs OUTSIDE the probe lock so a single
+    "leader" thread does it while latecomers wait for the cached result.
+    """
     last_status: int | None = None
     last_error: str | None = None
     gateway_api_key = _remote_gateway_api_key()
@@ -439,46 +460,108 @@ def _probe_remote_gateway(base_url: str, *, now: float | None = None) -> dict[st
             # An over-cap body (len > limit, i.e. the +1 sentinel byte was read)
             # is treated as "alive but no parseable gateway_state" — we still
             # report the gateway as up, just without the detailed state.
-            payload = {
+            return {
                 "alive": True,
                 "checked_at": _checked_at(),
                 "details": details,
             }
-            break
         # Remember the most informative failure signal we saw.
         if status is not None:
             last_status = status
         if err is not None:
             last_error = err
-    else:
-        details: dict[str, Any] = {
-            "state": "down",
-            "reason": "remote_gateway_unreachable",
-            "endpoint": base_url,
-        }
-        if last_status is not None:
-            details["status_code"] = last_status
-        if last_error is not None:
-            details["error"] = last_error
-        payload = {
-            "alive": False,
-            "checked_at": _checked_at(),
-            "details": details,
-        }
 
-    with _remote_probe_lock:
-        _remote_probe_cache["url"] = base_url
-        _remote_probe_cache["expires_at"] = current + _REMOTE_PROBE_CACHE_TTL_S
-        _remote_probe_cache["result"] = payload
+    details = {
+        "state": "down",
+        "reason": "remote_gateway_unreachable",
+        "endpoint": base_url,
+    }
+    if last_status is not None:
+        details["status_code"] = last_status
+    if last_error is not None:
+        details["error"] = last_error
+    return {
+        "alive": False,
+        "checked_at": _checked_at(),
+        "details": details,
+    }
+
+
+def _probe_remote_gateway(base_url: str, *, now: float | None = None) -> dict[str, Any]:
+    """Return an agent-health payload dict for a remote gateway base URL.
+
+    Result is cached for ``_REMOTE_PROBE_CACHE_TTL_S`` seconds per base_url.
+
+    Concurrency (single-flight): when several dashboard panels fan out on a cold
+    cache, only the first "leader" thread runs the ~2s-per-path network probe;
+    latecomers wait on ``_remote_probe_cond`` for the leader's cached result
+    rather than each hammering the (possibly dead) gateway (#5455, #2476). The
+    leader always clears the in-flight marker and wakes waiters — even on error —
+    so waiters can never deadlock.
+    """
+    current = time.monotonic() if now is None else now
+    with _remote_probe_cond:
+        cached = _cached_remote_result_locked(base_url, current)
+        if cached is not None:
+            # Refresh checked_at so the UI shows a current timestamp without
+            # actually re-hitting the gateway.
+            return {**cached, "checked_at": _checked_at()}
+
+        # A leader is already probing this base_url: wait for its result instead
+        # of starting a duplicate probe.
+        if base_url in _remote_probe_inflight:
+            deadline = current + _remote_probe_wait_budget_s()
+            while base_url in _remote_probe_inflight:
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    break  # leader stuck; fall through and probe ourselves
+                _remote_probe_cond.wait(remaining)
+            cached = _cached_remote_result_locked(base_url, time.monotonic())
+            if cached is not None:
+                return {**cached, "checked_at": _checked_at()}
+            # Woke without a usable cached result (leader failed, produced no
+            # cacheable result, or we timed out): become a leader ourselves.
+
+        # Become the leader for this base_url.
+        _remote_probe_inflight.add(base_url)
+
+    payload: dict[str, Any] | None = None
+    try:
+        payload = _run_remote_probe(base_url)
+    finally:
+        with _remote_probe_cond:
+            if payload is not None:
+                _remote_probe_cache["url"] = base_url
+                # Expire from the moment the probe COMPLETES, not from the
+                # leader's entry time: walking every path of a hung gateway
+                # takes len(paths) * timeout (~6s) which exceeds the 5s TTL, so
+                # `current + TTL` would write an already-expired cache line.
+                # Waiters woken right after this would then miss the cache and
+                # each re-probe the dead gateway — collapsing single-flight and
+                # regressing latency to worse-than-serial (#5455, #2476).
+                _remote_probe_cache["expires_at"] = (
+                    time.monotonic() + _REMOTE_PROBE_CACHE_TTL_S
+                )
+                _remote_probe_cache["result"] = payload
+            # Always release the in-flight marker and wake waiters, even if the
+            # probe raised — otherwise latecomers would wait out the full budget.
+            _remote_probe_inflight.discard(base_url)
+            _remote_probe_cond.notify_all()
+    # Only reached when _run_remote_probe returned normally (a raise would have
+    # propagated through the finally), so payload is always a dict here. The
+    # assert makes that explicit for the type checker (declared -> dict[str, Any]).
+    assert payload is not None
     return payload
 
 
 def _reset_remote_probe_cache_for_tests() -> None:
-    """Test hook: clear the in-process remote-probe cache."""
-    with _remote_probe_lock:
+    """Test hook: clear the in-process remote-probe cache and in-flight state."""
+    with _remote_probe_cond:
         _remote_probe_cache["url"] = None
         _remote_probe_cache["expires_at"] = 0.0
         _remote_probe_cache["result"] = None
+        _remote_probe_inflight.clear()
+        _remote_probe_cond.notify_all()
 
 
 def build_agent_health_payload() -> dict[str, Any]:

tests/test_5455_agent_health_single_flight.py

@@ -0,0 +1,173 @@
+"""Regression tests for single-flight remote-gateway health probing (#5455, #2476).
+
+Pre-fix behaviour: ``_probe_remote_gateway`` held ``_remote_probe_lock`` only
+around the cache read and the cache write; the network probe itself
+(``_http_probe`` per path, each up to ~2s) ran with no lock held.  On a cold
+cache a dashboard that fans out to N panels therefore fired N concurrent probe
+sets at the (possibly dead) gateway, each blocking for up to ~6s.
+
+Fix: single-flight.  The first "leader" thread marks the base_url in-flight and
+runs the probe; latecomers wait on ``_remote_probe_cond`` for the leader's
+cached result instead of probing themselves.  The leader always clears the
+marker and wakes waiters (even on error) so no waiter can deadlock.
+
+These tests pin that only ONE probe set runs under concurrency, that all callers
+observe the same result, and that a probe exception does not hang waiters.
+"""
+
+from __future__ import annotations
+
+import threading
+import time
+
+import pytest
+
+from api import agent_health
+
+
+@pytest.fixture(autouse=True)
+def _clear_cache():
+    agent_health._reset_remote_probe_cache_for_tests()
+    yield
+    agent_health._reset_remote_probe_cache_for_tests()
+
+
+def _run_concurrent(base_url: str, n: int, timeout: float = 20.0) -> list[object]:
+    """Fire *n* threads at ``_probe_remote_gateway`` simultaneously.
+
+    Returns each thread's result (or the exception it raised).  Uses a Barrier
+    so all threads reach the call together and genuinely contend on a cold cache.
+    """
+    barrier = threading.Barrier(n)
+    results: list[object] = [None] * n
+
+    def worker(idx: int) -> None:
+        barrier.wait()
+        try:
+            results[idx] = agent_health._probe_remote_gateway(base_url)
+        except BaseException as exc:  # noqa: BLE001 - recorded for the assertion
+            results[idx] = exc
+
+    threads = [threading.Thread(target=worker, args=(i,)) for i in range(n)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join(timeout)
+    # No thread may still be running — a hang would mean a single-flight deadlock.
+    assert not any(t.is_alive() for t in threads), "probe threads did not finish"
+    return results
+
+
+def test_concurrent_cold_cache_probes_only_once(monkeypatch):
+    base_url = "http://gw:9999"
+    call_count = 0
+    count_lock = threading.Lock()
+
+    def fake_http_probe(url, timeout_s, *, api_key=None):
+        nonlocal call_count
+        with count_lock:
+            call_count += 1
+        # Hold the leader in the probe long enough for latecomers to pile up and
+        # wait, so a lock-less stampede would show as call_count > 1.
+        time.sleep(0.3)
+        return (True, 200, None, b"{}")  # 2xx on the first path
+
+    monkeypatch.setattr(agent_health, "_http_probe", fake_http_probe)
+
+    results = _run_concurrent(base_url, n=8)
+
+    # Single-flight: exactly one leader ran the probe (first path returned 2xx).
+    assert call_count == 1, f"expected 1 probe, got {call_count}"
+    # Every caller observed the same alive result.
+    for r in results:
+        assert isinstance(r, dict), r
+        assert r["alive"] is True
+        assert r["details"]["endpoint"] == base_url + "/health/detailed"
+
+
+def test_concurrent_down_gateway_probes_paths_once_each(monkeypatch):
+    base_url = "http://gw:9998"
+    call_count = 0
+    count_lock = threading.Lock()
+
+    def fake_http_probe(url, timeout_s, *, api_key=None):
+        nonlocal call_count
+        with count_lock:
+            call_count += 1
+        time.sleep(0.2)
+        return (False, None, "URLError", None)  # never ok -> walks all paths
+
+    monkeypatch.setattr(agent_health, "_http_probe", fake_http_probe)
+
+    results = _run_concurrent(base_url, n=6)
+
+    # One leader walked every path exactly once; latecomers did not re-probe.
+    assert call_count == len(agent_health._REMOTE_PROBE_PATHS), (
+        f"expected {len(agent_health._REMOTE_PROBE_PATHS)} probes, got {call_count}"
+    )
+    for r in results:
+        assert isinstance(r, dict), r
+        assert r["alive"] is False
+        assert r["details"]["reason"] == "remote_gateway_unreachable"
+
+
+def test_slow_walk_exceeding_ttl_still_single_flight(monkeypatch):
+    """A probe walk LONGER than the cache TTL must not collapse single-flight.
+
+    Regression for a born-expired cache: the leader wrote
+    ``expires_at = entry_time + TTL``.  Walking every path of a hung gateway
+    takes len(paths) * timeout, which exceeds the TTL, so the cache line was
+    already stale the instant the leader stored it.  Latecomers woken right
+    after then missed the cache and each re-probed the dead gateway — the exact
+    #5455/#2476 fan-out the fix is meant to prevent (and worse: serialized).
+
+    Here the per-path sleep makes the total walk exceed a shrunk TTL.  The fix
+    expires from the COMPLETION time, so the freshly written result stays valid
+    long enough for every waiter to read it: still exactly one probe set.
+    """
+    base_url = "http://gw:9996"
+    # Walk = len(paths) * 0.25s; keep TTL well below that but comfortably above
+    # thread wake/reacquire latency so the completion-time cache is readable.
+    monkeypatch.setattr(agent_health, "_REMOTE_PROBE_CACHE_TTL_S", 0.4)
+    call_count = 0
+    count_lock = threading.Lock()
+
+    def fake_http_probe(url, timeout_s, *, api_key=None):
+        nonlocal call_count
+        with count_lock:
+            call_count += 1
+        time.sleep(0.25)
+        return (False, None, "URLError", None)  # down -> leader walks all paths
+
+    monkeypatch.setattr(agent_health, "_http_probe", fake_http_probe)
+
+    results = _run_concurrent(base_url, n=8)
+
+    # One leader walked every path once; no waiter re-probed despite walk > TTL.
+    assert call_count == len(agent_health._REMOTE_PROBE_PATHS), (
+        f"born-expired cache regressed single-flight: expected "
+        f"{len(agent_health._REMOTE_PROBE_PATHS)} probes, got {call_count}"
+    )
+    for r in results:
+        assert isinstance(r, dict), r
+        assert r["alive"] is False
+
+
+def test_probe_exception_wakes_waiters_without_deadlock(monkeypatch):
+    """A crash inside the leader's probe must not leave waiters hung on the cond."""
+    base_url = "http://gw:9997"
+    boom = RuntimeError("simulated probe crash")
+
+    def fake_http_probe(url, timeout_s, *, api_key=None):
+        time.sleep(0.2)
+        raise boom
+
+    monkeypatch.setattr(agent_health, "_http_probe", fake_http_probe)
+
+    results = _run_concurrent(base_url, n=5)
+
+    # Every caller surfaced the error (no result cached); crucially, the
+    # _run_concurrent join assertion already proved none of them deadlocked.
+    for r in results:
+        assert isinstance(r, RuntimeError), r
+        assert "simulated probe crash" in str(r)