diff --git a/api/agent_health.py b/api/agent_health.py index decb9675..276ca0d6 100644 --- a/api/agent_health.py +++ b/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]: diff --git a/tests/test_5455_agent_health_single_flight.py b/tests/test_5455_agent_health_single_flight.py new file mode 100644 index 00000000..a6cc1fb0 --- /dev/null +++ b/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)