fix(circuit_breaker): run a single probe per process when coroutines call the circuit concurrently

The probe lease is per thread, so concurrent coroutines in one event loop all passed
the owner check and probed together. Track in-flight probes per circuit name and give
the other callers the open response, matching the single-prober guarantee threads
already had.

Author: leandrodamascena

aws_lambda_powertools/utilities/circuit_breaker/base.py

@@ -27,6 +27,7 @@
     from aws_lambda_powertools.utilities.circuit_breaker.persistence.base import (
         CircuitBreakerPersistenceLayer,
     )
+    from aws_lambda_powertools.utilities.circuit_breaker.persistence.record import CircuitStateRecord
     from aws_lambda_powertools.utilities.circuit_breaker.states import CircuitInfo
 
 logger = logging.getLogger(__name__)
@@ -47,6 +48,11 @@
 # callbacks.
 _COUNTERS_LOCK = threading.Lock()
 
+# Circuit names with a probe currently running in this process. The probe lease is per
+# thread, so concurrent coroutines in one event loop would all pass the owner check and
+# probe together; this keeps the single-prober guarantee within a process too.
+_PROBES_IN_FLIGHT: set[str] = set()
+
 # Identifier used to claim the half-open probe lock, unique per thread so the store's
 # conditional election picks a single prober across threads as well as processes.
 _PROBE_OWNER = threading.local()
@@ -175,35 +181,59 @@ def _route(self) -> tuple[_Route, CircuitInfo]:
                 if prev is not None and prev != CircuitState.CLOSED:
                     _LOCAL_FAILURES[self.name] = 0
                 _LAST_OBSERVED_STATE[self.name] = CircuitState.CLOSED
-            return _Route.CLOSED, record.to_circuit_info()
+            route = _Route.CLOSED
+        elif record.state == CircuitState.OPEN:
+            route = self._route_open(record)
+        else:
+            route = self._route_half_open(record)
 
-        if record.state == CircuitState.OPEN:
-            with _COUNTERS_LOCK:
-                _LAST_OBSERVED_STATE[self.name] = CircuitState.OPEN
-            # ``opened_at`` may legitimately be 0 (epoch); treat only None as missing.
-            opened_at = record.opened_at if record.opened_at is not None else self._now()
-            if self._now() >= opened_at + self.config.recovery_timeout:
-                # Recovery window elapsed: try to become the single prober.
-                if self.persistence_store.try_acquire_half_open(self.name, _probe_owner_id(), opened_at):
-                    self._reset_probe_successes()
-                    self._notify(CircuitState.OPEN, CircuitState.HALF_OPEN, opened_at=opened_at)
-                    return _Route.PROBE, record.to_circuit_info()
-            return _Route.OPEN, record.to_circuit_info()
-
-        # HALF_OPEN: only the thread that owns the probe lock runs.
+        return route, record.to_circuit_info()
+
+    def _route_open(self, record: CircuitStateRecord) -> _Route:
+        """OPEN: reject, or elect this worker as the single prober once the recovery window elapsed."""
+        with _COUNTERS_LOCK:
+            _LAST_OBSERVED_STATE[self.name] = CircuitState.OPEN
+        # ``opened_at`` may legitimately be 0 (epoch); treat only None as missing.
+        opened_at = record.opened_at if record.opened_at is not None else self._now()
+        if self._now() < opened_at + self.config.recovery_timeout or not self._claim_probe():
+            return _Route.OPEN
+        # Recovery window elapsed: try to become the single prober.
+        if self.persistence_store.try_acquire_half_open(self.name, _probe_owner_id(), opened_at):
+            self._reset_probe_successes()
+            self._notify(CircuitState.OPEN, CircuitState.HALF_OPEN, opened_at=opened_at)
+            return _Route.PROBE
+        self._release_probe()
+        return _Route.OPEN
+
+    def _route_half_open(self, record: CircuitStateRecord) -> _Route:
+        """HALF_OPEN: only the probe owner runs, one probe at a time; expired leases can be taken over."""
         with _COUNTERS_LOCK:
             _LAST_OBSERVED_STATE[self.name] = CircuitState.HALF_OPEN
         if record.half_open_owner == _probe_owner_id():
-            return _Route.PROBE, record.to_circuit_info()
+            # Owner thread: one probe at a time, even when coroutines run concurrently.
+            return _Route.PROBE if self._claim_probe() else _Route.OPEN
 
         # If the probe lease has expired (owner recycled mid-probe), take over.
-        if record.probe_lease_expiry is not None and self._now() >= record.probe_lease_expiry:
+        lease_expired = record.probe_lease_expiry is not None and self._now() >= record.probe_lease_expiry
+        if lease_expired and self._claim_probe():
             logger.debug("Circuit '%s' probe lease expired; attempting takeover.", self.name)
             if self.persistence_store.try_acquire_half_open(self.name, _probe_owner_id(), record.opened_at or 0):
                 self._reset_probe_successes()
