celery: give sampler visibility on spans attributes

.changelog/5312.changed

@@ -0,0 +1 @@
+`opentelemetry-instrumentation-celery`: pass known attributes at span creation for sampler visibility

instrumentation/opentelemetry-instrumentation-celery/src/opentelemetry/instrumentation/celery/__init__.py

@@ -188,16 +188,26 @@ def _trace_prerun(self, *args, **kwargs):
         logger.debug("prerun signal start task_id=%s", task_id)
 
         operation_name = f"{_TASK_RUN}/{task.name}"
+        attributes = {
+            _TASK_TAG_KEY: _TASK_RUN,
+            **utils.get_attributes_from_context(kwargs),
+            **utils.get_attributes_from_context(request),
+            _TASK_NAME_KEY: task.name,
+        }
 
         if self._use_span_links and tracectx is not None:
             parent_span_context = trace.get_current_span(tracectx).get_span_context()
             links = [trace.Link(parent_span_context)] if parent_span_context.is_valid else None
-            span = self._tracer.start_span(operation_name, links=links, kind=trace.SpanKind.CONSUMER)
+            span = self._tracer.start_span(
+                operation_name, links=links, kind=trace.SpanKind.CONSUMER, attributes=attributes
+            )
             # Don't attach the context when using links to avoid parent-child relationship
             token = None
         else:
             token = context_api.attach(tracectx) if tracectx is not None else None
-            span = self._tracer.start_span(operation_name, context=tracectx, kind=trace.SpanKind.CONSUMER)
+            span = self._tracer.start_span(
+                operation_name, context=tracectx, kind=trace.SpanKind.CONSUMER, attributes=attributes
+            )
 
         activation = trace.use_span(span, end_on_exit=True)
         activation.__enter__()  # pylint: disable=unnecessary-dunder-call
@@ -221,12 +231,9 @@ def _trace_postrun(self, *args, **kwargs):
 
         span, activation, token = ctx
 
-        # request context tags
+        # The postrun signal supplies attributes only known after execution, such as state.
         if span.is_recording():
-            span.set_attribute(_TASK_TAG_KEY, _TASK_RUN)
-            utils.set_attributes_from_context(span, kwargs)
-            utils.set_attributes_from_context(span, task.request)
-            span.set_attribute(_TASK_NAME_KEY, task.name)
+            span.set_attributes(utils.get_attributes_from_context(kwargs))
 
         activation.__exit__(None, None, None)
         utils.detach_context(task, task_id)
@@ -254,14 +261,13 @@ def _trace_before_publish(self, *args, **kwargs):
         else:
             task_name = task.name
         operation_name = f"{_TASK_APPLY_ASYNC}/{task_name}"
-        span = self._tracer.start_span(operation_name, kind=trace.SpanKind.PRODUCER)
-
-        # apply some attributes here because most of the data is not available
-        if span.is_recording():
-            span.set_attribute(_TASK_TAG_KEY, _TASK_APPLY_ASYNC)
-            span.set_attribute(MESSAGING_MESSAGE_ID, task_id)
-            span.set_attribute(_TASK_NAME_KEY, task_name)
-            utils.set_attributes_from_context(span, kwargs)
+        attributes = {
+            _TASK_TAG_KEY: _TASK_APPLY_ASYNC,
+            MESSAGING_MESSAGE_ID: task_id,
+            _TASK_NAME_KEY: task_name,
+            **utils.get_attributes_from_context(kwargs),
+        }
+        span = self._tracer.start_span(operation_name, kind=trace.SpanKind.PRODUCER, attributes=attributes)
 
         activation = trace.use_span(span, end_on_exit=True)
         activation.__enter__()  # pylint: disable=unnecessary-dunder-call

instrumentation/opentelemetry-instrumentation-celery/src/opentelemetry/instrumentation/celery/utils.py

@@ -14,8 +14,9 @@
 from opentelemetry.semconv._incubating.attributes.messaging_attributes import (
     MESSAGING_MESSAGE_ID,
 )
-from opentelemetry.semconv.trace import SpanAttributes
+from opentelemetry.semconv.trace import MessagingDestinationKindValues, SpanAttributes
 from opentelemetry.trace import Span
+from opentelemetry.util.types import AttributeValue
 
 if TYPE_CHECKING:
     from contextlib import AbstractContextManager
@@ -59,13 +60,11 @@ def get(self, key: str, default: Any = None) -> Any: ...
 
 
 # pylint:disable=too-many-branches
