Reviewer suggestions and more tests

Author: lrafeei

newrelic/hooks/application_celery.py

@@ -64,6 +64,26 @@ def task_info(instance, *args, **kwargs):
     return task_name, task_source
 
 
+# def create_DT_headers_and_attach(wrapped, transaction):
+#     try:
+#         request = wrapped and wrapped.request_stack and wrapped.request_stack.top
+#         original_headers = getattr(request, "headers", None) or vars(request)
+#         dt_headers = MessageTrace.generate_request_headers(transaction)
+#         if dt_headers:
+#             if not original_headers:
+#                 wrapped.request.headers = dict(dt_headers)
+#             else:
+#                 dt_headers.update(original_headers)
+#                 wrapped.request.headers = dt_headers = dict(dt_headers)
+#                 # original_headers.update(dict(dt_headers))
+#                 # wrapped.request.headers = original_headers
+
+#     except Exception:
+#         pass
+
+#     return wrapped
+
+
 def wrap_task_call(wrapped, instance, args, kwargs):
     transaction = current_transaction(active_only=False)
 
@@ -104,22 +124,23 @@ def wrap_task_call(wrapped, instance, args, kwargs):
             try:
                 # Headers on earlier versions of Celery may end up as attributes
                 # on the request context instead of as custom headers. Handle this
-                # by defaulting to using vars() if headers is not available
-                request = wrapped.request
+                # by defaulting to using `vars()` if headers is not available
+
+                # If there is no request, the request property will return
+                # a new instance of `celery.Context()` instead of `None`, so
+                # this will be handled by accessing the request_stack directly.
+                request = wrapped and wrapped.request_stack and wrapped.request_stack.top
                 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:
+                        # Generate DT headers if they do not already exist in the incoming request
                         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:
-                                        wrapped.request.headers = dict(dt_headers)
-                                    else:
-                                        headers.update(dict(dt_headers))
-                                        wrapped.request.headers = headers
+                                    headers.update(dict(dt_headers))
                             except Exception:
                                 pass
                     elif transaction.settings.cross_application_tracer.enabled:
@@ -143,7 +164,6 @@ 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)
@@ -193,22 +213,23 @@ def wrapper(wrapped, instance, args, kwargs):
                 try:
                     # Headers on earlier versions of Celery may end up as attributes
                     # on the request context instead of as custom headers. Handle this
-                    # by defaulting to using vars() if headers is not available
-                    request = instance.request
+                    # by defaulting to using `vars()` if headers is not available
+
+                    # If there is no request, the request property will return
+                    # a new instance of `celery.Context()` instead of `None`, so
+                    # this will be handled by accessing the request_stack directly.
+                    request = instance and instance.request_stack and instance.request_stack.top
                     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:
+                            # Generate DT headers if they do not already exist in the incoming request
                             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
+                                        headers.update(dict(dt_headers))
                                 except Exception:
                                     pass
                         elif transaction.settings.cross_application_tracer.enabled:
@@ -238,7 +259,7 @@ def wrapper(wrapped, instance, args, kwargs):
     # Celery has included a monkey-patching provision which did not perform this
     # optimization on functions that were monkey-patched. Unfortunately, our
     # wrappers are too transparent for celery to detect that they've even been
-    # monky-patched. To circumvent this, we set the __module__ of our wrapped task
+    # monkey-patched. To circumvent this, we set the __module__ of our wrapped task
     # to this file which causes celery to properly detect that it has been patched.
     #
     # For versions of celery 2.5.3 to 2.5.5
@@ -252,7 +273,6 @@ def run(self, *args, **kwargs):
             return self.__call__(*args, **kwargs)
 
     wrapped_task = TaskWrapper(wrapped, wrapper)
-    # Reset __module__ to be less transparent so celery detects our monkey-patching
     wrapped_task.__module__ = CeleryTaskWrapper.__module__
 
     return wrapped_task
