fix(django) - Fix Django async views not behaving asyncronuously (#992)

ahmedetefy · sentry-bot · web-flow · commit 123f7af869a3 · 2021-02-03T16:16:43.000+01:00
* Refactored middlware span creation logic for middleware functions

* Added async instrumentation for django middlewares

* Added conditional that checks if async

* fix: Formatting

* Inherit from MiddlewareMixin for async behavior

* Refactored __call__ to be like __acall__ for better readability

* fix: Formatting

* Removed baseclass MiddlewareMixin for unecpected behavior

* fix: Formatting

* Added async_capable attribute to SentryWrappingMiddleware

* Added types to function signatures

* Refactored py3 logic to asgi module for py2 compat

* fix: Formatting

* Fixed function signature error

* fix: Formatting

* Refactored code to support both versions prior to Django 3.1 and after

* fix: Formatting

* Refactor middleware arg from asgi mixin factory

* fix: Formatting

* Added Types and documentation

* fix: Formatting

* Fixed py2 asgi mixin signature

* Added my_async_viewto myapp.views

* Added test to ensure concurrent behaviour in both ASGI and Django Channels

* Added urlpattern for my_async_view

* fix: Formatting

* Added test that ensures Performance timing spans are done correctly for async views

* Removed print statement

* Modified async_route_check function

* Added check for forwarding the async calls

* fix: Formatting

* Fixed django compat asgi_application import issue

* Fixed type import issues

* Linting changes

* fix: Formatting

* Fixed failing test by adding safeguard for middleware invocation for older django versions

* Removed unused import

* Removed redundant ASGI_APP global variable

* Added better documentation and modified method name for asgi middleware mixin factory

* Removed concurrency test for channels

* fix: Formatting

* Fixed typing and lint issues

Co-authored-by: sentry-bot &lt;markus+ghbot@sentry.io&gt;
diff --git a/sentry_sdk/integrations/django/asgi.py b/sentry_sdk/integrations/django/asgi.py
@@ -6,6 +6,8 @@
 `django.core.handlers.asgi`.
 """
 
+import asyncio
+
 from sentry_sdk import Hub, _functools
 from sentry_sdk._types import MYPY
 
@@ -14,6 +16,7 @@
 if MYPY:
     from typing import Any
     from typing import Union
+    from typing import Callable
 
     from django.http.response import HttpResponse
 
@@ -91,3 +94,52 @@ async def sentry_wrapped_callback(request, *args, **kwargs):
             return await callback(request, *args, **kwargs)
 
     return sentry_wrapped_callback
+
+
+def _asgi_middleware_mixin_factory(_check_middleware_span):
+    # type: (Callable[..., Any]) -> Any
+    """
+    Mixin class factory that generates a middleware mixin for handling requests
+    in async mode.
+    """
+
+    class SentryASGIMixin:
+        def __init__(self, get_response):
+            # type: (Callable[..., Any]) -> None
+            self.get_response = get_response
+            self._acall_method = None
+            self._async_check()
+
+        def _async_check(self):
+            # type: () -> None
+            """
+            If get_response is a coroutine function, turns us into async mode so
+            a thread is not consumed during a whole request.
+            Taken from django.utils.deprecation::MiddlewareMixin._async_check
+            """
+            if asyncio.iscoroutinefunction(self.get_response):
+                self._is_coroutine = asyncio.coroutines._is_coroutine  # type: ignore
+
+        def async_route_check(self):
+            # type: () -> bool
+            """
+            Function that checks if we are in async mode,
+            and if we are forwards the handling of requests to __acall__
+            """
+            return asyncio.iscoroutinefunction(self.get_response)
+
+        async def __acall__(self, *args, **kwargs):
+            # type: (*Any, **Any) -> Any
+            f = self._acall_method
+            if f is None:
+                self._acall_method = f = self._inner.__acall__  # type: ignore
+
+            middleware_span = _check_middleware_span(old_method=f)
+
+            if middleware_span is None:
+                return await f(*args, **kwargs)
+
+            with middleware_span:
+                return await f(*args, **kwargs)
+
+    return SentryASGIMixin
diff --git a/sentry_sdk/integrations/django/middleware.py b/sentry_sdk/integrations/django/middleware.py
@@ -16,8 +16,11 @@
 if MYPY:
     from typing import Any
     from typing import Callable
+    from typing import Optional
     from typing import TypeVar
 
+    from sentry_sdk.tracing import Span
+
     F = TypeVar("F", bound=Callable[..., Any])
 
 _import_string_should_wrap_middleware = ContextVar(
@@ -30,6 +33,12 @@
     import_string_name = "import_string"
 
 
+if DJANGO_VERSION < (3, 1):
+    _asgi_middleware_mixin_factory = lambda _: object
+else:
+    from .asgi import _asgi_middleware_mixin_factory
+
+
 def patch_django_middlewares():
     # type: () -> None
     from django.core.handlers import base
@@ -64,29 +73,40 @@ def _wrap_middleware(middleware, middleware_name):
     # type: (Any, str) -> Any
     from sentry_sdk.integrations.django import DjangoIntegration
 
+    def _check_middleware_span(old_method):
+        # type: (Callable[..., Any]) -> Optional[Span]
+        hub = Hub.current
+        integration = hub.get_integration(DjangoIntegration)
+        if integration is None or not integration.middleware_spans:
+            return None
+
+        function_name = transaction_from_function(old_method)
+
+        description = middleware_name
+        function_basename = getattr(old_method, "__name__", None)
+        if function_basename:
+            description = "{}.{}".format(description, function_basename)
+
+        middleware_span = hub.start_span(
+            op="django.middleware", description=description
+        )
+        middleware_span.set_tag("django.function_name", function_name)
+        middleware_span.set_tag("django.middleware_name", middleware_name)
+
+        return middleware_span
+
     def _get_wrapped_method(old_method):
         # type: (F) -> F
         with capture_internal_exceptions():
 
             def sentry_wrapped_method(*args, **kwargs):
                 # type: (*Any, **Any) -> Any
-                hub = Hub.current
-                integration = hub.get_integration(DjangoIntegration)
-                if integration is None or not integration.middleware_spans:
-                    return old_method(*args, **kwargs)
-
-                function_name = transaction_from_function(old_method)
+                middleware_span = _check_middleware_span(old_method)
 
-                description = middleware_name
-                function_basename = getattr(old_method, "__name__", None)
-                if function_basename:
-                    description = "{}.{}".format(description, function_basename)
+                if middleware_span is None:
+                    return old_method(*args, **kwargs)
 
-                with hub.start_span(
-                    op="django.middleware", description=description
-                ) as span:
-                    span.set_tag("django.function_name", function_name)
-                    span.set_tag("django.middleware_name", middleware_name)
+                with middleware_span:
                     return old_method(*args, **kwargs)
 
             try:
@@ -102,11 +122,22 @@ def sentry_wrapped_method(*args, **kwargs):
 
         return old_method
 
-    class SentryWrappingMiddleware(object):
-        def __init__(self, *args, **kwargs):
-            # type: (*Any, **Any) -> None
-            self._inner = middleware(*args, **kwargs)
+    class SentryWrappingMiddleware(
+        _asgi_middleware_mixin_factory(_check_middleware_span)  # type: ignore
+    ):
+
+        async_capable = getattr(middleware, "async_capable", False)
+
+        def __init__(self, get_response=None, *args, **kwargs):
+            # type: (Optional[Callable[..., Any]], *Any, **Any) -> None
+            if get_response:
+                self._inner = middleware(get_response, *args, **kwargs)
+            else:
+                self._inner = middleware(*args, **kwargs)
+            self.get_response = get_response
             self._call_method = None
+            if self.async_capable:
+                super(SentryWrappingMiddleware, self).__init__(get_response)
 
         # We need correct behavior for `hasattr()`, which we can only determine
         # when we have an instance of the middleware we're wrapping.
@@ -128,10 +159,20 @@ def __getattr__(self, method_name):
 
         def __call__(self, *args, **kwargs):
             # type: (*Any, **Any) -> Any
+            if hasattr(self, "async_route_check") and self.async_route_check():
+                return self.__acall__(*args, **kwargs)
+
             f = self._call_method
             if f is None:
-                self._call_method = f = _get_wrapped_method(self._inner.__call__)
-            return f(*args, **kwargs)
+                self._call_method = f = self._inner.__call__
+
+            middleware_span = _check_middleware_span(old_method=f)
+
+            if middleware_span is None:
+                return f(*args, **kwargs)
+
+            with middleware_span:
+                return f(*args, **kwargs)
 
     if hasattr(middleware, "__name__"):
         SentryWrappingMiddleware.__name__ = middleware.__name__
diff --git a/tests/integrations/django/asgi/test_asgi.py b/tests/integrations/django/asgi/test_asgi.py
@@ -68,3 +68,80 @@ async def test_async_views(sentry_init, capture_events, application):
         "query_string": None,
         "url": "/async_message",
     }
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    django.VERSION < (3, 1), reason="async views have been introduced in Django 3.1"
+)
+async def test_async_views_concurrent_execution(sentry_init, capture_events, settings):
+    import asyncio
+    import time
+
+    settings.MIDDLEWARE = []
+    asgi_application.load_middleware(is_async=True)
+
+    sentry_init(integrations=[DjangoIntegration()], send_default_pii=True)
+
+    comm = HttpCommunicator(asgi_application, "GET", "/my_async_view")
+    comm2 = HttpCommunicator(asgi_application, "GET", "/my_async_view")
+
+    loop = asyncio.get_event_loop()
+
+    start = time.time()
+
+    r1 = loop.create_task(comm.get_response(timeout=5))
+    r2 = loop.create_task(comm2.get_response(timeout=5))
+
+    (resp1, resp2), _ = await asyncio.wait({r1, r2})
+
+    end = time.time()
+
+    assert resp1.result()["status"] == 200
+    assert resp2.result()["status"] == 200
+
+    assert end - start < 1.5
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    django.VERSION < (3, 1), reason="async views have been introduced in Django 3.1"
+)
+async def test_async_middleware_spans(
+    sentry_init, render_span_tree, capture_events, settings
+):
+    settings.MIDDLEWARE = [
+        "django.contrib.sessions.middleware.SessionMiddleware",
+        "django.contrib.auth.middleware.AuthenticationMiddleware",
+        "django.middleware.csrf.CsrfViewMiddleware",
+        "tests.integrations.django.myapp.settings.TestMiddleware",
+    ]
+    asgi_application.load_middleware(is_async=True)
+
+    sentry_init(
+        integrations=[DjangoIntegration(middleware_spans=True)],
+        traces_sample_rate=1.0,
+        _experiments={"record_sql_params": True},
+    )
+
+    events = capture_events()
+
+    comm = HttpCommunicator(asgi_application, "GET", "/async_message")
+    response = await comm.get_response()
+    assert response["status"] == 200
+
+    await comm.wait()
+
+    message, transaction = events
+
+    assert (
+        render_span_tree(transaction)
+        == """\
+- op="http.server": description=null
+  - op="django.middleware": description="django.contrib.sessions.middleware.SessionMiddleware.__acall__"
+    - op="django.middleware": description="django.contrib.auth.middleware.AuthenticationMiddleware.__acall__"
+      - op="django.middleware": description="django.middleware.csrf.CsrfViewMiddleware.__acall__"
+        - op="django.middleware": description="tests.integrations.django.myapp.settings.TestMiddleware.__acall__"
+          - op="django.middleware": description="django.middleware.csrf.CsrfViewMiddleware.process_view"
+          - op="django.view": description="async_message\""""
+    )
diff --git a/tests/integrations/django/myapp/urls.py b/tests/integrations/django/myapp/urls.py
@@ -63,6 +63,9 @@ def path(path, *args, **kwargs):
 if views.async_message is not None:
     urlpatterns.append(path("async_message", views.async_message, name="async_message"))
 
+if views.my_async_view is not None:
+    urlpatterns.append(path("my_async_view", views.my_async_view, name="my_async_view"))
+
 # rest framework
 try:
     urlpatterns.append(
diff --git a/tests/integrations/django/myapp/views.py b/tests/integrations/django/myapp/views.py
@@ -141,5 +141,13 @@ def csrf_hello_not_exempt(*args, **kwargs):
     sentry_sdk.capture_message("hi")
     return HttpResponse("ok")"""
     )
+
+    exec(
+        """async def my_async_view(request):
+    import asyncio
+    await asyncio.sleep(1)
+    return HttpResponse('Hello World')"""
+    )
 else:
     async_message = None
+    my_async_view = None