-                return _Route.PROBE, record.to_circuit_info()
+                return _Route.PROBE
+            self._release_probe()
+        return _Route.OPEN
 
-        return _Route.OPEN, record.to_circuit_info()
+    def _claim_probe(self) -> bool:
+        """Mark a probe as in flight for this circuit; False when one is already running in this process."""
+        with _COUNTERS_LOCK:
+            if self.name in _PROBES_IN_FLIGHT:
+                return False
+            _PROBES_IN_FLIGHT.add(self.name)
+            return True
+
+    def _release_probe(self) -> None:
+        with _COUNTERS_LOCK:
+            _PROBES_IN_FLIGHT.discard(self.name)
 
     def _reset_probe_successes(self) -> None:
         """
@@ -270,6 +300,8 @@ def _call_probe(self) -> Any:
         except Exception as exc:
             self._on_probe_failure(exc)
             raise
+        finally:
+            self._release_probe()
         self._on_probe_success()
         return result
 
@@ -280,6 +312,8 @@ async def _call_probe_async(self) -> Any:
         except Exception as exc:
             self._on_probe_failure(exc)
             raise
+        finally:
+            self._release_probe()
         self._on_probe_success()
         return result
 

docs/utilities/circuit_breaker.md

@@ -191,6 +191,9 @@ After `recovery_timeout` seconds, the circuit moves to `HALF_OPEN` and elects a
     thread, so the single-prober guarantee spans threads as well as environments, and the in-memory failure counter is
     synchronized. Single-threaded functions (the normal Lambda model) are unaffected.
 
+    The same holds for coroutines: concurrent `await`s of the same circuit in one event loop run a single probe at a
+    time, and the others get the open response.
+
     Probe ownership belongs to the thread that won the election. If that thread never runs the circuit again (for
     example, a thread-per-request worker pool), recovery waits for the probe lease to expire before another thread or
     environment takes over.

tests/functional/circuit_breaker/conftest.py

@@ -102,6 +102,7 @@ def reset_local_counters():
     base_module._LOCAL_FAILURES.clear()
     base_module._LOCAL_SUCCESSES.clear()
     base_module._LAST_OBSERVED_STATE.clear()
+    base_module._PROBES_IN_FLIGHT.clear()
     yield
     base_module._LOCAL_FAILURES.clear()
     base_module._LOCAL_SUCCESSES.clear()

tests/functional/circuit_breaker/test_circuit_breaker.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import threading
 import time
 import warnings
@@ -1024,3 +1025,50 @@ async def call():
 
     assert result == "bypassed"
     assert any("Disabling the circuit breaker is intended for development" in str(w.message) for w in recorded)
+
+
+@pytest.mark.asyncio
+async def test_async_concurrent_calls_run_a_single_probe(store, now):
+    """Coroutines share the thread's probe owner id, so without a guard they would all probe at once."""
+    store.db["c"] = CircuitStateRecord(name="c", state=CircuitState.OPEN, failure_count=5, opened_at=now - 60)
+    config = CircuitBreakerConfig(recovery_timeout=30, success_threshold=2, local_cache_max_age=0)
+    probes = []
+
+    @circuit_breaker(name="c", persistence_store=store, config=config)
+    async def call(i):
+        probes.append(i)
+        await asyncio.sleep(0.01)
+        return "ok"
+
+    results = await asyncio.gather(*(call(i) for i in range(8)), return_exceptions=True)
+
+    assert len(probes) == 1, "exactly one coroutine must probe per recovery window"
+    assert sum(isinstance(r, CircuitBreakerOpenError) for r in results) == 7
+    assert store.db["c"].state == CircuitState.HALF_OPEN
+
+
+@pytest.mark.asyncio
+async def test_async_probe_slot_is_released_after_success_and_failure(store, now):
+    store.db["c"] = CircuitStateRecord(name="c", state=CircuitState.OPEN, failure_count=5, opened_at=now - 60)
+    config = CircuitBreakerConfig(recovery_timeout=30, success_threshold=3, local_cache_max_age=0)
+    calls = []
+
+    @circuit_breaker(name="c", persistence_store=store, config=config)
+    async def call(fail=False):
+        calls.append(fail)
+        if fail:
+            raise ValueError("probe failed")
+        return "ok"
+
+    # a successful probe keeps the circuit half-open and must leave the slot free for the next probe
+    assert await call() == "ok"
+    assert await call() == "ok"
+    assert store.db["c"].state == CircuitState.HALF_OPEN
+    # a failing probe reopens; once the recovery window elapses the next call must be able to probe again
+    with pytest.raises(ValueError):
+        await call(fail=True)
+    assert store.db["c"].state == CircuitState.OPEN
+    store.db["c"].opened_at = now - 60
+    assert await call() == "ok"
+    assert calls == [False, False, True, False]
+    assert "c" not in base_module._PROBES_IN_FLIGHT