Support Async Generators in Transaction Decorators (#1543)

* Add support for AsyncGenerators in transaction wrappers

* Add tests for async_generator_wrapper

* Add tests for async_generator_proxy

* Cleanup test names

* Add large explanatory comments

* Silence pytest warning for TestException containing Test

* Correct async generator aclose behavior and tests

* Correct tests for pypy garbage collection

---------

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: TimPansino

newrelic/common/async_proxy.py

@@ -15,7 +15,12 @@
 import logging
 import time
 
-from newrelic.common.coroutine import is_asyncio_coroutine, is_coroutine_callable, is_generator_function
+from newrelic.common.coroutine import (
+    is_async_generator_function,
+    is_asyncio_coroutine,
+    is_coroutine_callable,
+    is_generator_function,
+)
 from newrelic.common.object_wrapper import ObjectProxy
 from newrelic.core.trace_cache import trace_cache
 
@@ -29,6 +34,11 @@ def __init__(self, transaction_init):
         self.enter_time = None
         self.transaction = None
         self.transaction_init = transaction_init
+        # This flag is used to tell if the top level coroutine is an async generator.
+        # Depending on the conditions, the behavior when encountering
+        # StopIteration vs StopAsyncIteration is different. See the comments in __exit__
+        # for details on the expected behavior.
+        self.is_async_generator = False
 
     def pre_close(self):
         # If close is called prior to the start of the coroutine do not create
@@ -90,11 +100,63 @@ def __exit__(self, exc, value, tb):
             except:
                 CancelledError = GeneratorExit
 
-        # case: coroutine completed or cancelled
-        if exc is StopIteration or exc is GeneratorExit or exc is CancelledError:
+        # There are 3 separate cases we need to consider for when attempting to exit this context manager.
+        #
+        # When running as a generator, the execution path is simple. When StopIteration or GeneratorExit
+        # is raised, the generator has completed normally so we complete the transaction. Other exceptions
+        # are treated as errors and the transaction is closed with the exception details.
+        #
+        # When running as a coroutine, the execution path is only slightly changed. Calling __await__ on
+        # the coroutine will return a generator which is driven to completion. Since CoroutineProxy wraps
+        # this generator in a GeneratorProxy, the execution path ends up as the same as a generator, with
+        # the additional consideration of asyncio.CancelledError which is raised if the coroutine is cancelled.
+        # This is treated the same as GeneratorExit, completing the transaction without error.
+        #
+        # When running as an async generator, the execution path is considerably more complex. All of the
+        # interfaces for the async generator (__anext__, asend, athrow, aclose) are coroutines which themselves
+        # require context management. These coroutines are awaited, which again creates a generator that must also
+        # be tracked. This means that the TransactionContext will be entered and exited multiple times with
+        # StopIteration being thrown in an underlying generator for each item that's yielded from the async generator.
+        # We therefore need to avoid completing the transaction when StopIteration is raised, and only complete
+        # the transaction when StopAsyncIteration is raised which indicates the async generator has completed normally.
+        # GeneratorExit and asyncio.CancelledError are treated the same as the other cases, completing the transaction
+        # without error.
+        #
+        # The is_async_generator flag is used to tell if the top level coroutine is an async generator, which will
+        # then change the behavior with respect to StopIteration. As this TransactionContext object is shared
+        # between all the wrappers for async generators, the coroutines they create, and the generators created
+        # by awaiting those coroutines, we have to change the behavior of all of these types of wrappers when
+        # run under an async generator to avoid completing the transaction early.
+
+        if exc is StopAsyncIteration:
+            # If an async generator completes normally, complete the transaction without error.
+            # This condition is also manually run when an async generator is closed via AsyncGeneratorProxy.aclose().
+            if self.is_async_generator:
+                self.transaction.__exit__(None, None, None)
+            # If a non-async generator reaches this, complete the transaction and report as an error.
+            else:
+                self.transaction.__exit__(exc, value, tb)
+
+        elif exc is StopIteration:
+            # If a non-async generator completes normally, complete the transaction without error.
+            if not self.is_async_generator:
+                self.transaction.__exit__(None, None, None)
+
+            # If an async generator reaches this, don't complete the transaction as this is
+            # caused by yielding an item from the generator. This is due to completing the
+            # underlying coroutine which is a generator internally and will raise StopIteration.
+            # We need to wait until the async generator itself completes normally with
+            # a final StopAsyncIteration.
+
+            # Note: This block is equivalent to "else: pass".
+            # If this code block is unnested in the future, that should be made explicit.
+
+        # If coroutine was cancelled, either by asyncio.CancelledError, .close(),
+        # complete the transaction without error.
+        elif exc in (GeneratorExit, CancelledError):
             self.transaction.__exit__(None, None, None)
 
-        # case: coroutine completed because of error
+        # Unexpected exception, complete the transaction and report as an error.
         elif exc:
             self.transaction.__exit__(exc, value, tb)
 
@@ -151,9 +213,44 @@ def __await__(self):
         return GeneratorProxy(self.__wrapped__, self._nr_context)
 
 
+class AsyncGeneratorProxy(ObjectProxy):
+    def __init__(self, wrapped, context):
+        super().__init__(wrapped)
+        self._nr_context = context
+        # Set this flag to indicate that the top level coroutine is an async generator,
+        # which will change the behavior of __exit__ in TransactionContext for all the
+        # lower level wrappers.
+        self._nr_context.is_async_generator = True
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        return await self.asend(None)
+
+    async def asend(self, value):
+        return await CoroutineProxy(self.__wrapped__.asend(value), self._nr_context)
+
+    async def athrow(self, *args, **kwargs):
+        return await CoroutineProxy(self.__wrapped__.athrow(*args, **kwargs), self._nr_context)
+
+    async def aclose(self):
+        try:
+            return await CoroutineProxy(self.__wrapped__.aclose(), self._nr_context)
+        finally:
+            # There is nothing further down that can tell correctly that the async generator
+            # is being closed, so we call __exit__ with StopAsyncIteration manually to ensure
+            # the transaction is correctly completed without error. Otherwise this would rely on
+            # the garbage collector deleting the async generator to complete the transaction by
+            # calling Transaction.__del__, which will cause a hanging transaction and race conditions.
+            self._nr_context.__exit__(StopAsyncIteration, None, None)
+
+
 def async_proxy(wrapped):
     if is_coroutine_callable(wrapped):
         return CoroutineProxy
+    elif is_async_generator_function(wrapped):
+        return AsyncGeneratorProxy
     elif is_generator_function(wrapped):
         if is_asyncio_coroutine(wrapped):
             return AwaitableGeneratorProxy

tests/agent_features/test_async_generator_trace.py renamed to tests/agent_features/test_async_generator_trace_wrapper.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 import functools
+import gc
+import sys
 import time
 
 import pytest
@@ -84,7 +86,7 @@ class MyException(Exception):
     scoped_metrics=[("Function/agen", 1)],
     rollup_metrics=[("Function/agen", 1)],
 )
