wip

Author: gi0baro

src/sentry/hybridcloud/apigateway/apigateway.py

@@ -31,7 +31,9 @@ def _get_view_silo_mode(view_func: Callable[..., HttpResponseBase]) -> frozenset
 
 
 async def proxy_request_if_needed(
-    request: Request, view_func: Callable[..., HttpResponseBase], view_kwargs: dict[str, Any]
+    request: Request,
+    view_func: Callable[..., HttpResponseBase],
+    view_kwargs: dict[str, Any],
 ) -> HttpResponseBase | None:
     """
     Main execution flow for the API Gateway.
@@ -73,7 +75,7 @@ async def proxy_request_if_needed(
                 "kind": "error-embed",
             },
         )
-        return proxy_error_embed_request(request, dsn, url_name)
+        return await proxy_error_embed_request(request, dsn, url_name)
 
     if (
         request.resolver_match

src/sentry/hybridcloud/apigateway/middleware.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
+import asyncio
 from collections.abc import Callable
 from typing import Any
 
-# from asgiref.sync import iscoroutinefunction, markcoroutinefunction
+from asgiref.sync import async_to_sync, iscoroutinefunction, markcoroutinefunction
 from django.http.response import HttpResponseBase
 from rest_framework.request import Request
 
@@ -14,17 +15,57 @@ class ApiGatewayMiddleware:
     """Proxy requests intended for remote silos"""
 
     async_capable = True
-    sync_capable = False
+    sync_capable = True
 
     def __init__(self, get_response: Callable[[Request], HttpResponseBase]):
         self.get_response = get_response
-        # if iscoroutinefunction(self.get_response):
-        #    markcoroutinefunction(self)
+        if iscoroutinefunction(self.get_response):
+            markcoroutinefunction(self)
 
     def __call__(self, request: Request) -> HttpResponseBase:
+        if iscoroutinefunction(self):
+            return self.__acall__(request)
         return self.get_response(request)
 
-    async def process_view(
+    async def __acall__(self, request: Request) -> HttpResponseBase:
+        return await self.get_response(request)
+
+    def process_view(
+        self,
+        request: Request,
+        view_func: Callable[..., HttpResponseBase],
+        view_args: tuple[str],
+        view_kwargs: dict[str, Any],
+    ) -> HttpResponseBase | None:
+        return self._process_view_match(request, view_func, view_args, view_kwargs)
+
+    def _process_view_match(
+        self,
+        request: Request,
+        view_func: Callable[..., HttpResponseBase],
+        view_args: tuple[str],
+        view_kwargs: dict[str, Any],
+    ):
+        #: we check if we're in an async or sync runtime once, then
+        #  overwrite the method with the actual impl.
+        try:
+            asyncio.get_running_loop()
+            method = self._process_view_inner
+        except RuntimeError:
+            method = self._process_view_sync
+        setattr(self, "_process_view_match", method)
+        return method(request, view_func, view_args, view_kwargs)
+
+    def _process_view_sync(
+        self,
+        request: Request,
+        view_func: Callable[..., HttpResponseBase],
+        view_args: tuple[str],
+        view_kwargs: dict[str, Any],
+    ) -> HttpResponseBase | None:
+        return async_to_sync(self._process_view_inner)(request, view_func, view_args, view_kwargs)
+
+    async def _process_view_inner(
         self,
         request: Request,
         view_func: Callable[..., HttpResponseBase],
@@ -34,5 +75,4 @@ async def process_view(
         proxy_response = await proxy_request_if_needed(request, view_func, view_kwargs)
         if proxy_response is not None:
             return proxy_response
-        else:
-            return None
+        return None

src/sentry/hybridcloud/apigateway/proxy.py

@@ -6,22 +6,19 @@
 
 import asyncio
 import logging
-from collections.abc import Generator
+from collections.abc import AsyncGenerator
 from urllib.parse import urljoin, urlparse
 from wsgiref.util import is_hop_by_hop
 
+import httpx
+from asgiref.sync import sync_to_async
 from django.conf import settings
 from django.http import HttpRequest, HttpResponse, StreamingHttpResponse
 from django.http.response import HttpResponseBase
-from httpx import AsyncClient as HttpClient
-from httpx import TimeoutException
-from requests import Response as ExternalResponse
 
-# from requests import request as external_request
-# from requests.exceptions import Timeout
 from sentry import options
 from sentry.api.exceptions import RequestTimeout
-from sentry.objectstore.endpoints.organization import ChunkedEncodingDecoder, get_raw_body
+from sentry.objectstore.endpoints.organization import get_raw_body_async
 from sentry.silo.util import (
     PROXY_APIGATEWAY_HEADER,
     PROXY_DIRECT_LOCATION_HEADER,
@@ -35,11 +32,11 @@
     get_cell_for_organization,
 )
 from sentry.utils import metrics
-from sentry.utils.http import BodyWithLength
+from sentry.utils.http import BodyAsyncWrapper
 
 logger = logging.getLogger(__name__)
 
-proxy_client = HttpClient()
+proxy_client = httpx.AsyncClient()
 
 # Endpoints that handle uploaded files have higher timeouts configured
 # and we need to honor those timeouts when proxying.
@@ -59,21 +56,27 @@
 PROXY_CHUNK_SIZE = 512 * 1024
 
 
-def _parse_response(response: ExternalResponse, remote_url: str) -> StreamingHttpResponse:
-    """
-    Convert the Responses class from requests into the drf Response
-    """
+async def _stream_response_and_close(response: httpx.Response) -> AsyncGenerator[bytes]:
+    """Yield chunks from an httpx response and close the connection when done."""
+    try:
+        async for chunk in response.aiter_bytes(PROXY_CHUNK_SIZE):
+            yield chunk
+    finally:
+        await response.aclose()
+
+
+def _adapt_response(response: httpx.Response, remote_url: str) -> StreamingHttpResponse:
+    """Convert an httpx Response into a Django response."""
 
-    # def stream_response() -> Generator[bytes]:
-    #    yield from response.iter_content(PROXY_CHUNK_SIZE)
     new_headers = clean_outbound_headers(response.headers)
+    content_type = new_headers.pop("Content-Type", None)
 
     streamed_response = StreamingHttpResponse(
-        streaming_content=response.aiter_bytes(),
+        streaming_content=_stream_response_and_close(response),
         status=response.status_code,
-        content_type=new_headers.pop("Content-Type", None),
+        content_type=content_type,
     )
-    # Add Headers to response
+
     for header, value in new_headers.items():
         if not is_hop_by_hop(header):
             streamed_response[header] = value
@@ -82,13 +85,20 @@ def _parse_response(response: ExternalResponse, remote_url: str) -> StreamingHtt
     return streamed_response
 
 
+async def _stream_request(body) -> AsyncGenerator[bytes]:
+    async for chunk in body:
+        yield chunk
+
+
 async def proxy_request(
-    request: HttpRequest, org_id_or_slug: str, url_name: str
+    request: HttpRequest,
+    org_id_or_slug: str,
+    url_name: str,
 ) -> HttpResponseBase:
     """Take a django request object and proxy it to a remote location given an org_id_or_slug"""
 
     try:
-        cell = get_cell_for_organization(org_id_or_slug)
+        cell = await sync_to_async(get_cell_for_organization)(org_id_or_slug)
     except CellResolutionError as e:
         logger.info("region_resolution_error", extra={"org_slug": org_id_or_slug, "error": str(e)})
         return HttpResponse(status=404)
@@ -128,7 +138,9 @@ async def proxy_error_embed_request(
 
 
 async def proxy_cell_request(
-    request: HttpRequest, cell: Cell, url_name: str
+    request: HttpRequest,
+    cell: Cell,
+    url_name: str,
 ) -> StreamingHttpResponse:
     """Take a django request object and proxy it to a cell silo"""
     target_url = urljoin(cell.address, request.path)
@@ -137,7 +149,6 @@ async def proxy_cell_request(
     header_dict = clean_proxy_headers(request.headers)
     header_dict[PROXY_APIGATEWAY_HEADER] = "true"
 
-    # TODO: use requests session for connection pooling capabilities
     assert request.method is not None
     query_params = request.GET
 
@@ -148,31 +159,26 @@ async def proxy_cell_request(
     if settings.APIGATEWAY_PROXY_SKIP_RELAY and request.path.startswith("/api/0/relays/"):
         return StreamingHttpResponse(streaming_content="relay proxy skipped", status=404)
 
-    data: bytes | Generator[bytes] | ChunkedEncodingDecoder | BodyWithLength | None = None
+    data: AsyncGenerator[bytes] | None = None
     if url_name == "sentry-api-0-organization-objectstore":
         if content_encoding:
             header_dict["Content-Encoding"] = content_encoding
-        data = get_raw_body(request)
+        data = get_raw_body_async(request)
     else:
-        data = BodyWithLength(request)
+        data = BodyAsyncWrapper(request.body)
 
     try:
         with metrics.timer("apigateway.proxy_request.duration", tags=metric_tags):
-            async with proxy_client.stream(
+            req = proxy_client.build_request(
                 request.method,
                 target_url,
                 headers=header_dict,
                 params=dict(query_params) if query_params is not None else None,
-                data=data,
+                content=_stream_request(data),
                 timeout=timeout,
-                follow_redirects=False,
-            ) as resp:
-                return _parse_response(resp, target_url)
-    except (TimeoutException, asyncio.CancelledError):
+            )
+            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()
-
-    # new_headers = clean_outbound_headers(resp.headers)
-    # resp.headers.clear()
-    # resp.headers.update(new_headers)
-    # return _parse_response(resp, target_url)

src/sentry/objectstore/endpoints/organization.py

@@ -1,17 +1,18 @@
 from __future__ import annotations
 
-from collections.abc import Callable, Generator
+from collections.abc import AsyncGenerator, Callable, Generator
 from typing import Any
 from urllib.parse import urlparse
 from wsgiref.util import is_hop_by_hop
 
 import requests
+from asgiref.sync import sync_to_async
 from django.http import HttpRequest, StreamingHttpResponse
 from requests import Response as ExternalResponse
 from rest_framework.request import Request
 from rest_framework.response import Response
 
-from sentry.utils.http import BodyWithLength
+from sentry.utils.http import BodyAsyncWrapper, BodyWithLength, BodyWithLengthAiter
 
 # TODO(granian): Remove this and related code paths when we fully switch from uwsgi to granian
 uwsgi: Any = None
@@ -150,6 +151,32 @@ def stream_generator():
     return BodyWithLength(request)
 
 
+def get_raw_body_async(
+    request: HttpRequest,
+) -> AsyncGenerator[bytes] | ChunkedEncodingAsyncDecoder | BodyWithLengthAiter | None:
+    if request.body:
+        return BodyAsyncWrapper(request.body)
+
+    wsgi_input = request.META.get("wsgi.input")
+    if "granian" in request.META.get("SERVER_SOFTWARE", "").lower():
+        return BodyAsyncWrapper(wsgi_input)
+
+    # wsgiref will raise an exception and hang when attempting to read wsgi.input while there's no body.
+    # For now, support bodies only on PUT and POST requests when not using Granian.
+    if request.method not in ("PUT", "POST"):
+        return None
+
+    # wsgiref (dev/test server)
+    if (
+        hasattr(wsgi_input, "_read")
+        and request.headers.get("Transfer-Encoding", "").lower() == "chunked"
+    ):
+        return ChunkedEncodingAsyncDecoder(wsgi_input._read)  # type: ignore[union-attr]
+
+    # wsgiref and the request has been already proxied through control silo
+    return BodyWithLength(request).__aiter__()
+
+
 def get_target_url(path: str) -> str:
     base = options.get("objectstore.config")["base_url"].rstrip("/")
     # `path` should be a relative path, only grab that part
@@ -248,3 +275,10 @@ def read(self, size: int = -1) -> bytes:
                         raise ValueError("Malformed chunk encoded stream")
 
         return b"".join(buffer)
+
+
+class ChunkedEncodingAsyncDecoder(ChunkedEncodingDecoder):
+    async def __anext__(self) -> bytes:
+        if self._done:
+            raise StopIteration
+        return await sync_to_async(self.read)()