feat(apigateway): implement circuitbreakers in async variant

Author: gi0baro

src/sentry/conf/server.py

@@ -3183,6 +3183,9 @@ def custom_parameter_sort(parameter: dict) -> tuple[str, int]:
 }
 # Used in tests to skip forwarding relay paths to a region silo that does not exist.
 APIGATEWAY_PROXY_SKIP_RELAY = False
+APIGATEWAY_PROXY_MAX_CONCURRENCY = 100
+APIGATEWAY_PROXY_MAX_FAILURES = 100
+APIGATEWAY_PROXY_FAILURE_WINDOW = 60
 
 # Shared resource ids for accounting
 EVENT_PROCESSING_STORE = "rc_processing_redis"

src/sentry/hybridcloud/apigateway_async/circuitbreaker.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+import asyncio
+import time
+from collections import defaultdict
+
+from django.conf import settings
+
+
+class CircuitBreaker:
+    __slots__ = [
+        "concurrency",
+        "counter_window",
+        "failures",
+        "semaphore",
+        "_clock",
+        "_counters",
+        "_counter_idx",
+    ]
+
+    def __init__(self, concurrency, failures):
+        self.concurrency = concurrency
+        self.counter_window = failures[0]
+        self.failures = failures[1]
+        self.semaphore = asyncio.Semaphore(self.concurrency)
+        self._clock = 0
+        self._counters = [0, 0]
+        self._counter_idx = 0
+
+    def _counter_flip(self, clock):
+        self._clock = clock
+        prev = self._counter_idx
+        self._counter_idx = 1 - prev
+        self._counters[prev] = 0
+
+    def _maybe_counter_flip(self):
+        now = int(time.monotonic())
+        delta = now - self._clock
+        if delta > 0:
+            if delta // self.counter_window:
+                self._counter_flip(now)
+
+    def counter_incr(self):
+        self._maybe_counter_flip()
+        self._counters[self._counter_idx] += 1
+
+    def window_overflow(self) -> bool:
+        return self._counters[self._counter_idx] > self.failures
+
+    def overflow(self) -> bool:
+        return self.semaphore.locked()
+
+    def ctx(self) -> CircuitBreakerCtx:
+        return CircuitBreakerCtx(self)
+
+
+class CircuitBreakerOverflow(Exception): ...
+
+
+class CircuitBreakerWindowOverflow(Exception): ...
+
+
+class CircuitBreakerCtx:
+    __slots__ = ["cb"]
+
+    def __init__(self, cb: CircuitBreaker):
+        self.cb = cb
+
+    def incr_failures(self):
+        self.cb.counter_incr()
+
+    async def __aenter__(self):
+        if self.cb.overflow():
+            raise CircuitBreakerOverflow
+        await self.cb.semaphore.acquire()
+        if self.cb.window_overflow():
+            self.cb.semaphore.release()
+            raise CircuitBreakerWindowOverflow
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, exc_tb):
+        self.cb.semaphore.release()
+
+
+class CircuitBreakerManager:
+    __slots__ = ["objs"]
+
+    def __init__(
+        self,
+        max_concurrency: int | None = None,
+        failures: int | None = None,
+        failure_window: int | None = None,
+    ):
+        concurrency = max_concurrency or settings.APIGATEWAY_PROXY_MAX_CONCURRENCY
+        failures = failures or settings.APIGATEWAY_PROXY_MAX_FAILURES
+        failure_window = failure_window or settings.APIGATEWAY_PROXY_FAILURE_WINDOW
+        self.objs = defaultdict(lambda: CircuitBreaker(concurrency, (failure_window, failures)))
+
+    def get(self, key: str) -> CircuitBreakerCtx:
+        return self.objs[key].ctx()

src/sentry/hybridcloud/apigateway_async/proxy.py

@@ -13,7 +13,7 @@
 import httpx
 from asgiref.sync import sync_to_async
 from django.conf import settings
-from django.http import HttpRequest, HttpResponse, StreamingHttpResponse
+from django.http import HttpRequest, HttpResponse, JsonResponse, StreamingHttpResponse
 from django.http.response import HttpResponseBase
 
 from sentry import options
@@ -34,9 +34,16 @@
 from sentry.utils import metrics
 from sentry.utils.http import BodyAsyncWrapper
 
+from .circuitbreaker import (
+    CircuitBreakerManager,
+    CircuitBreakerOverflow,
+    CircuitBreakerWindowOverflow,
+)
+
 logger = logging.getLogger(__name__)
 
 proxy_client = httpx.AsyncClient()