-def set_attributes_from_context(
-    span: Span,
+def get_attributes_from_context(
     context: ContextCarrier,
-) -> None:
+) -> dict[str, AttributeValue]:
     """Helper to extract meta values from a Celery Context"""
-    if not span.is_recording():
-        return
+    attributes: dict[str, AttributeValue] = {}
     for key in CELERY_CONTEXT_ATTRIBUTES:
         value = context.get(key)
 
@@ -92,7 +91,7 @@ def set_attributes_from_context(
             routing_key = value.get("routing_key")
 
             if routing_key is not None:
-                span.set_attribute(SpanAttributes.MESSAGING_DESTINATION, routing_key)
+                attributes[SpanAttributes.MESSAGING_DESTINATION] = routing_key
 
             value = str(value)
 
@@ -110,10 +109,10 @@ def set_attributes_from_context(
             attribute_name = SpanAttributes.MESSAGING_DESTINATION_KIND
             for declare in value:
                 if declare.exchange.type == "direct":
-                    value = "queue"
+                    value = MessagingDestinationKindValues.QUEUE.value
                     break
                 if declare.exchange.type == "topic":
-                    value = "topic"
+                    value = MessagingDestinationKindValues.TOPIC.value
                     break
 
         # set attribute name if not set specially for a key
@@ -125,7 +124,9 @@ def set_attributes_from_context(
             if value == "":
                 continue
 
-        span.set_attribute(attribute_name, value)
+        attributes[attribute_name] = value
+
+    return attributes
 
 
 def attach_context(

instrumentation/opentelemetry-instrumentation-celery/tests/test_tasks.py

@@ -4,19 +4,21 @@
 import threading
 import time
 
-from celery.signals import after_task_publish, before_task_publish
-from kombu import Exchange
+from celery.signals import after_task_publish, before_task_publish, task_postrun, task_prerun
+from kombu import Exchange, Queue
 from wrapt import wrap_function_wrapper
 
 from opentelemetry import baggage, context
 from opentelemetry.instrumentation.celery import CeleryInstrumentor, utils
 from opentelemetry.instrumentation.utils import unwrap
+from opentelemetry.semconv._incubating.attributes.messaging_attributes import MESSAGING_MESSAGE_ID
 from opentelemetry.semconv.attributes.exception_attributes import (
     EXCEPTION_MESSAGE,
     EXCEPTION_STACKTRACE,
     EXCEPTION_TYPE,
 )
 from opentelemetry.semconv.trace import SpanAttributes
+from opentelemetry.test.samplertestutil import CapturingSampler
 from opentelemetry.test.test_base import TestBase
 from opentelemetry.trace import SpanKind, StatusCode
 
@@ -54,6 +56,130 @@ def test_before_task_publish_with_kombu_exchange(self):
         self.assertEqual(spans[0].attributes["celery.exchange"], "celery_delayed_27")
 
 
+class TestCelerySampling(TestBase):
+    def tearDown(self) -> None:
+        CeleryInstrumentor().uninstrument()
+        super().tearDown()
+
+    def _instrument_with_sampler(self, sampler: CapturingSampler, use_span_links: bool = False) -> None:
+        self.tracer_provider, self.memory_exporter = self.create_tracer_provider(sampler=sampler)
+        self.addCleanup(self.tracer_provider.shutdown)
+        instrumentor = CeleryInstrumentor()
+        if instrumentor.is_instrumented_by_opentelemetry:
+            instrumentor.uninstrument()
+        instrumentor.instrument(tracer_provider=self.tracer_provider, use_span_links=use_span_links)
+
+    def test_publish_attributes_passed_to_sampler(self) -> None:
+        sampler = CapturingSampler()
+        self._instrument_with_sampler(sampler)
+        signal_kwargs = {
+            "sender": task_add.name,
+            "headers": {"id": "task-id"},
+            "body": None,
+            "exchange": Exchange("celery", type="direct"),
+            "routing_key": "celery",
+            "declare": [Queue("celery", exchange=Exchange("celery", type="direct"))],
+            "priority": 4,
+        }
+
+        before_task_publish.send(**signal_kwargs)
+        after_task_publish.send(**signal_kwargs)
+
+        expected = {
+            "celery.action": "apply_async",
+            MESSAGING_MESSAGE_ID: "task-id",
+            "celery.task_name": task_add.name,
+            "celery.exchange": "celery",
+            SpanAttributes.MESSAGING_DESTINATION: "celery",
+            SpanAttributes.MESSAGING_DESTINATION_KIND: "queue",
+            "celery.priority": 4,
+        }
+        self.assertEqual(sampler.kind, SpanKind.PRODUCER)
+        self.assertEqual(sampler.attributes, expected)
+        for key, value in expected.items():
+            self.assertIs(type(sampler.attributes[key]), type(value))
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
+        self.assertEqual(spans[0].attributes, expected)
+
+    def test_prerun_attributes_passed_to_sampler(self) -> None:
+        for use_span_links in (False, True):
+            with self.subTest(use_span_links=use_span_links):
+                sampler = CapturingSampler()
+                self._instrument_with_sampler(sampler, use_span_links=use_span_links)
+                task_add.push_request(
+                    id="task-id",
+                    correlation_id="correlation-id",
+                    hostname="worker",
+                    origin="publisher",
+                    delivery_info={"routing_key": "celery"},
+                    retries=2,
+                    timelimit=(None, 10),
+                    traceparent="00-03afa25236b8cd948fa853d67038ac79-405ff022e8247c46-01",
+                )
+                self.addCleanup(task_add.pop_request)
+                signal_kwargs = {"sender": task_add, "task": task_add, "task_id": "task-id"}
+
+                task_prerun.send(**signal_kwargs)
+                task_postrun.send(**signal_kwargs, state="SUCCESS")
+
+                expected = {
+                    "celery.action": "run",
+                    "celery.task_name": task_add.name,
+                    MESSAGING_MESSAGE_ID: "task-id",
+                    SpanAttributes.MESSAGING_CONVERSATION_ID: "correlation-id",
+                    "celery.hostname": "worker",
+                    "celery.origin": "publisher",
+                    "celery.delivery_info": str({"routing_key": "celery"}),
+                    SpanAttributes.MESSAGING_DESTINATION: "celery",
+                    "celery.retries": 2,
+                    "celery.timelimit": ["", "10"],
+                }
+                self.assertEqual(sampler.kind, SpanKind.CONSUMER)
+                self.assertEqual(sampler.attributes, expected)
+                for key, value in expected.items():
+                    self.assertIs(type(sampler.attributes[key]), type(value))
+                spans = self.memory_exporter.get_finished_spans()
+                self.assertEqual(len(spans), 1)
+                expected["celery.timelimit"] = ("", "10")
+                expected["celery.state"] = "SUCCESS"
+                self.assertEqual(spans[0].attributes, expected)
+                if use_span_links:
+                    self.assertIsNone(spans[0].parent)
+                    self.assertEqual(len(spans[0].links), 1)
+                else:
+                    self.assertEqual(spans[0].parent.span_id, 0x405FF022E8247C46)
+
+    def test_attributes_control_sampling_apply_async(self) -> None:
+        for routing_key, expected_count in (("sample", 1), ("drop", 0)):
+            with self.subTest(routing_key=routing_key):
+                sampler = CapturingSampler(required_attributes={SpanAttributes.MESSAGING_DESTINATION: "sample"})
+                self._instrument_with_sampler(sampler)
+                signal_kwargs = {
+                    "sender": task_add.name,
+                    "headers": {"id": "task-id"},
+                    "body": None,
+                    "routing_key": routing_key,
+                }
+                before_task_publish.send(**signal_kwargs)
+                after_task_publish.send(**signal_kwargs)
+                self.assertEqual(len(self.memory_exporter.get_finished_spans()), expected_count)
+                self.assertIsNone(utils.retrieve_context(task_add, "task-id", is_publish=True))
+
+    def test_attributes_control_sampling_run(self) -> None:
+        for routing_key, expected_count in (("sample", 1), ("drop", 0)):
+            with self.subTest(routing_key=routing_key):
+                sampler = CapturingSampler(required_attributes={SpanAttributes.MESSAGING_DESTINATION: "sample"})
+                self._instrument_with_sampler(sampler)
+                task_add.push_request(id="task-id", hostname="worker", delivery_info={"routing_key": routing_key})
+                self.addCleanup(task_add.pop_request)
+                signal_kwargs = {"sender": task_add, "task": task_add, "task_id": "task-id"}
+                task_prerun.send(**signal_kwargs)
+                task_postrun.send(**signal_kwargs, state="SUCCESS")
+                self.assertEqual(len(self.memory_exporter.get_finished_spans()), expected_count)
+                self.assertIsNone(utils.retrieve_context(task_add, "task-id", is_publish=False))
+
+
 class TestCeleryInstrumentation(TestBase):
     def setUp(self):
         super().setUp()