STASH fixes to circuit breaker

Author: lobsterkatie

src/sentry/utils/circuit_breaker2.py

@@ -1,8 +1,7 @@
 import logging
 import time
 from enum import Enum
-from math import ceil
-from typing import Any, Literal, NotRequired, TypedDict, cast, overload
+from typing import Literal, NotRequired, TypedDict, overload
 
 from django.core.cache import cache
 
@@ -15,6 +14,9 @@
 
 logger = logging.getLogger(__name__)
 
+# How many times stricter to be with the error limit during recovery
+DEFAULT_RECOVERY_STRICTNESS = 10
+
 
 class CircuitBreakerState(Enum):
     CLOSED = "circuit_closed"
@@ -27,41 +29,35 @@ class CircuitBreakerConfig(TypedDict):
     error_limit: int
     # The time period, in seconds, over which we're tracking errors
     error_limit_window: int
-    # The length, in seconds, of each time bucket ("granule") used by the underlying rate limiter -
-    # effectively the resolution of the time window. Will be set automatically based on
-    # `error_limit_window` if not provided.
-    error_limit_window_granularity: NotRequired[int]
     # How long, in seconds, to stay in the broken state (blocking all requests) before entering the
     # recovery phase
     broken_state_duration: int
-    # The number of errors within the given time period necessary to trip the breaker while in recovery
-    recovery_error_limit: int
-    # The time period, in seconds, over which we're tracking errors in recovery
-    recovery_error_limit_window: int
+    # The number of errors within the given time period necessary to trip the breaker while in
+    # recovery. Will be set automatically based on `error_limit` if not provided
+    recovery_error_limit: NotRequired[int]
     # The length, in seconds, of each time bucket ("granule") used by the underlying rate limiter -
     # effectively the resolution of the time window. Will be set automatically based on
-    # `recovery_error_limit_window` if not provided.
-    recovery_error_limit_window_granularity: NotRequired[int]
+    # `error_limit_window` if not provided.
+    error_limit_window_granularity: NotRequired[int]
     # How long, in seconds, to stay in the recovery state (allowing requests but with a stricter
-    # error limit) before returning to normal operation.
-    recovery_duration: int
+    # error limit) before returning to normal operation. Will be set based on `error_limit_window`
+    # if not provided.
+    recovery_duration: NotRequired[int]
 
 
-# TODO: These limits were estimated based on EA traffic. (In an average 10 min period, there are
-# roughly 35K events without matching hashes. About 2% of orgs are EA, so for simplicity, assume 2%
-# of those events are from EA orgs. If we're willing to tolerate up to a 95% failure rate, then we
-# need 35K * 0.02 * 0.95 events to fail to trip the breaker. Technically that's 665, not 666, but
-# we're talking about everything going to hell, so the bump to 666 seemed appropriate!)
+# TODO: The default error limit here was estimated based on EA traffic. (In an average 10 min
+# period, there are roughly 35K events without matching hashes. About 2% of orgs are EA, so for
+# simplicity, assume 2% of those events are from EA orgs. If we're willing to tolerate up to a 95%
+# failure rate, then we need 35K * 0.02 * 0.95 events to fail to trip the breaker. Technically
+# that's 665, not 666, but we're talking about everything going to hell, so the bump to 666 seemed
+# appropriate!)
 #
 # When we GA, we should multiply both the limits by 50 (to remove the 2% part of the current
-# calculation).
+# calculation), and remove this TODO.
 CIRCUIT_BREAKER_DEFAULT_CONFIG: CircuitBreakerConfig = {
     "error_limit": 666,
     "error_limit_window": 600,  # 10 min
     "broken_state_duration": 300,  # 5 min
-    "recovery_error_limit": 3,  # In recovery, we're twice as strict as normal limit
-    "recovery_error_limit_window": 60,  # And we bail much more quickly
-    "recovery_duration": 300,  # 5 min
 }
 
 
@@ -108,66 +104,80 @@ def get_top_dogs(payload):
 
         return format_hof_entries(response)
 