+circuitbreakers = CircuitBreakerManager()
 
 # Endpoints that handle uploaded files have higher timeouts configured
 # and we need to honor those timeouts when proxying.
@@ -141,8 +148,9 @@ async def proxy_cell_request(
     request: HttpRequest,
     cell: Cell,
     url_name: str,
-) -> StreamingHttpResponse:
+) -> HttpResponseBase:
     """Take a django request object and proxy it to a cell silo"""
+    metric_tags = {"region": cell.name, "url_name": url_name}
     target_url = urljoin(cell.address, request.path)
 
     content_encoding = request.headers.get("Content-Encoding")
@@ -154,37 +162,59 @@ async def proxy_cell_request(
     query_params = request.GET
 
     timeout = ENDPOINT_TIMEOUT_OVERRIDE.get(url_name, settings.GATEWAY_PROXY_TIMEOUT)
-    metric_tags = {"region": cell.name, "url_name": url_name}
 
     # XXX: See sentry.testutils.pytest.sentry for more information
     if settings.APIGATEWAY_PROXY_SKIP_RELAY and request.path.startswith("/api/0/relays/"):
         return StreamingHttpResponse(streaming_content="relay proxy skipped", status=404)
 
-    if url_name == "sentry-api-0-organization-objectstore":
-        if content_encoding:
-            header_dict["Content-Encoding"] = content_encoding
-        data = get_raw_body_async(request)
-    else:
-        data = BodyAsyncWrapper(request.body)
-        # With request streaming, and without `Content-Length` header,
-        # `httpx` will set chunked transfer encoding.
-        # Upstream doesn't necessarily support this,
-        # thus we re-add the header if it was present in the original request.
-        if content_length:
-            header_dict["Content-Length"] = content_length
-
     try:
-        with metrics.timer("apigateway.proxy_request.duration", tags=metric_tags):
-            req = proxy_client.build_request(
-                request.method,
-                target_url,
-                headers=header_dict,
-                params=dict(query_params) if query_params is not None else None,
-                content=_stream_request(data) if data else None,  # type: ignore[arg-type]
-                timeout=timeout,
-            )
-            resp = await proxy_client.send(req, stream=True, follow_redirects=False)
-            return _adapt_response(resp, target_url)
-    except (httpx.TimeoutException, asyncio.CancelledError):
-        # remote silo timeout. Use DRF timeout instead
-        raise RequestTimeout()
+        async with circuitbreakers.get(cell.name) as circuitbreaker:
+            if url_name == "sentry-api-0-organization-objectstore":
+                if content_encoding:
+                    header_dict["Content-Encoding"] = content_encoding
+                data = get_raw_body_async(request)
+            else:
+                data = BodyAsyncWrapper(request.body)
+                # With request streaming, and without `Content-Length` header,
+                # `httpx` will set chunked transfer encoding.
+                # Upstream doesn't necessarily support this,
+                # thus we re-add the header if it was present in the original request.
+                if content_length:
+                    header_dict["Content-Length"] = content_length
+
+            try:
+                with metrics.timer("apigateway.proxy_request.duration", tags=metric_tags):
+                    req = proxy_client.build_request(
+                        request.method,
+                        target_url,
+                        headers=header_dict,
+                        params=dict(query_params) if query_params is not None else None,
+                        content=_stream_request(data) if data else None,  # type: ignore[arg-type]
+                        timeout=timeout,
+                    )
+                    resp = await proxy_client.send(req, stream=True, follow_redirects=False)
+                    if resp.status_code >= 502:
+                        metrics.incr("apigateway.proxy.request_failed", tags=metric_tags)
+                        circuitbreaker.incr_failures()
+                    return _adapt_response(resp, target_url)
+            except (httpx.TimeoutException, asyncio.CancelledError):
+                metrics.incr("apigateway.proxy.request_timeout", tags=metric_tags)
+                circuitbreaker.incr_failures()
+                # remote silo timeout. Use DRF timeout instead
+                raise RequestTimeout()
+            except httpx.RequestError:
+                metrics.incr("apigateway.proxy.request_failed", tags=metric_tags)
+                circuitbreaker.incr_failures()
+                raise
+    except CircuitBreakerOverflow:
+        metrics.incr("apigateway.proxy.circuit_breaker.overflow", tags=metric_tags)
+        return JsonResponse(
+            {"error": "apigateway", "detail": "Too many requests"},
+            status=429,
+        )
+    except CircuitBreakerWindowOverflow:
+        metrics.incr("apigateway.proxy.circuit_breaker.rejected", tags=metric_tags)
+        return JsonResponse(
+            {"error": "apigateway", "detail": "Downstream service temporarily unavailable"},
+            status=503,
+        )