diff --git a/.changelog/5312.changed b/.changelog/5312.changed new file mode 100644 index 0000000000..d9ae4c39fa --- /dev/null +++ b/.changelog/5312.changed @@ -0,0 +1 @@ +`opentelemetry-instrumentation-celery`: pass known attributes at span creation for sampler visibility diff --git a/instrumentation/opentelemetry-instrumentation-celery/src/opentelemetry/instrumentation/celery/__init__.py b/instrumentation/opentelemetry-instrumentation-celery/src/opentelemetry/instrumentation/celery/__init__.py index ebca623585..bb232e1ce2 100644 --- a/instrumentation/opentelemetry-instrumentation-celery/src/opentelemetry/instrumentation/celery/__init__.py +++ b/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 diff --git a/instrumentation/opentelemetry-instrumentation-celery/src/opentelemetry/instrumentation/celery/utils.py b/instrumentation/opentelemetry-instrumentation-celery/src/opentelemetry/instrumentation/celery/utils.py index a489ea3f87..3033908484 100644 --- a/instrumentation/opentelemetry-instrumentation-celery/src/opentelemetry/instrumentation/celery/utils.py +++ b/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( diff --git a/instrumentation/opentelemetry-instrumentation-celery/tests/test_tasks.py b/instrumentation/opentelemetry-instrumentation-celery/tests/test_tasks.py index 3e750eb727..6ca3c59aeb 100644 --- a/instrumentation/opentelemetry-instrumentation-celery/tests/test_tasks.py +++ b/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() diff --git a/instrumentation/opentelemetry-instrumentation-celery/tests/test_utils.py b/instrumentation/opentelemetry-instrumentation-celery/tests/test_utils.py index 5904f515b3..da34855b7b 100644 --- a/instrumentation/opentelemetry-instrumentation-celery/tests/test_utils.py +++ b/instrumentation/opentelemetry-instrumentation-celery/tests/test_utils.py @@ -20,7 +20,7 @@ class TestUtils(unittest.TestCase): def setUp(self): self.app = Celery("celery.test_app") - def test_set_attributes_from_context(self): + def test_get_attributes_from_context(self): # it should extract only relevant keys context = { "correlation_id": "44b7f305", @@ -37,41 +37,39 @@ def test_set_attributes_from_context(self): "routing_key": "celery", } - span = trace._Span("name", mock.Mock(spec=trace_api.SpanContext)) - utils.set_attributes_from_context(span, context) + attributes = utils.get_attributes_from_context(context) self.assertEqual( - span.attributes.get(MESSAGING_MESSAGE_ID), + attributes.get(MESSAGING_MESSAGE_ID), "44b7f305", ) self.assertEqual( - span.attributes.get(SpanAttributes.MESSAGING_CONVERSATION_ID), + attributes.get(SpanAttributes.MESSAGING_CONVERSATION_ID), "44b7f305", ) - self.assertEqual(span.attributes.get(SpanAttributes.MESSAGING_DESTINATION), "celery") + self.assertEqual(attributes.get(SpanAttributes.MESSAGING_DESTINATION), "celery") - self.assertEqual(span.attributes["celery.delivery_info"], str({"eager": True})) - self.assertEqual(span.attributes.get("celery.eta"), "soon") - self.assertEqual(span.attributes.get("celery.expires"), "later") - self.assertEqual(span.attributes.get("celery.hostname"), "localhost") - self.assertEqual(span.attributes.get("celery.origin"), "gen8@localhost") + self.assertEqual(attributes["celery.delivery_info"], str({"eager": True})) + self.assertEqual(attributes.get("celery.eta"), "soon") + self.assertEqual(attributes.get("celery.expires"), "later") + self.assertEqual(attributes.get("celery.hostname"), "localhost") + self.assertEqual(attributes.get("celery.origin"), "gen8@localhost") - self.assertEqual(span.attributes.get("celery.reply_to"), "44b7f305") - self.assertEqual(span.attributes.get("celery.retries"), 4) - self.assertEqual(span.attributes.get("celery.timelimit"), ("now", "later")) - self.assertNotIn("custom_meta", span.attributes) + self.assertEqual(attributes.get("celery.reply_to"), "44b7f305") + self.assertEqual(attributes.get("celery.retries"), 4) + self.assertEqual(attributes.get("celery.timelimit"), ("now", "later")) + self.assertNotIn("custom_meta", attributes) - def test_set_attributes_from_context_origin_not_aliased_to_hostname(self): + def test_get_attributes_from_context_origin_not_aliased_to_hostname(self): # `origin` is the publisher's nodename, not the worker hostname context = {"origin": "gen8@localhost"} - span = trace._Span("name", mock.Mock(spec=trace_api.SpanContext)) - utils.set_attributes_from_context(span, context) + attributes = utils.get_attributes_from_context(context) - self.assertEqual(span.attributes.get("celery.origin"), "gen8@localhost") - self.assertNotIn("celery.hostname", span.attributes) + self.assertEqual(attributes.get("celery.origin"), "gen8@localhost") + self.assertNotIn("celery.hostname", attributes) - def test_set_attributes_from_context_kombu_routing_objects(self): + def test_get_attributes_from_context_kombu_routing_objects(self): context = { "exchange": Exchange("celery_delayed_27", type="topic"), "queue": Queue("celery_queue"), @@ -79,64 +77,37 @@ def test_set_attributes_from_context_kombu_routing_objects(self): "timelimit": ("soft", "hard"), } - span = trace._Span("name", mock.Mock(spec=trace_api.SpanContext)) with self.assertNoLogs("opentelemetry.attributes", level="WARNING"): - utils.set_attributes_from_context(span, context) + attributes = utils.get_attributes_from_context(context) self.assertEqual( - span.attributes.get("celery.exchange"), + attributes.get("celery.exchange"), "celery_delayed_27", ) - self.assertEqual(span.attributes.get("celery.queue"), "celery_queue") - self.assertEqual(span.attributes.get("celery.priority"), 4) - self.assertEqual(span.attributes.get("celery.timelimit"), ("soft", "hard")) + self.assertEqual(attributes.get("celery.queue"), "celery_queue") + self.assertEqual(attributes.get("celery.priority"), 4) + self.assertEqual(attributes.get("celery.timelimit"), ("soft", "hard")) - def test_set_attributes_from_context_string_exchange(self): + def test_get_attributes_from_context_string_exchange(self): context = {"exchange": "celery"} - span = trace._Span("name", mock.Mock(spec=trace_api.SpanContext)) - utils.set_attributes_from_context(span, context) + attributes = utils.get_attributes_from_context(context) - self.assertEqual(span.attributes.get("celery.exchange"), "celery") + self.assertEqual(attributes.get("celery.exchange"), "celery") - def test_set_attributes_from_context_empty_kombu_names(self): + def test_get_attributes_from_context_empty_kombu_names(self): context = { "exchange": Exchange(""), "queue": Queue(""), } - span = trace._Span("name", mock.Mock(spec=trace_api.SpanContext)) with self.assertNoLogs("opentelemetry.attributes", level="WARNING"): - utils.set_attributes_from_context(span, context) - - self.assertNotIn("celery.exchange", span.attributes) - self.assertNotIn("celery.queue", span.attributes) - - def test_set_attributes_not_recording(self): - # it should extract only relevant keys - context = { - "correlation_id": "44b7f305", - "delivery_info": {"eager": True}, - "eta": "soon", - "expires": "later", - "hostname": "localhost", - "id": "44b7f305", - "reply_to": "44b7f305", - "retries": 4, - "timelimit": ("now", "later"), - "custom_meta": "custom_value", - "routing_key": "celery", - } + attributes = utils.get_attributes_from_context(context) - mock_span = mock.Mock() - mock_span.is_recording.return_value = False - utils.set_attributes_from_context(mock_span, context) - self.assertFalse(mock_span.is_recording()) - self.assertTrue(mock_span.is_recording.called) - self.assertFalse(mock_span.set_attribute.called) - self.assertFalse(mock_span.set_status.called) + self.assertNotIn("celery.exchange", attributes) + self.assertNotIn("celery.queue", attributes) - def test_set_attributes_partial_timelimit_hard_limit(self): + def test_get_attributes_partial_timelimit_hard_limit(self): # it should extract only relevant keys context = { "correlation_id": "44b7f305", @@ -151,11 +122,10 @@ def test_set_attributes_partial_timelimit_hard_limit(self): "custom_meta": "custom_value", "routing_key": "celery", } - span = trace._Span("name", mock.Mock(spec=trace_api.SpanContext)) - utils.set_attributes_from_context(span, context) - self.assertEqual(span.attributes.get("celery.timelimit"), ("now", "")) + attributes = utils.get_attributes_from_context(context) + self.assertEqual(attributes.get("celery.timelimit"), ["now", ""]) - def test_set_attributes_partial_timelimit_soft_limit(self): + def test_get_attributes_partial_timelimit_soft_limit(self): # it should extract only relevant keys context = { "correlation_id": "44b7f305", @@ -170,11 +140,10 @@ def test_set_attributes_partial_timelimit_soft_limit(self): "custom_meta": "custom_value", "routing_key": "celery", } - span = trace._Span("name", mock.Mock(spec=trace_api.SpanContext)) - utils.set_attributes_from_context(span, context) - self.assertEqual(span.attributes.get("celery.timelimit"), ("", "later")) + attributes = utils.get_attributes_from_context(context) + self.assertEqual(attributes.get("celery.timelimit"), ["", "later"]) - def test_set_attributes_from_context_empty_keys(self): + def test_get_attributes_from_context_empty_keys(self): # it should not extract empty keys context = { "correlation_id": None, @@ -183,18 +152,17 @@ def test_set_attributes_from_context_empty_keys(self): "retries": 0, } - span = trace._Span("name", mock.Mock(spec=trace_api.SpanContext)) - utils.set_attributes_from_context(span, context) + attributes = utils.get_attributes_from_context(context) - self.assertEqual(len(span.attributes), 0) + self.assertEqual(len(attributes), 0) # edge case: `timelimit` can also be a list of None values context = { "timelimit": [None, None], } - utils.set_attributes_from_context(span, context) + attributes = utils.get_attributes_from_context(context) - self.assertEqual(len(span.attributes), 0) + self.assertEqual(len(attributes), 0) def test_span_propagation(self): # ensure spans getter and setter works properly