-@validate_transaction_errors(errors=["test_async_generator_trace:MyException"])
+@validate_transaction_errors(errors=["test_async_generator_trace_wrapper:MyException"])
 def test_async_generator_error(event_loop):
     @function_trace(name="agen")
     async def agen():
@@ -182,30 +184,36 @@ async def _test():
 
 
 @validate_transaction_metrics(
-    "test_async_generator_close_ends_trace",
+    "test_async_generator_aclose_ends_trace",
     background_task=True,
     scoped_metrics=[("Function/agen", 1)],
     rollup_metrics=[("Function/agen", 1)],
 )
-def test_async_generator_close_ends_trace(event_loop):
+def test_async_generator_aclose_ends_trace(event_loop):
     @function_trace(name="agen")
     async def agen():
         yield
 
-    @background_task(name="test_async_generator_close_ends_trace")
-    async def _test():
-        gen = agen()
+    # Save a reference to the generator and run the validations before that
+    # is garbage collected to avoid this test becoming a duplicate
+    # of the test "test_incomplete_async_generator"
+    gen = agen()
 
-        # kickstart the coroutine
-        await gen.asend(None)
+    @background_task(name="test_async_generator_aclose_ends_trace")
+    def _test_async_generator_aclose_ends_trace():
+        async def _test():
+            # kickstart the coroutine
+            await gen.asend(None)
 
-        # trace should be ended/recorded by close
-        await gen.aclose()
+            # trace should be ended/recorded by close
+            await gen.aclose()
 
-        # We may call gen.close as many times as we want
-        await gen.aclose()
+            # We may call gen.close as many times as we want
+            await gen.aclose()
 
-    event_loop.run_until_complete(_test())
+        event_loop.run_until_complete(_test())
+
+    _test_async_generator_aclose_ends_trace()
 
 
 @validate_tt_parenting(("TransactionNode", [("FunctionNode", [("FunctionNode", [])])]))
@@ -493,6 +501,10 @@ async def _test():
             async for _ in c:
                 break
 
+            # This test differs from the test for async_proxy in that the generator
+            # going out of scope does not immediately close the trace. Instead, it's
+            # the transaction ending that closes the trace. No need to call gc.collect().
+
         if nr_transaction:
             _test = background_task(name="test_incomplete_async_generator")(_test)