@@ -262,7 +282,7 @@ def instrument_celery_local(module):
     if hasattr(module, "Proxy"):
         # This is used in the case where the function is
         # called directly on the Proxy object (rather than
-        # using "delay" or "apply_async")
+        # using `delay` or `apply_async`)
         module.Proxy.__call__ = CeleryTaskWrapper(module.Proxy.__call__)
 
 

tests/application_celery/_target_application.py

@@ -26,14 +26,33 @@
 )
 
 
+@app.task
+def add(x, y):
+    return x + y
+
+
+@app.task
+def tsum(nums):
+    return sum(nums)
+
+
+@app.task
+def nested_add(x, y):
+    return add(x, y)
+
+
+@shared_task
+def shared_task_add(x, y):
+    return x + y
+
+
 class CustomCeleryTaskWithSuper(Task):
     def __call__(self, *args, **kwargs):
         transaction = current_transaction()
         if transaction:
             transaction.add_custom_attribute("custom_task_attribute", "Called with super")
         return super().__call__(*args, **kwargs)
 
-
 class CustomCeleryTaskWithRun(Task):
     def __call__(self, *args, **kwargs):
         transaction = current_transaction()
@@ -42,31 +61,11 @@ def __call__(self, *args, **kwargs):
         return self.run(*args, **kwargs)
 
 
-@app.task
-def add(x, y):
-    return x + y
-
-
 @app.task(base=CustomCeleryTaskWithSuper)
 def add_with_super(x, y):
     return x + y
 
 
 @app.task(base=CustomCeleryTaskWithRun)
 def add_with_run(x, y):
-    return x + y
-
-
-@app.task
-def tsum(nums):
-    return sum(nums)
-
-
-@app.task
-def nested_add(x, y):
-    return add(x, y)
-
-
-@shared_task
-def shared_task_add(x, y):
-    return x + y
+    return x + y

tests/application_celery/conftest.py

@@ -11,7 +11,9 @@
 # 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.
+
 import pytest
+from celery.app.trace import setup_worker_optimizations, reset_worker_optimizations
 from testing_support.fixtures import collector_agent_registration_fixture, collector_available_fixture
 
 _default_settings = {
@@ -27,7 +29,6 @@
     app_name="Python Agent Test (application_celery)", default_settings=_default_settings
 )
 
-
 @pytest.fixture(scope="session")
 def celery_config():
     # Used by celery pytest plugin to configure Celery instance
@@ -43,3 +44,12 @@ def celery_worker_parameters():
 @pytest.fixture(scope="session", autouse=True)
 def celery_worker_available(celery_session_worker):
     return celery_session_worker
+
+
+@pytest.fixture(scope="session", autouse=True, params=[False, True], ids=["unpatched", "patched"])
+def with_worker_optimizations(request, celery_worker_available):
+    if request.param:
+        setup_worker_optimizations(celery_worker_available.app)
+
+    yield request.param
+    reset_worker_optimizations()

tests/application_celery/test_distributed_tracing.py

@@ -19,10 +19,11 @@
 from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics
 
 from newrelic.api.background_task import background_task
+from newrelic.api.transaction import insert_distributed_trace_headers
 
 
 @pytest.mark.parametrize("dt_enabled", [True, False])
