From be60e06a7a93ee5840daa948eb810acd469dd628 Mon Sep 17 00:00:00 2001 From: ASHMEET555 Date: Thu, 29 Oct 2026 21:53:52 +0530 Subject: [PATCH 1/4] fix(circuit_breaker): await async coroutines inside circuit breaker failure tracking --- .../utilities/circuit_breaker/base.py | 116 +++++++++++++++ .../circuit_breaker/circuit_breaker.py | 28 +++- .../circuit_breaker/test_circuit_breaker.py | 136 ++++++++++++++++++ 3 files changed, 279 insertions(+), 1 deletion(-) diff --git a/aws_lambda_powertools/utilities/circuit_breaker/base.py b/aws_lambda_powertools/utilities/circuit_breaker/base.py index 3458a67aae8..c01e1215d40 100644 --- a/aws_lambda_powertools/utilities/circuit_breaker/base.py +++ b/aws_lambda_powertools/utilities/circuit_breaker/base.py @@ -9,6 +9,7 @@ from __future__ import annotations import datetime as dt +import inspect import logging import os import threading @@ -271,6 +272,121 @@ def _open_response(self, circuit: CircuitInfo) -> Any: circuit=circuit, ) + async def _open_response_async(self, circuit: CircuitInfo) -> Any: + """Produce the response for an open circuit in an async context.""" + if self.on_circuit_open is not None: + res = self.on_circuit_open(*self.fn_args, **self.fn_kwargs, circuit=circuit) + if inspect.iscoroutine(res): + return await res + return res + raise CircuitBreakerOpenError( + f"Circuit '{self.name}' is open.", + circuit=circuit, + ) + + async def handle_async(self) -> Any: + """ + Evaluate the circuit and route the call asynchronously. + + Returns + ------- + Any + The protected async function's result when the call runs, or the + ``on_circuit_open`` callback's return value when the circuit is open. + + Raises + ------ + CircuitBreakerOpenError + If the circuit is open and no callback is registered. + """ + record = self.persistence_store.get_state(self.name) + + if record.state == CircuitState.CLOSED: + with _COUNTERS_LOCK: + prev = _LAST_OBSERVED_STATE.get(self.name) + if prev is not None and prev != CircuitState.CLOSED: + _LOCAL_FAILURES[self.name] = 0 + _LAST_OBSERVED_STATE[self.name] = CircuitState.CLOSED + return await self._call_closed_async() + + if record.state == CircuitState.OPEN: + with _COUNTERS_LOCK: + _LAST_OBSERVED_STATE[self.name] = CircuitState.OPEN + opened_at = record.opened_at if record.opened_at is not None else self._now() + if self._now() >= opened_at + self.config.recovery_timeout: + 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 await self._call_probe_async() + return await self._open_response_async(record.to_circuit_info()) + + with _COUNTERS_LOCK: + _LAST_OBSERVED_STATE[self.name] = CircuitState.HALF_OPEN + if record.half_open_owner == _probe_owner_id(): + return await self._call_probe_async() + + if record.probe_lease_expiry is not None and self._now() >= record.probe_lease_expiry: + 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 await self._call_probe_async() + + return await self._open_response_async(record.to_circuit_info()) + + async def _call_closed_async(self) -> Any: + """Run the protected async call while the circuit is closed, tracking failures.""" + try: + result = await self.function(*self.fn_args, **self.fn_kwargs) + except Exception as exc: + if not self.config.counts_as_failure(exc): + raise + with _COUNTERS_LOCK: + failures = _LOCAL_FAILURES.get(self.name, 0) + 1 + tripped = failures >= self.config.failure_threshold + _LOCAL_FAILURES[self.name] = 0 if tripped else failures + if tripped: + logger.debug("Circuit '%s' tripping CLOSED to OPEN after %d failures.", self.name, failures) + opened_at = self._now() + self._safe_persist( + self.persistence_store.save_open, + self.name, + failure_count=failures, + opened_at=opened_at, + ) + self._notify(CircuitState.CLOSED, CircuitState.OPEN, opened_at=opened_at) + raise + else: + with _COUNTERS_LOCK: + _LOCAL_FAILURES[self.name] = 0 + return result + + async def _call_probe_async(self) -> Any: + """Run an async probe during half-open, closing or reopening based on the outcome.""" + try: + result = await self.function(*self.fn_args, **self.fn_kwargs) + except Exception as exc: + if not self.config.counts_as_failure(exc): + raise + logger.debug("Circuit '%s' probe failed; reopening.", self.name) + opened_at = self._now() + self._safe_persist(self.persistence_store.save_reopen, self.name, opened_at=opened_at) + with _COUNTERS_LOCK: + _LOCAL_SUCCESSES[self.name] = 0 + self._notify(CircuitState.HALF_OPEN, CircuitState.OPEN, opened_at=opened_at) + raise + else: + with _COUNTERS_LOCK: + successes = _LOCAL_SUCCESSES.get(self.name, 0) + 1 + closed = successes >= self.config.success_threshold + _LOCAL_SUCCESSES[self.name] = 0 if closed else successes + if closed: + _LOCAL_FAILURES[self.name] = 0 + if closed: + logger.debug("Circuit '%s' closing after %d probe successes.", self.name, successes) + self._safe_persist(self.persistence_store.save_closed, self.name) + self._notify(CircuitState.HALF_OPEN, CircuitState.CLOSED) + return result + def _notify(self, from_state: CircuitState, to_state: CircuitState, opened_at: int | None = None) -> None: """ Fire the ``on_transition`` hook for a state change. diff --git a/aws_lambda_powertools/utilities/circuit_breaker/circuit_breaker.py b/aws_lambda_powertools/utilities/circuit_breaker/circuit_breaker.py index f0d65542607..c3d5bf9d591 100644 --- a/aws_lambda_powertools/utilities/circuit_breaker/circuit_breaker.py +++ b/aws_lambda_powertools/utilities/circuit_breaker/circuit_breaker.py @@ -5,6 +5,7 @@ from __future__ import annotations import functools +import inspect import logging import os import warnings @@ -90,6 +91,31 @@ def charge(order: dict) -> dict: config = config or CircuitBreakerConfig() def decorator(function: Callable) -> Callable: + is_async = inspect.iscoroutinefunction(function) + + @functools.wraps(function) + async def async_wrapper(*args, **kwargs) -> Any: + if strtobool(os.getenv(constants.CIRCUIT_BREAKER_DISABLED_ENV, "false")): + warnings.warn( + message="Disabling the circuit breaker is intended for development environments only " + "and should not be used in production.", + category=PowertoolsUserWarning, + stacklevel=2, + ) + return await function(*args, **kwargs) + + handler = CircuitBreakerHandler( + function=function, + name=name, + config=config, + persistence_store=persistence_store, + on_circuit_open=on_circuit_open, + on_transition=on_transition, + function_args=args, + function_kwargs=kwargs, + ) + return await handler.handle_async() + @functools.wraps(function) def wrapper(*args, **kwargs) -> Any: # Skip the circuit entirely when disabled (development only). @@ -114,6 +140,6 @@ def wrapper(*args, **kwargs) -> Any: ) return handler.handle() - return wrapper + return async_wrapper if is_async else wrapper return decorator diff --git a/tests/functional/circuit_breaker/test_circuit_breaker.py b/tests/functional/circuit_breaker/test_circuit_breaker.py index 75b39fd5c9f..be73b4a71bc 100644 --- a/tests/functional/circuit_breaker/test_circuit_breaker.py +++ b/tests/functional/circuit_breaker/test_circuit_breaker.py @@ -888,3 +888,139 @@ def test_settings_are_kept_per_circuit_on_a_shared_layer(store, now): lease = store.db["a"].probe_lease_expiry assert lease is not None assert lease <= now + 30 + 2, "the lease must derive from circuit 'a' recovery_timeout, not 'b'" + + +@pytest.mark.asyncio +async def test_async_closed_circuit_returns_result_and_writes_nothing(store): + @circuit_breaker(name="async-c", persistence_store=store, config=CircuitBreakerConfig(local_cache_max_age=0)) + async def call(value): + return value * 2 + + assert await call(21) == 42 + assert store.db == {}, "healthy path must not write to the store" + + +@pytest.mark.asyncio +async def test_async_trips_open_after_failure_threshold(store): + config = CircuitBreakerConfig(failure_threshold=2, local_cache_max_age=0) + + @circuit_breaker(name="async-trip", persistence_store=store, config=config) + async def failing_call(): + raise ConnectionError("downstream down") + + for _ in range(2): + with pytest.raises(ConnectionError): + await failing_call() + + assert store.db["async-trip"].state == CircuitState.OPEN + assert store.db["async-trip"].opened_at is not None + + +@pytest.mark.asyncio +async def test_async_open_with_callback_returns_callback_value(store, now): + store.db["async-open"] = CircuitStateRecord( + name="async-open", + state=CircuitState.OPEN, + failure_count=5, + opened_at=now, + ) + + def sync_on_open(order, circuit): + return {"buffered": order, "state": str(circuit.state)} + + config = CircuitBreakerConfig(recovery_timeout=9999, local_cache_max_age=0) + + @circuit_breaker(name="async-open", persistence_store=store, on_circuit_open=sync_on_open, config=config) + async def charge(order): + return f"charged {order}" + + result = await charge({"id": 1}) + assert result == {"buffered": {"id": 1}, "state": "OPEN"} + + +@pytest.mark.asyncio +async def test_async_open_with_async_callback_returns_callback_value(store, now): + store.db["async-open-2"] = CircuitStateRecord( + name="async-open-2", + state=CircuitState.OPEN, + failure_count=5, + opened_at=now, + ) + + async def async_on_open(order, circuit): + return {"async_buffered": order, "state": str(circuit.state)} + + config = CircuitBreakerConfig(recovery_timeout=9999, local_cache_max_age=0) + + @circuit_breaker(name="async-open-2", persistence_store=store, on_circuit_open=async_on_open, config=config) + async def charge(order): + return f"charged {order}" + + result = await charge({"id": 2}) + assert result == {"async_buffered": {"id": 2}, "state": "OPEN"} + + +@pytest.mark.asyncio +async def test_async_open_without_callback_raises_with_circuit_info(store, now): + store.db["async-err"] = CircuitStateRecord( + name="async-err", + state=CircuitState.OPEN, + failure_count=3, + opened_at=now, + ) + config = CircuitBreakerConfig(recovery_timeout=9999, local_cache_max_age=0) + + @circuit_breaker(name="async-err", persistence_store=store, config=config) + async def call(): + return "success" + + with pytest.raises(CircuitBreakerOpenError) as exc_info: + await call() + + assert exc_info.value.circuit.name == "async-err" + + +@pytest.mark.asyncio +async def test_async_half_open_probe_success_closes(store, now): + store.db["async-probe"] = CircuitStateRecord(name="async-probe", state=CircuitState.OPEN, opened_at=now - 100) + config = CircuitBreakerConfig(recovery_timeout=30, success_threshold=1, local_cache_max_age=0) + + @circuit_breaker(name="async-probe", persistence_store=store, config=config) + async def call(): + return "healed" + + res = await call() + assert res == "healed" + assert store.db["async-probe"].state == CircuitState.CLOSED + + +@pytest.mark.asyncio +async def test_async_half_open_probe_failure_reopens(store, now): + store.db["async-reopen"] = CircuitStateRecord(name="async-reopen", state=CircuitState.OPEN, opened_at=now - 100) + config = CircuitBreakerConfig(recovery_timeout=30, success_threshold=1, local_cache_max_age=0) + + @circuit_breaker(name="async-reopen", persistence_store=store, config=config) + async def call(): + raise ConnectionError("still down") + + with pytest.raises(ConnectionError): + await call() + + assert store.db["async-reopen"].state == CircuitState.OPEN + + +@pytest.mark.asyncio +async def test_async_disabled_env_bypasses_circuit(store, monkeypatch): + monkeypatch.setenv("POWERTOOLS_CIRCUIT_BREAKER_DISABLED", "true") + config = CircuitBreakerConfig(failure_threshold=1, local_cache_max_age=0) + + @circuit_breaker(name="async-disabled", persistence_store=store, config=config) + async def call(): + return "bypassed" + + with warnings.catch_warnings(record=True) as recorded: + warnings.simplefilter("always") + result = await call() + + assert result == "bypassed" + assert any("Disabling the circuit breaker is intended for development" in str(w.message) for w in recorded) From dd58df610ce47fd0d34ed85c9e45a5cfb73b180d Mon Sep 17 00:00:00 2001 From: Leandro Damascena Date: Thu, 29 Oct 2026 16:43:39 +0000 Subject: [PATCH 2/4] refactor(circuit_breaker): share the state machine between sync and async paths _route() decides CLOSED/PROBE/OPEN once; _on_closed_*/_on_probe_* hold the bookkeeping; the sync and async call methods only differ in how they invoke the protected function. Also dedupes the decorator wrappers and documents async support. --- .../utilities/circuit_breaker/base.py | 280 ++++++++---------- .../circuit_breaker/circuit_breaker.py | 61 ++-- docs/utilities/circuit_breaker.md | 10 + .../src/working_with_async_functions.py | 38 +++ 4 files changed, 193 insertions(+), 196 deletions(-) create mode 100644 examples/circuit_breaker/src/working_with_async_functions.py diff --git a/aws_lambda_powertools/utilities/circuit_breaker/base.py b/aws_lambda_powertools/utilities/circuit_breaker/base.py index c01e1215d40..c5047da281f 100644 --- a/aws_lambda_powertools/utilities/circuit_breaker/base.py +++ b/aws_lambda_powertools/utilities/circuit_breaker/base.py @@ -14,6 +14,7 @@ import os import threading import uuid +from enum import Enum from typing import TYPE_CHECKING, Any from aws_lambda_powertools.utilities.circuit_breaker.exceptions import CircuitBreakerOpenError @@ -51,6 +52,14 @@ _PROBE_OWNER = threading.local() +class _Route(Enum): + """How a call is dispatched after reading the persisted circuit state.""" + + CLOSED = "closed" + PROBE = "probe" + OPEN = "open" + + def _probe_owner_id() -> str: """ Return this thread's stable probe-owner identifier, minting it on first use. @@ -131,6 +140,30 @@ def handle(self) -> Any: CircuitBreakerOpenError If the circuit is open and no callback is registered. """ + route, circuit = self._route() + if route is _Route.CLOSED: + return self._call_closed() + if route is _Route.PROBE: + return self._call_probe() + return self._open_response(circuit) + + async def handle_async(self) -> Any: + """Async counterpart of :meth:`handle` for coroutine functions; same routing and bookkeeping.""" + route, circuit = self._route() + if route is _Route.CLOSED: + return await self._call_closed_async() + if route is _Route.PROBE: + return await self._call_probe_async() + return await self._open_response_async(circuit) + + def _route(self) -> tuple[_Route, CircuitInfo]: + """ + Read the persisted state and decide how this call runs. + + Shared by the sync and async paths so the state machine lives in one place. Side + effects (probe election, OPEN to HALF_OPEN notification) happen here; only the + protected call itself differs between the two paths. + """ record = self.persistence_store.get_state(self.name) if record.state == CircuitState.CLOSED: @@ -142,7 +175,7 @@ def handle(self) -> Any: if prev is not None and prev != CircuitState.CLOSED: _LOCAL_FAILURES[self.name] = 0 _LAST_OBSERVED_STATE[self.name] = CircuitState.CLOSED - return self._call_closed() + return _Route.CLOSED, record.to_circuit_info() if record.state == CircuitState.OPEN: with _COUNTERS_LOCK: @@ -154,23 +187,23 @@ def handle(self) -> Any: 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 self._call_probe() - return self._open_response(record.to_circuit_info()) + 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. with _COUNTERS_LOCK: _LAST_OBSERVED_STATE[self.name] = CircuitState.HALF_OPEN if record.half_open_owner == _probe_owner_id(): - return self._call_probe() + return _Route.PROBE, record.to_circuit_info() # 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: 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 self._call_probe() + return _Route.PROBE, record.to_circuit_info() - return self._open_response(record.to_circuit_info()) + return _Route.OPEN, record.to_circuit_info() def _reset_probe_successes(self) -> None: """ @@ -189,57 +222,90 @@ def _call_closed(self) -> Any: try: result = self.function(*self.fn_args, **self.fn_kwargs) except Exception as exc: - if not self.config.counts_as_failure(exc): - raise - # Increment and reset atomically so exactly one thread observes the threshold - # crossing; racing threads would otherwise lose increments (tripping late) or - # each persist the same transition. - with _COUNTERS_LOCK: - failures = _LOCAL_FAILURES.get(self.name, 0) + 1 - tripped = failures >= self.config.failure_threshold - _LOCAL_FAILURES[self.name] = 0 if tripped else failures - if tripped: - logger.debug("Circuit '%s' tripping CLOSED to OPEN after %d failures.", self.name, failures) - opened_at = self._now() - self._safe_persist( - self.persistence_store.save_open, - self.name, - failure_count=failures, - opened_at=opened_at, - ) - self._notify(CircuitState.CLOSED, CircuitState.OPEN, opened_at=opened_at) + self._on_closed_failure(exc) raise - else: - with _COUNTERS_LOCK: - _LOCAL_FAILURES[self.name] = 0 - return result + self._on_closed_success() + return result + + async def _call_closed_async(self) -> Any: + """Await the protected coroutine while the circuit is closed, tracking failures.""" + try: + result = await self.function(*self.fn_args, **self.fn_kwargs) + except Exception as exc: + self._on_closed_failure(exc) + raise + self._on_closed_success() + return result + + def _on_closed_failure(self, exc: Exception) -> None: + """Count a qualifying failure while closed and trip the circuit at the threshold.""" + if not self.config.counts_as_failure(exc): + return + # Increment and reset atomically so exactly one thread observes the threshold + # crossing; racing threads would otherwise lose increments (tripping late) or + # each persist the same transition. + with _COUNTERS_LOCK: + failures = _LOCAL_FAILURES.get(self.name, 0) + 1 + tripped = failures >= self.config.failure_threshold + _LOCAL_FAILURES[self.name] = 0 if tripped else failures + if tripped: + logger.debug("Circuit '%s' tripping CLOSED to OPEN after %d failures.", self.name, failures) + opened_at = self._now() + self._safe_persist( + self.persistence_store.save_open, + self.name, + failure_count=failures, + opened_at=opened_at, + ) + self._notify(CircuitState.CLOSED, CircuitState.OPEN, opened_at=opened_at) + + def _on_closed_success(self) -> None: + with _COUNTERS_LOCK: + _LOCAL_FAILURES[self.name] = 0 def _call_probe(self) -> Any: """Run a probe during half-open, closing or reopening based on the outcome.""" try: result = self.function(*self.fn_args, **self.fn_kwargs) except Exception as exc: - if not self.config.counts_as_failure(exc): - raise - logger.debug("Circuit '%s' probe failed; reopening.", self.name) - opened_at = self._now() - self._safe_persist(self.persistence_store.save_reopen, self.name, opened_at=opened_at) - with _COUNTERS_LOCK: - _LOCAL_SUCCESSES[self.name] = 0 - self._notify(CircuitState.HALF_OPEN, CircuitState.OPEN, opened_at=opened_at) + self._on_probe_failure(exc) raise - else: - with _COUNTERS_LOCK: - successes = _LOCAL_SUCCESSES.get(self.name, 0) + 1 - closed = successes >= self.config.success_threshold - _LOCAL_SUCCESSES[self.name] = 0 if closed else successes - if closed: - _LOCAL_FAILURES[self.name] = 0 + self._on_probe_success() + return result + + async def _call_probe_async(self) -> Any: + """Await a probe during half-open, closing or reopening based on the outcome.""" + try: + result = await self.function(*self.fn_args, **self.fn_kwargs) + except Exception as exc: + self._on_probe_failure(exc) + raise + self._on_probe_success() + return result + + def _on_probe_failure(self, exc: Exception) -> None: + """A qualifying probe failure reopens the circuit.""" + if not self.config.counts_as_failure(exc): + return + logger.debug("Circuit '%s' probe failed; reopening.", self.name) + opened_at = self._now() + self._safe_persist(self.persistence_store.save_reopen, self.name, opened_at=opened_at) + with _COUNTERS_LOCK: + _LOCAL_SUCCESSES[self.name] = 0 + self._notify(CircuitState.HALF_OPEN, CircuitState.OPEN, opened_at=opened_at) + + def _on_probe_success(self) -> None: + """Count a probe success and close the circuit at the threshold.""" + with _COUNTERS_LOCK: + successes = _LOCAL_SUCCESSES.get(self.name, 0) + 1 + closed = successes >= self.config.success_threshold + _LOCAL_SUCCESSES[self.name] = 0 if closed else successes if closed: - logger.debug("Circuit '%s' closing after %d probe successes.", self.name, successes) - self._safe_persist(self.persistence_store.save_closed, self.name) - self._notify(CircuitState.HALF_OPEN, CircuitState.CLOSED) - return result + _LOCAL_FAILURES[self.name] = 0 + if closed: + logger.debug("Circuit '%s' closing after %d probe successes.", self.name, successes) + self._safe_persist(self.persistence_store.save_closed, self.name) + self._notify(CircuitState.HALF_OPEN, CircuitState.CLOSED) def _safe_persist(self, fn: Callable, *args: Any, **kwargs: Any) -> None: """ @@ -273,119 +339,11 @@ def _open_response(self, circuit: CircuitInfo) -> Any: ) async def _open_response_async(self, circuit: CircuitInfo) -> Any: - """Produce the response for an open circuit in an async context.""" - if self.on_circuit_open is not None: - res = self.on_circuit_open(*self.fn_args, **self.fn_kwargs, circuit=circuit) - if inspect.iscoroutine(res): - return await res - return res - raise CircuitBreakerOpenError( - f"Circuit '{self.name}' is open.", - circuit=circuit, - ) - - async def handle_async(self) -> Any: - """ - Evaluate the circuit and route the call asynchronously. - - Returns - ------- - Any - The protected async function's result when the call runs, or the - ``on_circuit_open`` callback's return value when the circuit is open. - - Raises - ------ - CircuitBreakerOpenError - If the circuit is open and no callback is registered. - """ - record = self.persistence_store.get_state(self.name) - - if record.state == CircuitState.CLOSED: - with _COUNTERS_LOCK: - prev = _LAST_OBSERVED_STATE.get(self.name) - if prev is not None and prev != CircuitState.CLOSED: - _LOCAL_FAILURES[self.name] = 0 - _LAST_OBSERVED_STATE[self.name] = CircuitState.CLOSED - return await self._call_closed_async() - - if record.state == CircuitState.OPEN: - with _COUNTERS_LOCK: - _LAST_OBSERVED_STATE[self.name] = CircuitState.OPEN - opened_at = record.opened_at if record.opened_at is not None else self._now() - if self._now() >= opened_at + self.config.recovery_timeout: - 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 await self._call_probe_async() - return await self._open_response_async(record.to_circuit_info()) - - with _COUNTERS_LOCK: - _LAST_OBSERVED_STATE[self.name] = CircuitState.HALF_OPEN - if record.half_open_owner == _probe_owner_id(): - return await self._call_probe_async() - - if record.probe_lease_expiry is not None and self._now() >= record.probe_lease_expiry: - 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 await self._call_probe_async() - - return await self._open_response_async(record.to_circuit_info()) - - async def _call_closed_async(self) -> Any: - """Run the protected async call while the circuit is closed, tracking failures.""" - try: - result = await self.function(*self.fn_args, **self.fn_kwargs) - except Exception as exc: - if not self.config.counts_as_failure(exc): - raise - with _COUNTERS_LOCK: - failures = _LOCAL_FAILURES.get(self.name, 0) + 1 - tripped = failures >= self.config.failure_threshold - _LOCAL_FAILURES[self.name] = 0 if tripped else failures - if tripped: - logger.debug("Circuit '%s' tripping CLOSED to OPEN after %d failures.", self.name, failures) - opened_at = self._now() - self._safe_persist( - self.persistence_store.save_open, - self.name, - failure_count=failures, - opened_at=opened_at, - ) - self._notify(CircuitState.CLOSED, CircuitState.OPEN, opened_at=opened_at) - raise - else: - with _COUNTERS_LOCK: - _LOCAL_FAILURES[self.name] = 0 - return result - - async def _call_probe_async(self) -> Any: - """Run an async probe during half-open, closing or reopening based on the outcome.""" - try: - result = await self.function(*self.fn_args, **self.fn_kwargs) - except Exception as exc: - if not self.config.counts_as_failure(exc): - raise - logger.debug("Circuit '%s' probe failed; reopening.", self.name) - opened_at = self._now() - self._safe_persist(self.persistence_store.save_reopen, self.name, opened_at=opened_at) - with _COUNTERS_LOCK: - _LOCAL_SUCCESSES[self.name] = 0 - self._notify(CircuitState.HALF_OPEN, CircuitState.OPEN, opened_at=opened_at) - raise - else: - with _COUNTERS_LOCK: - successes = _LOCAL_SUCCESSES.get(self.name, 0) + 1 - closed = successes >= self.config.success_threshold - _LOCAL_SUCCESSES[self.name] = 0 if closed else successes - if closed: - _LOCAL_FAILURES[self.name] = 0 - if closed: - logger.debug("Circuit '%s' closing after %d probe successes.", self.name, successes) - self._safe_persist(self.persistence_store.save_closed, self.name) - self._notify(CircuitState.HALF_OPEN, CircuitState.CLOSED) - return result + """Produce the response for an open circuit, awaiting an async ``on_circuit_open`` callback.""" + response = self._open_response(circuit) + if inspect.isawaitable(response): + return await response + return response def _notify(self, from_state: CircuitState, to_state: CircuitState, opened_at: int | None = None) -> None: """ diff --git a/aws_lambda_powertools/utilities/circuit_breaker/circuit_breaker.py b/aws_lambda_powertools/utilities/circuit_breaker/circuit_breaker.py index c3d5bf9d591..7c28c04d05e 100644 --- a/aws_lambda_powertools/utilities/circuit_breaker/circuit_breaker.py +++ b/aws_lambda_powertools/utilities/circuit_breaker/circuit_breaker.py @@ -91,20 +91,20 @@ def charge(order: dict) -> dict: config = config or CircuitBreakerConfig() def decorator(function: Callable) -> Callable: - is_async = inspect.iscoroutinefunction(function) + def disabled() -> bool: + # Skip the circuit entirely when disabled (development only). + if not strtobool(os.getenv(constants.CIRCUIT_BREAKER_DISABLED_ENV, "false")): + return False + warnings.warn( + message="Disabling the circuit breaker is intended for development environments only " + "and should not be used in production.", + category=PowertoolsUserWarning, + stacklevel=3, + ) + return True - @functools.wraps(function) - async def async_wrapper(*args, **kwargs) -> Any: - if strtobool(os.getenv(constants.CIRCUIT_BREAKER_DISABLED_ENV, "false")): - warnings.warn( - message="Disabling the circuit breaker is intended for development environments only " - "and should not be used in production.", - category=PowertoolsUserWarning, - stacklevel=2, - ) - return await function(*args, **kwargs) - - handler = CircuitBreakerHandler( + def build_handler(args: tuple, kwargs: dict) -> CircuitBreakerHandler: + return CircuitBreakerHandler( function=function, name=name, config=config, @@ -114,32 +114,23 @@ async def async_wrapper(*args, **kwargs) -> Any: function_args=args, function_kwargs=kwargs, ) - return await handler.handle_async() + + if inspect.iscoroutinefunction(function): + + @functools.wraps(function) + async def async_wrapper(*args, **kwargs) -> Any: + if disabled(): + return await function(*args, **kwargs) + return await build_handler(args, kwargs).handle_async() + + return async_wrapper @functools.wraps(function) def wrapper(*args, **kwargs) -> Any: - # Skip the circuit entirely when disabled (development only). - if strtobool(os.getenv(constants.CIRCUIT_BREAKER_DISABLED_ENV, "false")): - warnings.warn( - message="Disabling the circuit breaker is intended for development environments only " - "and should not be used in production.", - category=PowertoolsUserWarning, - stacklevel=2, - ) + if disabled(): return function(*args, **kwargs) + return build_handler(args, kwargs).handle() - handler = CircuitBreakerHandler( - function=function, - name=name, - config=config, - persistence_store=persistence_store, - on_circuit_open=on_circuit_open, - on_transition=on_transition, - function_args=args, - function_kwargs=kwargs, - ) - return handler.handle() - - return async_wrapper if is_async else wrapper + return wrapper return decorator diff --git a/docs/utilities/circuit_breaker.md b/docs/utilities/circuit_breaker.md index ab219c942c3..b6637b506b3 100644 --- a/docs/utilities/circuit_breaker.md +++ b/docs/utilities/circuit_breaker.md @@ -121,6 +121,16 @@ There is no wrapper type to inspect. The contract is: | **Open**, `on_circuit_open` set | Whatever the callback returns | | **Open**, no callback | Raises `CircuitBreakerOpenError` (with the `CircuitInfo` attached) | +### Protecting async functions + +`@circuit_breaker` also wraps `async def` functions. The coroutine is awaited inside the circuit, so failures count towards `failure_threshold`, probes run on the awaited result, and an `on_circuit_open` callback can be a coroutine function too. Call it with `await` as you normally would. + +=== "working_with_async_functions.py" + + ```python hl_lines="17 22-28 29" + --8<-- "examples/circuit_breaker/src/working_with_async_functions.py" + ``` + ## Handling an open circuit ### With a callback diff --git a/examples/circuit_breaker/src/working_with_async_functions.py b/examples/circuit_breaker/src/working_with_async_functions.py new file mode 100644 index 00000000000..38b4e0aeafb --- /dev/null +++ b/examples/circuit_breaker/src/working_with_async_functions.py @@ -0,0 +1,38 @@ +import asyncio +import os + +import httpx + +from aws_lambda_powertools.utilities.circuit_breaker import CircuitBreakerConfig, CircuitInfo, circuit_breaker +from aws_lambda_powertools.utilities.circuit_breaker.persistence import ( + CircuitBreakerDynamoDBPersistence, +) +from aws_lambda_powertools.utilities.typing import LambdaContext + +table = os.getenv("CIRCUIT_BREAKER_TABLE", "") +persistence = CircuitBreakerDynamoDBPersistence(table_name=table) +config = CircuitBreakerConfig(failure_threshold=3, recovery_timeout=30) + + +async def cached_quote(symbol: str, circuit: CircuitInfo) -> dict: + # Async callbacks are awaited; sync callbacks work too. + return {"symbol": symbol, "price": None, "circuit": circuit.state.value} + + +@circuit_breaker( + name="quotes-api", + persistence_store=persistence, + config=config, + on_circuit_open=cached_quote, +) +async def get_quote(symbol: str) -> dict: + # The coroutine is awaited inside the circuit, so a failure raised here counts + # towards failure_threshold just like it does for a regular function. + async with httpx.AsyncClient(timeout=2) as client: + response = await client.get(f"https://quotes.example.com/{symbol}") + response.raise_for_status() + return response.json() + + +def lambda_handler(event: dict, context: LambdaContext): + return asyncio.run(get_quote(event["symbol"])) From 57aef795c536182a8ca43fc63e862e96fa9c203a Mon Sep 17 00:00:00 2001 From: Leandro Damascena Date: Thu, 29 Oct 2026 18:39:44 +0000 Subject: [PATCH 3/4] 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. --- .../utilities/circuit_breaker/base.py | 72 ++++++++++++++----- docs/utilities/circuit_breaker.md | 3 + tests/functional/circuit_breaker/conftest.py | 1 + .../circuit_breaker/test_circuit_breaker.py | 48 +++++++++++++ 4 files changed, 105 insertions(+), 19 deletions(-) diff --git a/aws_lambda_powertools/utilities/circuit_breaker/base.py b/aws_lambda_powertools/utilities/circuit_breaker/base.py index c5047da281f..28539349f0c 100644 --- a/aws_lambda_powertools/utilities/circuit_breaker/base.py +++ b/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 diff --git a/docs/utilities/circuit_breaker.md b/docs/utilities/circuit_breaker.md index b6637b506b3..66b49744650 100644 --- a/docs/utilities/circuit_breaker.md +++ b/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. diff --git a/tests/functional/circuit_breaker/conftest.py b/tests/functional/circuit_breaker/conftest.py index 30925dfc5e8..c9370d13831 100644 --- a/tests/functional/circuit_breaker/conftest.py +++ b/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() diff --git a/tests/functional/circuit_breaker/test_circuit_breaker.py b/tests/functional/circuit_breaker/test_circuit_breaker.py index be73b4a71bc..d6563361ca4 100644 --- a/tests/functional/circuit_breaker/test_circuit_breaker.py +++ b/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 From 6bcc9f3629a72234d769504981b2ab6eccd6dba8 Mon Sep 17 00:00:00 2001 From: Leandro Damascena Date: Thu, 29 Oct 2026 18:43:52 +0000 Subject: [PATCH 4/4] fix(circuit_breaker): ignore an in-flight probe entry inherited through fork --- .../utilities/circuit_breaker/base.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/aws_lambda_powertools/utilities/circuit_breaker/base.py b/aws_lambda_powertools/utilities/circuit_breaker/base.py index 28539349f0c..7f5b2dc1672 100644 --- a/aws_lambda_powertools/utilities/circuit_breaker/base.py +++ b/aws_lambda_powertools/utilities/circuit_breaker/base.py @@ -48,10 +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() +# Circuit names with a probe currently running in this process, mapped to the pid that +# claimed it. 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. The pid lets a forked child ignore an entry it inherited. +_PROBES_IN_FLIGHT: dict[str, int] = {} # 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. @@ -226,14 +227,14 @@ def _route_half_open(self, record: CircuitStateRecord) -> _Route: 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: + if _PROBES_IN_FLIGHT.get(self.name) == os.getpid(): return False - _PROBES_IN_FLIGHT.add(self.name) + _PROBES_IN_FLIGHT[self.name] = os.getpid() return True def _release_probe(self) -> None: with _COUNTERS_LOCK: - _PROBES_IN_FLIGHT.discard(self.name) + _PROBES_IN_FLIGHT.pop(self.name, None) def _reset_probe_successes(self) -> None: """