Distributed Tracing tests

Author: lrafeei

newrelic/config.py

@@ -4007,7 +4007,7 @@ def _process_module_builtin_defaults():
         "instrument_rest_framework_decorators",
     )
 
-    _process_module_definition("celery.app.base", "newrelic.hooks.application_celery", "instrument_celery_app_base")
+    # _process_module_definition("celery.app.base", "newrelic.hooks.application_celery", "instrument_celery_app_base")
     _process_module_definition("celery.local", "newrelic.hooks.application_celery", "instrument_celery_local")
     # _process_module_definition("celery.task.base", "newrelic.hooks.application_celery", "instrument_celery_app_task")
     # _process_module_definition("celery.app.task", "newrelic.hooks.application_celery", "instrument_celery_app_task")

newrelic/hooks/application_celery.py

@@ -95,6 +95,10 @@ def wrap_task_call(wrapped, instance, args, kwargs):
 
     elif transaction:
         with FunctionTrace(_name, source=_source):
+            # Do we accept distributed tracing headers here?
+            # request = wrapped.request
+            # headers = getattr(request, "headers", None) or vars(request)
+            # transaction.accept_distributed_trace_headers(headers, transport_type="AMQP")
             return wrapped(*args, **kwargs)
 
     else:
@@ -104,14 +108,28 @@ def wrap_task_call(wrapped, instance, args, kwargs):
                 # Headers on earlier versions of Celery may end up as attributes
                 # on the request context instead of as custom headers. Handler this
                 # by defaulting to using vars() if headers is not available
-                # print(f"wrapped.request: {wrapped.request}")
                 request = wrapped.request
                 headers = getattr(request, "headers", None) or vars(request)
 
                 settings = transaction.settings
                 if headers is not None and settings is not None:
                     if settings.distributed_tracing.enabled:
