Merge pull request #1111 from newrelic/fix-asgiref-instrumentation

Fix Asgiref Instrumentation

Author: hmstepanek

newrelic/hooks/adapter_asgiref.py

@@ -14,16 +14,23 @@
 
 from newrelic.api.time_trace import current_trace
 from newrelic.common.object_wrapper import wrap_function_wrapper
+from newrelic.common.signature import bind_args
 from newrelic.core.context import ContextOf, context_wrapper_async
 
 
-def _bind_thread_handler(loop, source_task, *args, **kwargs):
-    return source_task
+async def wrap_SyncToAsync__call__(wrapped, instance, args, kwargs):
+    kwargs["_nr_current_trace"] = current_trace()
+    return await wrapped(*args, **kwargs)
 
 
 def thread_handler_wrapper(wrapped, instance, args, kwargs):
-    task = _bind_thread_handler(*args, **kwargs)
-    with ContextOf(trace_cache_id=id(task), strict=False):
+    try:
+        bound_args = bind_args(wrapped, args, kwargs)
+        trace = bound_args["args"][0].keywords.pop("_nr_current_trace", None)
+    except Exception:
+        trace = None
+
+    with ContextOf(trace=trace, strict=False):
         return wrapped(*args, **kwargs)
 
 
@@ -34,4 +41,5 @@ def main_wrap_wrapper(wrapped, instance, args, kwargs):
 
 def instrument_asgiref_sync(module):
     wrap_function_wrapper(module, "SyncToAsync.thread_handler", thread_handler_wrapper)
+    wrap_function_wrapper(module, "SyncToAsync.__call__", wrap_SyncToAsync__call__)
     wrap_function_wrapper(module, "AsyncToSync.main_wrap", main_wrap_wrapper)

tests/adapter_asgiref/conftest.py

@@ -0,0 +1,32 @@
+# Copyright 2010 New Relic, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from testing_support.fixture.event_loop import (  # noqa: F401; pylint: disable=W0611
+    event_loop as loop,
+)
+from testing_support.fixtures import (  # noqa: F401; pylint: disable=W0611
+    collector_agent_registration_fixture,
+    collector_available_fixture,
+)
+
+_default_settings = {
+    "transaction_tracer.explain_threshold": 0.0,
+    "transaction_tracer.transaction_threshold": 0.0,
+    "transaction_tracer.stack_trace_threshold": 0.0,
+    "debug.log_data_collector_payloads": True,
+    "debug.record_transaction_failure": True,
+}
+
+collector_agent_registration = collector_agent_registration_fixture(
+    app_name="Python Agent Test (adapter_asgiref)", default_settings=_default_settings
+)

tests/adapter_asgiref/test_context_propagation.py

@@ -0,0 +1,63 @@
+# Copyright 2010 New Relic, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from asgiref.sync import async_to_sync, sync_to_async
+from testing_support.validators.validate_transaction_metrics import (
+    validate_transaction_metrics,
+)
+
+from newrelic.api.background_task import background_task
+from newrelic.api.function_trace import function_trace
+from newrelic.api.transaction import current_transaction
+
+
+@sync_to_async
+@function_trace()
+def sync_func():
+    assert current_transaction()
+    return 1
+
+
+@validate_transaction_metrics(
+    "test_context_propagation:test_sync_to_async_context_propagation",
+    scoped_metrics=[("Function/test_context_propagation:sync_func", 1)],
+    rollup_metrics=[("Function/test_context_propagation:sync_func", 1)],
+    background_task=True,
+)
+@background_task()
+def test_sync_to_async_context_propagation(loop):
+    async def _test():
+        return_val = await sync_func()
+        assert return_val == 1, "Sync function failed to return"
+
+    loop.run_until_complete(_test())
+
+
+@async_to_sync
+@function_trace()
+async def async_func():
+    assert current_transaction()
+    return 1
+
+
+@validate_transaction_metrics(
+    "test_context_propagation:test_async_to_sync_context_propagation",
+    scoped_metrics=[("Function/test_context_propagation:async_func", 1)],
+    rollup_metrics=[("Function/test_context_propagation:async_func", 1)],
+    background_task=True,
+)
+@background_task()
+def test_async_to_sync_context_propagation():
+    return_val = async_func()
+    assert return_val == 1, "Async function failed to return"

tox.ini

@@ -71,6 +71,8 @@ envlist =
     postgres-datastore_psycopg2-{py27,py37,py38,py39,py310,py311,py312}-psycopg2latest,
     postgres-datastore_psycopg2cffi-{py27,pypy27,py37,py38,py39,py310,py311,py312}-psycopg2cffilatest,
     postgres-datastore_pyodbc-{py27,py37,py38,py39,py310,py311,py312}-pyodbclatest,
+    python-adapter_asgiref-{py37,py38,py39,py310,py311,py312,pypy310}-asgireflatest,
+    python-adapter_asgiref-py310-asgiref{0303,0304,0305,0306,0307},
     python-adapter_cheroot-{py27,py37,py38,py39,py310,py311,py312},
     python-adapter_daphne-{py37,py38,py39,py310,py311,py312}-daphnelatest,
     python-adapter_gevent-{py27,py37,py38,py310,py311,py312},
@@ -163,6 +165,12 @@ deps =
     WebTest==2.0.35
 
     # Test Suite Dependencies
+    adapter_asgiref-asgireflatest: asgiref
+    adapter_asgiref-asgiref0303: asgiref<3.4
+    adapter_asgiref-asgiref0304: asgiref<3.5
+    adapter_asgiref-asgiref0305: asgiref<3.6
+    adapter_asgiref-asgiref0306: asgiref<3.7
+    adapter_asgiref-asgiref0307: asgiref<3.8
     adapter_cheroot: cheroot
     adapter_daphne-daphnelatest: daphne
     adapter_gevent: WSGIProxy2
@@ -386,6 +394,7 @@ extras =
     agent_streaming: infinite-tracing
 
 changedir =
+    adapter_asgiref: tests/adapter_asgiref
     adapter_cheroot: tests/adapter_cheroot
     adapter_daphne: tests/adapter_daphne
     adapter_gevent: tests/adapter_gevent