Add support for instrumenting async functions via decorator.

This is only for python versions 3.6 and up.

Fixes #633

Author: tim-schilling

CHANGELOG.md

@@ -6,6 +6,10 @@
 
 - Tested with Django 3.2. Only define ``default_app_config`` when using
   a version of Django earlier than 3.2.
+- Support instrumentation and transaction decorators for asynchronous
+  functions via ``@instrument.async_``, ``@WebTransaction.async_`` and
+  ``@BackgroundTransaction.async_``.
+  ([PR #633](https://github.com/scoutapp/scout_apm_python/issues/633))
 
 ### Fixed
 

src/scout_apm/api/__init__.py

@@ -6,6 +6,15 @@
 from scout_apm.core.config import ScoutConfig
 from scout_apm.core.tracked_request import TrackedRequest
 
+# The async_ module can only be shipped on Python 3.6+
+try:
+    from scout_apm.async_.api import AsyncDecoratorMixin
+except ImportError:
+
+    class AsyncDecoratorMixin(object):
+        pass
+
+
 __all__ = [
     "BackgroundTransaction",
     "Config",
@@ -41,7 +50,7 @@ def ignore_transaction():
     TrackedRequest.instance().tag("ignore_transaction", True)
 
 
-class instrument(ContextDecorator):
+class instrument(AsyncDecoratorMixin, ContextDecorator):
     def __init__(self, operation, kind="Custom", tags=None):
         self.operation = text(kind) + "/" + text(operation)
         if tags is None:
@@ -66,7 +75,7 @@ def tag(self, key, value):
             self.span.tag(key, value)
 
 
-class Transaction(ContextDecorator):
+class Transaction(AsyncDecoratorMixin, ContextDecorator):
     """
     This Class is not meant to be used directly.
     Use one of the subclasses

src/scout_apm/async_/api.py

@@ -0,0 +1,42 @@
+# coding=utf-8
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+from functools import wraps
+
+
+class AsyncDecoratorMixin(object):
+    """Provide the ability to decorate both sync and async functions."""
+
+    is_async = False
+
+    @classmethod
+    def async_(cls, operation, tags=None, **kwargs):
+        """
+        Instrument an async function via a decorator.
+
+        This will return an awaitable which must be awaited.
+        Using this on a synchronous function will raise a
+        RuntimeError.
+
+        ``
+           @instrument.async_("Foo")
+           async def foo():
+              ...
+        ``
+        """
+        instance = cls(operation, tags=tags, **kwargs)
+        instance.is_async = True
+        return instance
+
+    def __call__(self, func):
+        if self.is_async:
+            # Until https://bugs.python.org/issue37398 has a resolution,
+            # manually wrap the async function
+            @wraps(func)
+            async def decorated(*args, **kwds):
+                with self._recreate_cm():
+                    return await func(*args, **kwds)
+
+            return decorated
+        else:
+            return super().__call__(func)

src/scout_apm/compat.py

@@ -34,13 +34,15 @@ def iteritems(dictionary):
 if sys.version_info >= (3, 2):
     from contextlib import ContextDecorator
 else:
-    import functools
 
     class ContextDecorator(object):
+        def _recreate_cm(self):
+            return self
+
         def __call__(self, f):
-            @functools.wraps(f)
+            @wraps(f)
             def decorated(*args, **kwds):
-                with self:
+                with self._recreate_cm():
                     return f(*args, **kwds)
 
             return decorated

tests/integration/instruments/test_asyncio_py36plus.py

@@ -130,6 +130,7 @@ async def test_task_awaited(tracked_request):
 
 @async_test
 async def test_nested_tasks(tracked_request):
+    @instrument.async_("orchestrator")
     async def orchestrator():
         await coro("1")
         await asyncio.gather(
@@ -143,11 +144,12 @@ async def orchestrator():
         await orchestrator()
 
     spans = tracked_request.complete_spans
-    assert len(spans) == 4
+    assert len(spans) == 5
     assert [span.operation for span in spans] == [
         "Custom/coro",
         "Custom/coro",
         "Custom/coro",
+        "Custom/orchestrator",
         "Job/test",
     ]
     # Verify the order of the coroutines
@@ -156,6 +158,7 @@ async def orchestrator():
         "2b",
         "2a",
         None,
+        None,
     ]
 
 

tests/integration/test_api_py36plus.py

@@ -0,0 +1,248 @@
+# coding=utf-8
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import pytest
+
+from scout_apm.api import BackgroundTransaction, WebTransaction, instrument
+from tests.tools import async_test
+
+
+@async_test
+async def test_instrument_decorator_async(tracked_request):
+    @instrument.async_("Foo")
+    async def foo():
+        pass
+
+    @instrument.async_("Bar")
+    async def example():
+        await foo()
+
+    await example()
+
+    assert len(tracked_request.active_spans) == 0
+    assert len(tracked_request.complete_spans) == 2
+    assert tracked_request.complete_spans[0].operation == "Custom/Foo"
+    assert tracked_request.complete_spans[1].operation == "Custom/Bar"
+
+
+def test_instrument_decorator_async_for_sync_function(tracked_request):
+    @instrument.async_("Bar")
+    def example():
+        pass
+
+    with pytest.warns(RuntimeWarning):
+        example()
+
+    assert len(tracked_request.active_spans) == 0
+    assert len(tracked_request.complete_spans) == 0
+
+
+@async_test
+async def test_instrument_decorator_async_misconfigured(tracked_request):
+    """Test case where .async_ isn't used from parent instrument"""
+
+    @instrument.async_("Foo")
+    async def foo():
+        pass
+
+    @instrument("Bar")
+    async def example():
+        await foo()
+
+    await example()
+
+    assert len(tracked_request.active_spans) == 0
+    assert len(tracked_request.complete_spans) == 1
+    assert tracked_request.complete_spans[0].operation == "Custom/Bar"
+
+
+@async_test
+async def test_instrument_decorator_async_classmethod(tracked_request):
+    class Example(object):
+        @classmethod
+        @instrument.async_("Test Decorator")
+        async def method(cls):
+            pass
+
+    await Example.method()
+
+    assert len(tracked_request.active_spans) == 0
+    assert len(tracked_request.complete_spans) == 1
+    assert tracked_request.complete_spans[0].operation == "Custom/Test Decorator"
+
+
+@async_test
+async def test_instrument_decorator_async_staticmethod(tracked_request):
+    class Example(object):
+        @staticmethod
+        @instrument.async_("Test Decorator")
+        async def method():
+            pass
+
+    await Example.method()
+
+    assert len(tracked_request.active_spans) == 0
+    assert len(tracked_request.complete_spans) == 1
+    assert tracked_request.complete_spans[0].operation == "Custom/Test Decorator"
+
+
+@async_test
+async def test_instrument_decorator_async_return_awaitable(tracked_request):
+    @instrument.async_("Foo")
+    async def foo():
+        pass
+
+    @instrument.async_("Bar")
+    def return_awaitable():
+        return foo()
+
+    await return_awaitable()
+
+    assert len(tracked_request.active_spans) == 0
+    assert len(tracked_request.complete_spans) == 2
+    assert tracked_request.complete_spans[0].operation == "Custom/Foo"
+    assert tracked_request.complete_spans[1].operation == "Custom/Bar"
+
+
+@async_test
+async def test_instrument_decorator_async_return_awaitable_misconfigured(
+    tracked_request,
+):
+    """Test case where .async_ isn't used from parent instrument"""
+
+    @instrument.async_("Foo")
+    async def foo():
+        pass
+
+    @instrument("Bar")
+    def return_awaitable():
+        return foo()
+
+    await return_awaitable()
+
+    assert len(tracked_request.active_spans) == 0
+    assert len(tracked_request.complete_spans) == 1
+    assert tracked_request.complete_spans[0].operation == "Custom/Bar"
+
+
+@async_test
+async def test_instrument_context_manager_async_await_later(tracked_request):
+    """
+    Test proving that if an awaitable goes unawaited in a context manager,
+    the spans are lost.
+    """
+
+    @instrument.async_("Outer")
+    async def foo():
+        with instrument("Inner"):
+            pass
+
+    async def example():
+        await foo()
+
+    with instrument("Test Decorator"):
+        awaitable = example()
+
+    await awaitable
+
+    assert len(tracked_request.active_spans) == 0
+    assert len(tracked_request.complete_spans) == 1
+    assert tracked_request.complete_spans[0].operation == "Custom/Test Decorator"
+
+
+@async_test
+async def test_web_transaction_decorator_async(tracked_request):
+    @instrument.async_("Foo")
+    async def foo():
+        pass
+
+    @WebTransaction.async_("Bar")
+    async def my_transaction():
+        await foo()
+
+    await my_transaction()
+
+    assert len(tracked_request.active_spans) == 0
+    assert len(tracked_request.complete_spans) == 2
+    assert tracked_request.complete_spans[0].operation == "Custom/Foo"
+    assert tracked_request.complete_spans[1].operation == "Controller/Bar"
+
+
+@async_test
+async def test_web_transaction_decorator_async_misconfigured(tracked_request):
+    """Test case where .async_ isn't used from WebTransaction"""
+
+    @instrument.async_("Foo")
+    async def foo():
+        pass
+
+    @WebTransaction("Bar")
+    async def my_transaction():
+        await foo()
+
+    await my_transaction()
+
+    assert len(tracked_request.active_spans) == 0
+    assert len(tracked_request.complete_spans) == 1
+    assert tracked_request.complete_spans[0].operation == "Controller/Bar"
+
+
+def test_web_transaction_decorator_async_for_sync_function(tracked_request):
+    @WebTransaction.async_("Bar")
+    def example():
+        pass
+
+    with pytest.warns(RuntimeWarning):
+        example()
+
+    assert len(tracked_request.active_spans) == 0
+    assert len(tracked_request.complete_spans) == 0
+
+
+@async_test
+async def test_background_transaction_decorator_async(tracked_request):
+    @instrument.async_("Foo")
+    async def foo():
+        pass
+
+    @BackgroundTransaction.async_("Bar")
+    async def my_transaction():
+        await foo()
+
+    await my_transaction()
+
+    assert len(tracked_request.active_spans) == 0
+    assert len(tracked_request.complete_spans) == 2
+    assert tracked_request.complete_spans[0].operation == "Custom/Foo"
+    assert tracked_request.complete_spans[1].operation == "Job/Bar"
+
+
+@async_test
+async def test_background_transaction_decorator_async_misconfigured(tracked_request):
+    """Test case where .async_ isn't used from BackgroundTransaction"""
+
+    @instrument.async_("Foo")
+    async def foo():
+        pass
+
+    @BackgroundTransaction("Bar")
+    async def my_transaction():
+        await foo()
+
+    await my_transaction()
+
+    assert len(tracked_request.active_spans) == 0
+    assert len(tracked_request.complete_spans) == 1
+    assert tracked_request.complete_spans[0].operation == "Job/Bar"
+
+
+def test_background_transaction_decorator_async_for_sync_function(tracked_request):
+    @BackgroundTransaction.async_("Bar")
+    def example():
+        pass
+
+    with pytest.warns(RuntimeWarning):
+        example()
+
+    assert len(tracked_request.active_spans) == 0
+    assert len(tracked_request.complete_spans) == 0