fix(celery): account for a chains scenario

ddtrace/contrib/internal/celery/app.py

@@ -133,10 +133,6 @@ def _traced_apply_async_inner(func, instance, args, kwargs):
                 if task_span:
                     task_span.set_exc_info(*sys.exc_info())
 
-                prerun_span = core.get_item("prerun_span")
-                if prerun_span:
-                    prerun_span.set_exc_info(*sys.exc_info())
-
                 raise
             finally:
                 task_span = core.get_item("task_span")
@@ -147,11 +143,4 @@ def _traced_apply_async_inner(func, instance, args, kwargs):
                     )
                     task_span.finish()
 
-                prerun_span = core.get_item("prerun_span")
-                if prerun_span:
-                    log.debug(
-                        "The task_postrun signal was not called, so manually closing span: %s", prerun_span._pprint()
-                    )
-                    prerun_span.finish()
-
     return _traced_apply_async_inner

ddtrace/contrib/internal/celery/signals.py

@@ -52,9 +52,6 @@ def trace_prerun(*args, **kwargs):
     service = config.celery["worker_service_name"]
     span = pin.tracer.trace(c.WORKER_ROOT_SPAN, service=service, resource=task.name, span_type=SpanTypes.WORKER)
 
-    # Store an item called "prerun span" in case task_postrun doesn't get called
-    core.set_item("prerun_span", span)
-
     # set span.kind to the type of request being performed
     span.set_tag_str(SPAN_KIND, SpanKind.CONSUMER)
 

tests/contrib/celery/test_integration.py

@@ -477,6 +477,32 @@ def fn_a(self, user, force_logout=False):
         trace_id = traces[0][0].trace_id
         assert all(trace_id == span.trace_id for trace in traces for span in trace)
 
+    # Ensure that when we call Celery chains, the root span has celery specific span tags
+    def test_task_chain_task_call_task(self):
+        @self.app.task
+        def fn_a():
+            return 1
+
+        @self.app.task
+        def fn_b():
+            return 2
+
+        traces = None
+        try:
+            # See https://github.com/DataDog/dd-trace-py/issues/11479
+            (fn_a.si() | fn_b.si()).delay()
+            import time
+
+            time.sleep(10)
+        except Exception:
+            pass
+
+        traces = self.pop_traces()
+        assert len(traces) == 3
+        # Assert that each celery related span has a celery specific tag like celery.action
+        # TODO figure out why this test passes on the main branch in pytest but not a real app
+        assert all("celery.action" in span._meta.keys() for trace in traces for span in trace)
+
     @mock.patch("kombu.messaging.Producer.publish", mock.Mock(side_effect=ValueError))
     def test_fn_task_apply_async_soft_exception(self):
         # If the underlying library runs into an exception that doesn't crash the app