-                        transaction.accept_distributed_trace_headers(headers, transport_type="AMQP")
+                        if not transaction.accept_distributed_trace_headers(headers, transport_type="AMQP"):
+                            try:
+                                dt_headers = MessageTrace.generate_request_headers(transaction)
+                                # original_headers = kwargs.get("headers", None)
+                                if dt_headers:
+                                    if not headers:
+                                        wrapped.request.headers = dict(dt_headers)
+                                        # kwargs["headers"] = dict(dt_headers)
+                                    else:
+                                        headers.update(dict(dt_headers))
+                                        wrapped.request.headers = headers
+                                        # wrapped.request.headers.update(dict(dt_headers))
+                                        # kwargs["headers"] = dt_headers = dict(dt_headers)
+                                        # dt_headers.update(dict(headers))
+                            except Exception:
+                                pass
                     elif transaction.settings.cross_application_tracer.enabled:
                         transaction._process_incoming_cat_headers(
                             headers.get(MessageTrace.cat_id_key, None),
@@ -133,6 +151,7 @@ def run(self, *args, **kwargs):
         task = bound_args.get("task", None)
 
         task = TaskWrapper(task, wrap_task_call)
+        task.__module__ = wrapped.__module__  # Ensure module is set for monkeypatching detection
         bound_args["task"] = task
 
         return wrapped(**bound_args)
@@ -189,7 +208,17 @@ def wrapper(wrapped, instance, args, kwargs):
                     settings = transaction.settings
                     if headers is not None and settings is not None:
                         if settings.distributed_tracing.enabled:
-                            transaction.accept_distributed_trace_headers(headers, transport_type="AMQP")
+                            if not transaction.accept_distributed_trace_headers(headers, transport_type="AMQP"):
+                                try:
+                                    dt_headers = MessageTrace.generate_request_headers(transaction)
+                                    if dt_headers:
+                                        if not headers:
+                                            instance.request.headers = dict(dt_headers)
+                                        else:
+                                            headers.update(dict(dt_headers))
+                                            instance.request.headers = headers
+                                except Exception:
+                                    pass
                         elif transaction.settings.cross_application_tracer.enabled:
                             transaction._process_incoming_cat_headers(
                                 headers.get(MessageTrace.cat_id_key, None),
@@ -237,29 +266,29 @@ def run(self, *args, **kwargs):
     return wrapped_task
 
 
-# This will not work with the current version of Celery
-# This only gets called during the async execution of a task
-# and the task is wrapped later in the process to accomodate
-# custom task classes.
-def wrap_Celery_send_task(wrapped, instance, args, kwargs):
-    transaction = current_transaction()
-    if not transaction:
-        return wrapped(*args, **kwargs)
-
-    # Merge distributed tracing headers into outgoing task headers
-    try:
-        dt_headers = MessageTrace.generate_request_headers(transaction)
-        original_headers = kwargs.get("headers", None)
-        if dt_headers:
-            if not original_headers:
-                kwargs["headers"] = dict(dt_headers)
-            else:
-                kwargs["headers"] = dt_headers = dict(dt_headers)
-                dt_headers.update(dict(original_headers))
-    except Exception:
-        pass
-
-    return wrapped(*args, **kwargs)
+# # This will not work with the current version of Celery
+# # This only gets called during the async execution of a task
+# # and the task is wrapped later in the process to accomodate
+# # custom task classes.
+# def wrap_Celery_send_task(wrapped, instance, args, kwargs):
+#     transaction = current_transaction()
+#     if not transaction:
+#         return wrapped(*args, **kwargs)
+
+#     # Merge distributed tracing headers into outgoing task headers
+#     try:
+#         dt_headers = MessageTrace.generate_request_headers(transaction)
+#         original_headers = kwargs.get("headers", None)
+#         if dt_headers:
+#             if not original_headers:
+#                 kwargs["headers"] = dict(dt_headers)
+#             else:
+#                 kwargs["headers"] = dt_headers = dict(dt_headers)
+#                 dt_headers.update(dict(original_headers))
+#     except Exception:
+#         pass
+
+#     return wrapped(*args, **kwargs)
 
 
 def wrap_worker_optimizations(wrapped, instance, args, kwargs):
@@ -290,9 +319,9 @@ def instrument_celery_local(module):
         module.Proxy.__call__ = CeleryTaskWrapper(module.Proxy.__call__)
 
 
-def instrument_celery_app_base(module):
-    if hasattr(module, "Celery") and hasattr(module.Celery, "send_task"):
-        wrap_function_wrapper(module, "Celery.send_task", wrap_Celery_send_task)
+# def instrument_celery_app_base(module):
+#     if hasattr(module, "Celery") and hasattr(module.Celery, "send_task"):
+#         wrap_function_wrapper(module, "Celery.send_task", wrap_Celery_send_task)
 
 
 def instrument_celery_worker(module):
@@ -347,4 +376,5 @@ def instrument_celery_app_trace(module):
 
     if hasattr(module, "build_tracer"):
         wrap_function_wrapper(module, "build_tracer", wrap_build_tracer)
+        
 

tests/application_celery/test_application.py

@@ -24,9 +24,11 @@
 @validate_transaction_metrics(
     name="test_application:test_celery_task_as_function_trace",
     scoped_metrics=[("Function/_target_application.add", 1)],
+    # scoped_metrics=[("Function/_target_application.add", 2)],
     background_task=True,
 )
 @validate_code_level_metrics("_target_application", "add")
+# @validate_code_level_metrics("_target_application", "add", count=2)
 @background_task()
 def test_celery_task_as_function_trace():
     """

tests/application_celery/test_distributed_tracing.py

@@ -12,52 +12,103 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from _target_application import add, assert_dt
+import pytest
+from _target_application import add
 from testing_support.fixtures import override_application_settings
 from testing_support.validators.validate_transaction_count import validate_transaction_count
 from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics
 
 from newrelic.api.background_task import background_task
 
 
-@validate_transaction_metrics(
-    name="_target_application.assert_dt",
-    group="Celery",
-    rollup_metrics=[
-        ("Supportability/DistributedTrace/AcceptPayload/Success", None),
-        ("Supportability/TraceContext/Accept/Success", 1),
-    ],
-    background_task=True,
-    index=-2,
-)
-@validate_transaction_metrics(
-    name="test_distributed_tracing:test_celery_task_distributed_tracing_enabled", background_task=True
-)
-@validate_transaction_count(2)
-@background_task()
-def test_celery_task_distributed_tracing_enabled():
-    result = assert_dt.apply_async()
-    result = result.get()
-    assert result == 1
-
-
-@override_application_settings({"distributed_tracing.enabled": False})
-@validate_transaction_metrics(
-    name="_target_application.add",
-    group="Celery",
-    rollup_metrics=[
-        ("Supportability/DistributedTrace/AcceptPayload/Success", None),
-        ("Supportability/TraceContext/Accept/Success", None),  # No trace context should be accepted
-    ],
-    background_task=True,
-    index=-2,
-)
-@validate_transaction_metrics(
-    name="test_distributed_tracing:test_celery_task_distributed_tracing_disabled", background_task=True
-)
-@validate_transaction_count(2)
-@background_task()
-def test_celery_task_distributed_tracing_disabled():
-    result = add.apply_async((1, 2))
-    result = result.get()
-    assert result == 3
+@pytest.mark.parametrize("dt_enabled", [True, False])
+def test_celery_task_distributed_tracing_inside_background_task(dt_enabled):
+    @override_application_settings({"distributed_tracing.enabled": dt_enabled})
+    @validate_transaction_metrics(
+        name="_target_application.add",
+        group="Celery",
+        rollup_metrics=[
+            ("Supportability/TraceContext/Accept/Success", 1 if dt_enabled else None),
+            ("Supportability/TraceContext/TraceParent/Accept/Success", 1 if dt_enabled else None),
+        ],
+        background_task=True,
+        index=-2,
+    )
+    @validate_transaction_metrics(
+        name="test_distributed_tracing:test_celery_task_distributed_tracing_inside_background_task.<locals>._test",
+        rollup_metrics=[
+            ("Supportability/TraceContext/Create/Success", 1 if dt_enabled else None),
+            ("Supportability/DistributedTrace/CreatePayload/Success", 1 if dt_enabled else None),
+        ],
+        background_task=True,
+    )
+    @validate_transaction_count(2)  # One for the background task, one for the Celery task.  Runs in different processes.
+    @background_task()
+    def _test():
+        result = add.apply_async((1, 2))
+        result = result.get()
+        assert result == 3
+
+    _test()
+    
+    
+@pytest.mark.parametrize("dt_enabled", [True, False])
+def test_celery_task_distributed_tracing_outside_background_task(dt_enabled):
+    @override_application_settings({"distributed_tracing.enabled": dt_enabled})
+    @validate_transaction_metrics(
+        name="_target_application.add",
+        group="Celery",
+        rollup_metrics=[
+            ("Supportability/TraceContext/Create/Success", 1 if dt_enabled else None),
+            ("Supportability/DistributedTrace/CreatePayload/Success", 1 if dt_enabled else None),
+        ],
+        background_task=True,
+    )
+    @validate_transaction_count(1)
+    def _test():
+        result = add.apply_async((1, 2))
+        result = result.get()
+        assert result == 3
+
+    _test()
+    
+
+# In this case, the background task creating the transaction 
+# has not generated a distributed trace header, so the Celery 
+# task will not have a distributed trace header to accept.
+@pytest.mark.parametrize("dt_enabled", [True, False])
+def test_celery_task_distributed_tracing_inside_background_task_apply(dt_enabled):
+    @override_application_settings({"distributed_tracing.enabled": dt_enabled})
+    @validate_transaction_metrics(
+        name="test_distributed_tracing:test_celery_task_distributed_tracing_inside_background_task_apply.<locals>._test",
+        background_task=True,
+    )
+    @validate_transaction_count(1)  # In the same process, so only one transaction
+    @background_task()
+    def _test():
+        result = add.apply((1, 2))
+        result = result.get()
+        assert result == 3
+
+    _test()
+    
+    
+@pytest.mark.parametrize("dt_enabled", [True, False])
+def test_celery_task_distributed_tracing_outside_background_task_apply(dt_enabled):
+    @override_application_settings({"distributed_tracing.enabled": dt_enabled})
+    @validate_transaction_metrics(
+        name="_target_application.add",
+        group="Celery",
+        rollup_metrics=[
+            ("Supportability/TraceContext/Create/Success", 1 if dt_enabled else None),
+            ("Supportability/DistributedTrace/CreatePayload/Success", 1 if dt_enabled else None),
+        ],
+        background_task=True,
+    )
+    @validate_transaction_count(1)  # In the same process, so only one transaction
+    def _test():
+        result = add.apply((1, 2))
+        result = result.get()
+        assert result == 3
+
+    _test()