-def test_celery_task_distributed_tracing_inside_background_task(dt_enabled):
+def test_DT_inside_transaction_delay(dt_enabled):
     @override_application_settings({"distributed_tracing.enabled": dt_enabled})
     @validate_transaction_metrics(
         name="_target_application.add",
@@ -35,7 +36,7 @@ def test_celery_task_distributed_tracing_inside_background_task(dt_enabled):
         index=-2,
     )
     @validate_transaction_metrics(
-        name="test_distributed_tracing:test_celery_task_distributed_tracing_inside_background_task.<locals>._test",
+        name="test_distributed_tracing:test_DT_inside_transaction_delay.<locals>._test",
         rollup_metrics=[
             ("Supportability/TraceContext/Create/Success", 1 if dt_enabled else None),
             ("Supportability/DistributedTrace/CreatePayload/Success", 1 if dt_enabled else None),
@@ -46,15 +47,15 @@ def test_celery_task_distributed_tracing_inside_background_task(dt_enabled):
     # 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 = add.delay(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):
+def test_DT_outside_transaction_delay(dt_enabled):
     @override_application_settings({"distributed_tracing.enabled": dt_enabled})
     @validate_transaction_metrics(
         name="_target_application.add",
@@ -67,21 +68,23 @@ def test_celery_task_distributed_tracing_outside_background_task(dt_enabled):
     )
     @validate_transaction_count(1)
     def _test():
-        result = add.apply_async((1, 2))
+        result = add.delay(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):
+def test_DT_inside_transaction_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",
+        name="test_distributed_tracing:test_DT_inside_transaction_apply.<locals>._test",
+        rollup_metrics=[
+            ("Function/_target_application.add", 1),
+        ],
+        scoped_metrics=[
+            ("Function/_target_application.add", 1),
+        ],
         background_task=True,
     )
     @validate_transaction_count(1)  # In the same process, so only one transaction
@@ -95,7 +98,34 @@ def _test():
 
 
 @pytest.mark.parametrize("dt_enabled", [True, False])
-def test_celery_task_distributed_tracing_outside_background_task_apply(dt_enabled):
+def test_DT_inside_transaction_apply_with_added_headers(dt_enabled):
+    @override_application_settings({"distributed_tracing.enabled": dt_enabled})
+    @validate_transaction_metrics(
+        name="test_distributed_tracing:test_DT_inside_transaction_apply_with_added_headers.<locals>._test",
+        rollup_metrics=[
+            ("Function/_target_application.add", 1),
+            ("Supportability/TraceContext/Create/Success", 1 if dt_enabled else None),
+            ("Supportability/DistributedTrace/CreatePayload/Success", 1 if dt_enabled else None),
+        ],
+        scoped_metrics=[
+            ("Function/_target_application.add", 1),
+        ],
+        background_task=True,
+    )
+    @validate_transaction_count(1)  # In the same process, so only one transaction
+    @background_task()
+    def _test():
+        headers = []
+        insert_distributed_trace_headers(headers)
+        result = add.apply((1, 2), headers=headers)
+        result = result.get()
+        assert result == 3
+
+    _test()
+
+
+@pytest.mark.parametrize("dt_enabled", [True, False])
+def test_DT_outside_transaction_apply(dt_enabled):
     @override_application_settings({"distributed_tracing.enabled": dt_enabled})
     @validate_transaction_metrics(
         name="_target_application.add",
@@ -113,3 +143,41 @@ def _test():
         assert result == 3
 
     _test()
+
+
+@pytest.mark.parametrize("dt_enabled", [True, False])
+def test_DT_inside_transaction__call__(dt_enabled):
+    @override_application_settings({"distributed_tracing.enabled": dt_enabled})
+    @validate_transaction_metrics(
+        name="test_distributed_tracing:test_DT_inside_transaction__call__.<locals>._test",
+        rollup_metrics=[
+            ("Function/_target_application.add", 1),
+        ],
+        scoped_metrics=[
+            ("Function/_target_application.add", 1),
+        ],
+        background_task=True,
+    )
+    @validate_transaction_count(1)  # In the same process, so only one transaction
+    @background_task()
+    def _test():
+        result = add(1, 2)
+        assert result == 3
+
+    _test()
+    
+    
+@pytest.mark.parametrize("dt_enabled", [True, False])
+def test_DT_outside_transaction__call__(dt_enabled):
+    @override_application_settings({"distributed_tracing.enabled": dt_enabled})
+    @validate_transaction_metrics(
+        name="_target_application.add",
+        group="Celery",
+        background_task=True,
+    )
+    @validate_transaction_count(1)  # In the same process, so only one transaction
+    def _test():
+        result = add(1, 2)
+        assert result == 3
+
+    _test()