-    The `breaker.should_allow_request()` check can alternatively be used outside of
-    `get_top_dogs`, to prevent calls to it. In that case, the original `breaker` object can be
-    imported alongside `get_top_dogs` or reinstantiated with the same config - it has no state of
-    its own, instead relying on redis and the cache to track error count and breaker status.
+    The `breaker.should_allow_request()` check can alternatively be used outside of `get_top_dogs`,
+    to prevent calls to it. In that case, the original `breaker` object can be imported alongside
+    `get_top_dogs` or reinstantiated with the same config - it has no state of its own, instead
+    relying on redis-backed rate limiters and the cache to track error count and breaker status.
     """
 
-    def __init__(self, key: str, config: CircuitBreakerConfig | None = None):
+    def __init__(self, key: str, config: CircuitBreakerConfig = CIRCUIT_BREAKER_DEFAULT_CONFIG):
         self.key = key
         self.broken_state_key = f"{key}.circuit_breaker.broken"
         self.recovery_state_key = f"{key}.circuit_breaker.in_recovery"
 
-        final_config: CircuitBreakerConfig = {
-            **CIRCUIT_BREAKER_DEFAULT_CONFIG,
-            **(config or cast(Any, {})),
-        }
-        default_window_granularity = self._get_default_window_granularity(
-            final_config["error_limit_window"]
+        self.error_limit = config["error_limit"]
+        self.recovery_error_limit = config.get(
+            "recovery_error_limit", max(self.error_limit // DEFAULT_RECOVERY_STRICTNESS, 1)
         )
-        default_recovery_window_granularity = self._get_default_window_granularity(
-            final_config["recovery_error_limit_window"]
+        self.window = config["error_limit_window"]
+        self.window_granularity = config.get(
+            "error_limit_window_granularity", max(self.window // 10, 5)
         )
+        self.broken_state_duration = config["broken_state_duration"]
+        self.recovery_duration = config.get("recovery_duration", self.window * 2)
 
         self.limiter = RedisSlidingWindowRateLimiter()
         self.primary_quota = Quota(
-            final_config["error_limit_window"],
-            final_config.get("error_limit_window_granularity", default_window_granularity),
-            final_config["error_limit"],
+            self.window,
+            self.window_granularity,
+            self.error_limit,
             f"{key}.circuit_breaker",
         )
         self.recovery_quota = Quota(
-            final_config["recovery_error_limit_window"],
-            final_config.get(
-                "recovery_error_limit_window_granularity", default_recovery_window_granularity
-            ),
-            final_config["recovery_error_limit"],
+            self.window,
+            self.window_granularity,
+            self.recovery_error_limit,
             f"{key}.circuit_breaker_recovery",
         )
 
-        self.broken_state_duration = final_config["broken_state_duration"]
-        self.recovery_duration = final_config["recovery_duration"]
-
-        if self.recovery_duration < final_config["recovery_error_limit_window"]:
+        if self.recovery_error_limit >= self.error_limit:
+            self.recovery_error_limit = max(self.error_limit // DEFAULT_RECOVERY_STRICTNESS, 1)
             logger.warning(
-                "Circuit breaker %s has `recovery_duration` < `recovery_error_limit_window`."
-                + " Recovery duration has been reset to match the window.",
+                "Circuit breaker '%s' has a recovery error limit (%d) greater than"
+                + " or equal to its error limit (%d). Using the stricter error-limit-based default"
+                + " (%d) instead.",
                 key,
+                config["recovery_error_limit"],
+                self.error_limit,
+                self.recovery_error_limit,
             )
-            self.recovery_duration = final_config["recovery_error_limit_window"]
-        if self.recovery_duration < final_config["recovery_error_limit_window"]:
+
+        # XXX: If we discover we have a config where we want this combo to work, we can consider
+        # using the `MockCircuitBreaker._clear_quota` helper, which is currently only used in tests,
+        # to clear out the main quota when we switch to the broken state. (It will need tests of its
+        # own if so.)
+        if self.broken_state_duration + self.recovery_duration < self.window:
             logger.warning(
-                "Circuit breaker %s has `recovery_duration` < `recovery_error_limit_window`."
-                + " Recovery duration has been reset to match the window.",
+                "Circuit breaker '%s' has broken state and recovery durations (%d and %d sec)"
+                + " which together are less than the main error limit window (%d sec). This"
+                + " can lead to the breaker getting tripped unexpectedly, until the original"
+                + " spike in errors clears the main time window.",
                 key,
+                self.broken_state_duration,
+                self.recovery_duration,
+                self.window,
             )
-            self.recovery_duration = final_config["recovery_error_limit_window"]
 
     def record_error(self) -> None:
+        """
+        Record a single error towards the breaker's quota, and handle the case where that error puts
+        us over the limit.
+        """
         state, seconds_left_in_state = self._get_state_and_remaining_time()
 
         if state == CircuitBreakerState.BROKEN:
+            assert seconds_left_in_state is not None  # mypy appeasement
+
             # If the circuit is broken, and `should_allow_request` is being used correctly, requests
             # should be blocked and we shouldn't even be here. That said, maybe there was a race
             # condition, so make sure the circuit hasn't just broken before crying foul.
@@ -222,6 +232,10 @@ def record_error(self) -> None:
             cache.set(self.recovery_state_key, recovery_state_expiry, recovery_state_timeout)
 
     def should_allow_request(self) -> bool:
+        """
+        Determine, based on the current state of the breaker and the number of allowable errors
+        remaining, whether requests should be allowed through.
+        """
         state, _ = self._get_state_and_remaining_time()
 
         if state == CircuitBreakerState.BROKEN:
@@ -255,8 +269,13 @@ def get_remaining_error_quota(
         self, quota: Quota | None = None, window_end: int | None = None
     ) -> int | None:
         """
