diff --git a/sentinel/enrichment/circuit_breaker.py b/sentinel/enrichment/circuit_breaker.py index 0e6f628..9fae868 100644 --- a/sentinel/enrichment/circuit_breaker.py +++ b/sentinel/enrichment/circuit_breaker.py @@ -47,6 +47,9 @@ def __init__( self._state: State = "closed" self._opened_at: float | None = None self._lock = asyncio.Lock() + # True while a single half-open trial is in flight; gates concurrent + # callers so only one probes the unhealthy dependency (see `call`). + self._half_open_in_flight = False @property def name(self) -> str: @@ -75,6 +78,7 @@ def _prune(self, now: float) -> None: self._failures.popleft() async def call(self, fn: Callable[[], Awaitable[T]]) -> T: + probing = False async with self._lock: now = self._time() if self._state == "open": @@ -85,12 +89,29 @@ async def call(self, fn: Callable[[], Awaitable[T]]) -> T: self._transition("half_open") else: raise CircuitOpenError(f"breaker {self._name} open") + if self._state == "half_open": + # Single-flight the trial: only the first post-cooldown caller + # probes the unhealthy dependency. Concurrent callers short-circuit + # until that probe closes or re-opens the breaker, instead of all + # hammering the dependency at once. + if self._half_open_in_flight: + raise CircuitOpenError(f"breaker {self._name} half-open probe in flight") + self._half_open_in_flight = True + probing = True try: result = await fn() except asyncio.CancelledError: + # Release the probe slot without counting a failure. A bare assignment + # has no await point, so it is race-free even off-lock — important here + # because acquiring the lock during cancellation could itself be + # cancelled and leak the slot, wedging the breaker half-open forever. + if probing: + self._half_open_in_flight = False raise except BaseException: async with self._lock: + if probing: + self._half_open_in_flight = False now2 = self._time() if self._state == "half_open": self._failures.append(now2) @@ -103,6 +124,8 @@ async def call(self, fn: Callable[[], Awaitable[T]]) -> T: raise else: async with self._lock: + if probing: + self._half_open_in_flight = False if self._state == "half_open": self._transition("closed") return result diff --git a/tests/unit/enrichment/test_circuit_breaker.py b/tests/unit/enrichment/test_circuit_breaker.py index e0e6ac8..7e1f278 100644 --- a/tests/unit/enrichment/test_circuit_breaker.py +++ b/tests/unit/enrichment/test_circuit_breaker.py @@ -130,6 +130,117 @@ async def cancelled() -> None: assert cb.state == "closed" +@pytest.mark.asyncio +async def test_half_open_single_flights_concurrent_probes() -> None: + """In half_open, exactly ONE caller probes the unhealthy dependency; callers + arriving while the probe is in flight short-circuit with CircuitOpenError + instead of all hammering the dependency (issue #65).""" + clock = _Clock() + cb = CircuitBreaker("x", threshold=1, window_s=60.0, cooldown_s=30.0, time_fn=clock) + + async def boom() -> None: + raise RuntimeError("boom") + + with pytest.raises(RuntimeError): + await cb.call(boom) # threshold=1 → opens immediately + assert cb.state == "open" + clock.advance(31.0) # cooldown elapsed → next caller may probe + + probe_in_flight = asyncio.Event() + release = asyncio.Event() + probe_calls = 0 + + async def slow_probe() -> str: + nonlocal probe_calls + probe_calls += 1 + probe_in_flight.set() + await release.wait() + return "ok" + + async def must_not_run() -> str: # pragma: no cover - asserts it never runs + raise AssertionError("concurrent caller should have been short-circuited") + + probe_task = asyncio.create_task(cb.call(slow_probe)) + await asyncio.wait_for(probe_in_flight.wait(), timeout=1.0) + + # While the single probe is in flight, concurrent callers must be rejected. + results = await asyncio.gather( + cb.call(must_not_run), cb.call(must_not_run), return_exceptions=True + ) + assert all(isinstance(r, CircuitOpenError) for r in results), results + + release.set() + assert await probe_task == "ok" + assert probe_calls == 1 + assert cb.state == "closed" # type: ignore[comparison-overlap] + + +@pytest.mark.asyncio +async def test_half_open_probe_reopens_then_allows_next_probe_after_cooldown() -> None: + """Guards the *failure-path* flag release: a failed half-open probe re-opens + the breaker AND clears the in-flight flag, so a fresh probe is admitted after + the next cooldown. (Drop the flag-clear from the failure branch and the final + call short-circuits with CircuitOpenError instead of probing.)""" + clock = _Clock() + cb = CircuitBreaker("x", threshold=1, window_s=60.0, cooldown_s=30.0, time_fn=clock) + + async def boom() -> None: + raise RuntimeError("boom") + + async def ok() -> int: + return 1 + + with pytest.raises(RuntimeError): + await cb.call(boom) + clock.advance(31.0) + # Single-flight probe fails → reopen (and in-flight flag must be released). + with pytest.raises(RuntimeError): + await cb.call(boom) + assert cb.state == "open" + # Next cooldown → a fresh probe is allowed (proves the flag was cleared). + clock.advance(31.0) + assert await cb.call(ok) == 1 + assert cb.state == "closed" # type: ignore[comparison-overlap] + + +@pytest.mark.asyncio +async def test_half_open_cancelled_probe_releases_slot() -> None: + """Guards the cancellation-path flag release: if the single half-open probe + is cancelled (e.g. the caller's task is torn down), the in-flight slot must be + freed so a later caller can probe — never wedged half-open forever (issue #65). + """ + clock = _Clock() + cb = CircuitBreaker("x", threshold=1, window_s=60.0, cooldown_s=30.0, time_fn=clock) + + async def boom() -> None: + raise RuntimeError("boom") + + async def ok() -> int: + return 1 + + with pytest.raises(RuntimeError): + await cb.call(boom) + clock.advance(31.0) + + probe_in_flight = asyncio.Event() + + async def blocking_probe() -> None: + probe_in_flight.set() + await asyncio.Event().wait() # block until cancelled + + probe_task = asyncio.create_task(cb.call(blocking_probe)) + await asyncio.wait_for(probe_in_flight.wait(), timeout=1.0) + probe_task.cancel() + with pytest.raises(asyncio.CancelledError): + await probe_task + + # State is still half_open (cancellation isn't a failure); only the flag gates + # the next caller. If it leaked True this would raise CircuitOpenError. + assert cb.state == "half_open" + assert await cb.call(ok) == 1 + assert cb.state == "closed" # type: ignore[comparison-overlap] + + @pytest.mark.asyncio async def test_state_change_callback_fires_on_transitions() -> None: clock = _Clock()