From 87e187bba578cad8eaf73ebc1204d9828e737ceb Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Wed, 9 Aug 2023 10:29:18 -0700 Subject: [PATCH 1/2] Fix async wrapper implementations --- newrelic/common/async_wrapper.py | 86 +++++++++++++++++--------------- 1 file changed, 47 insertions(+), 39 deletions(-) diff --git a/newrelic/common/async_wrapper.py b/newrelic/common/async_wrapper.py index 575b30e308..2d3db2b4be 100644 --- a/newrelic/common/async_wrapper.py +++ b/newrelic/common/async_wrapper.py @@ -20,6 +20,7 @@ is_generator_function, is_async_generator_function, ) +from newrelic.packages import six def evaluate_wrapper(wrapper_string, wrapped, trace): @@ -61,24 +62,41 @@ def wrapper(*args, **kwargs): return wrapped -def generator_wrapper(wrapped, trace): - @functools.wraps(wrapped) - def wrapper(*args, **kwargs): - g = wrapped(*args, **kwargs) - value = None - with trace: - while True: +if six.PY3: + def generator_wrapper(wrapped, trace): + WRAPPER = textwrap.dedent(""" + @functools.wraps(wrapped) + def wrapper(*args, **kwargs): + with trace: + result = yield from wrapped(*args, **kwargs) + return result + """) + + try: + return evaluate_wrapper(WRAPPER, wrapped, trace) + except: + return wrapped +else: + def generator_wrapper(wrapped, trace): + @functools.wraps(wrapped) + def wrapper(*args, **kwargs): + g = wrapped(*args, **kwargs) + with trace: try: - yielded = g.send(value) + yielded = g.send(None) + while True: + try: + sent = yield yielded + except GeneratorExit as e: + g.close() + raise + except BaseException as e: + yielded = g.throw(e) + else: + yielded = g.send(sent) except StopIteration: - break - - try: - value = yield yielded - except BaseException as e: - value = yield g.throw(type(e), e) - - return wrapper + return + return wrapper def async_generator_wrapper(wrapped, trace): @@ -86,31 +104,21 @@ def async_generator_wrapper(wrapped, trace): @functools.wraps(wrapped) async def wrapper(*args, **kwargs): g = wrapped(*args, **kwargs) - value = None with trace: - while True: - try: - yielded = await g.asend(value) - except StopAsyncIteration as e: - # The underlying async generator has finished, return propagates a new StopAsyncIteration - return - except StopIteration as e: - # The call to async_generator_asend.send() should raise a StopIteration containing the yielded value - yielded = e.value - - try: - value = yield yielded - except BaseException as e: - # An exception was thrown with .athrow(), propagate to the original async generator. - # Return value logic must be identical to .asend() + try: + yielded = await g.asend(None) + while True: try: - value = yield await g.athrow(type(e), e) - except StopAsyncIteration as e: - # The underlying async generator has finished, return propagates a new StopAsyncIteration - return - except StopIteration as e: - # The call to async_generator_athrow.send() should raise a StopIteration containing a yielded value - value = yield e.value + sent = yield yielded + except GeneratorExit as e: + await g.aclose() + raise + except BaseException as e: + yielded = await g.athrow(e) + else: + yielded = await g.asend(sent) + except StopAsyncIteration: + return """) try: From 917918309d9076847748c5e09149f18ac712351a Mon Sep 17 00:00:00 2001 From: Tim Pansino Date: Wed, 9 Aug 2023 10:40:55 -0700 Subject: [PATCH 2/2] Add regression testing --- .../_test_async_generator_trace.py | 35 +++++++++++++++++++ tests/agent_features/test_coroutine_trace.py | 31 ++++++++++++++++ 2 files changed, 66 insertions(+) diff --git a/tests/agent_features/_test_async_generator_trace.py b/tests/agent_features/_test_async_generator_trace.py index f6097bc434..30b970c372 100644 --- a/tests/agent_features/_test_async_generator_trace.py +++ b/tests/agent_features/_test_async_generator_trace.py @@ -323,6 +323,41 @@ async def _test(): event_loop.run_until_complete(_test()) +@validate_transaction_metrics( + "test_multiple_throws_yield_a_value", + background_task=True, + scoped_metrics=[("Function/agen", 1)], + rollup_metrics=[("Function/agen", 1)], +) +def test_multiple_throws_yield_a_value(event_loop): + @function_trace(name="agen") + async def agen(): + value = None + for _ in range(4): + try: + yield value + value = "bar" + except MyException: + value = "foo" + + + @background_task(name="test_multiple_throws_yield_a_value") + async def _test(): + gen = agen() + + # kickstart the coroutine + assert await gen.asend(None) is None + assert await gen.athrow(MyException) == "foo" + assert await gen.athrow(MyException) == "foo" + assert await gen.asend(None) == "bar" + + # finish consumption of the coroutine if necessary + async for _ in gen: + pass + + event_loop.run_until_complete(_test()) + + @validate_transaction_metrics( "test_athrow_does_not_yield_a_value", background_task=True, diff --git a/tests/agent_features/test_coroutine_trace.py b/tests/agent_features/test_coroutine_trace.py index a30f1e70a8..2043f13268 100644 --- a/tests/agent_features/test_coroutine_trace.py +++ b/tests/agent_features/test_coroutine_trace.py @@ -340,6 +340,37 @@ def coro(): pass +@validate_transaction_metrics( + "test_multiple_throws_yield_a_value", + background_task=True, + scoped_metrics=[("Function/coro", 1)], + rollup_metrics=[("Function/coro", 1)], +) +@background_task(name="test_multiple_throws_yield_a_value") +def test_multiple_throws_yield_a_value(): + @function_trace(name="coro") + def coro(): + value = None + for _ in range(4): + try: + yield value + value = "bar" + except MyException: + value = "foo" + + c = coro() + + # kickstart the coroutine + assert next(c) is None + assert c.throw(MyException) == "foo" + assert c.throw(MyException) == "foo" + assert next(c) == "bar" + + # finish consumption of the coroutine if necessary + for _ in c: + pass + + @pytest.mark.parametrize( "trace", [