-        # TODO: write me
-        returns None when in broken state
+        Get the number of allowable errors remaining in the given quota for the time window ending
+        at the given time.
+
+        If no quota is given, in closed and recovery states, return the current controlling quota's
+        remaining errors. In broken state, return None.
+
+        If no time window end is given, return the current amount of quota remaining.
         """
         if not quota:
             quota = self._get_controlling_quota()
@@ -273,12 +292,6 @@ def get_remaining_error_quota(
 
         return result[0].granted
 
-    def _get_default_window_granularity(self, window_duration: int) -> int:
-        # Never more than 10 buckets, and no bucket smaller than 5 seconds. If greater precision is
-        # needed, the `error_limit_window_granularity` and `recovery_error_limit_window_granularity`
-        # config options can be used.
-        return max(ceil(window_duration / 10), 5)
-
     @overload
     def _get_controlling_quota(
         self, state: Literal[CircuitBreakerState.CLOSED, CircuitBreakerState.RECOVERY]
@@ -295,8 +308,8 @@ def _get_controlling_quota(self) -> Quota | None:
 
     def _get_controlling_quota(self, state: CircuitBreakerState | None = None) -> Quota | None:
         """
-        # TODO: write me
-        returns None when in broken state
+        Return the Quota corresponding to the given breaker state (or the current breaker state, if
+        no state is provided). If the state is question is the broken state, return None.
         """
         controlling_quota_by_state = {
             CircuitBreakerState.CLOSED: self.primary_quota,
@@ -308,10 +321,13 @@ def _get_controlling_quota(self, state: CircuitBreakerState | None = None) -> Qu
 
         return controlling_quota_by_state[_state]
 
-    def _get_state_and_remaining_time(self) -> tuple[CircuitBreakerState, int]:
+    def _get_state_and_remaining_time(
+        self,
+    ) -> tuple[CircuitBreakerState, int | None]:
         """
         Return the current state of the breaker (closed, broken, or in recovery), along with the
-        number of seconds until that state expires.
+        number of seconds until that state expires (or `None` when in closed state, as it has no
+        expiry).
         """
         now = int(time.time())
 
@@ -322,7 +338,4 @@ def _get_state_and_remaining_time(self) -> tuple[CircuitBreakerState, int]:
         if cache.has_key(self.recovery_state_key):
             return (CircuitBreakerState.RECOVERY, cache.get(self.recovery_state_key) - now)
 
-        # TODO Fix this with overloads?
-        # 0 here is just a placeholder, as "remaining seconds" doesn't really apply to a state we
-        # hope to stay in indefinitely
-        return (CircuitBreakerState.CLOSED, 0)
+        return (CircuitBreakerState.CLOSED, None)