pika: add visibility to sampler to messagging spans attributes

.changelog/5313.changed

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

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

@@ -27,6 +27,7 @@
 )
 from opentelemetry.trace import SpanKind, Tracer
 from opentelemetry.trace.span import Span
+from opentelemetry.util.types import AttributeValue
 
 _LOG = getLogger(__name__)
 
@@ -146,13 +147,12 @@ def _get_span(
     if not is_instrumentation_enabled():
         return None
     task_name = properties.type if properties.type else task_name
-    span = tracer.start_span(
+    attributes = _get_attributes(channel, properties, destination, operation)
+    return tracer.start_span(
         name=_generate_span_name(destination, operation),
         kind=span_kind,
+        attributes=attributes,
     )
-    if span.is_recording():
-        _enrich_span(span, channel, properties, destination, operation)
-    return span
 
 
 def _generate_span_name(task_name: str, operation: MessagingOperationValues | None) -> str:
@@ -161,34 +161,33 @@ def _generate_span_name(task_name: str, operation: MessagingOperationValues | No
     return f"{task_name} {operation.value}"
 
 
-def _enrich_span(
-    span: Span,
+def _get_attributes(
     channel: Channel | None,
     properties: BasicProperties,
     task_destination: str,
     operation: MessagingOperationValues | None = None,
-) -> None:
-    span.set_attribute(messaging_attributes.MESSAGING_SYSTEM, "rabbitmq")
+) -> dict[str, AttributeValue]:
+    attributes: dict[str, AttributeValue] = {
+        messaging_attributes.MESSAGING_SYSTEM: messaging_attributes.MessagingSystemValues.RABBITMQ.value,
+        SpanAttributes.MESSAGING_DESTINATION: task_destination,
+    }
     if operation:
-        span.set_attribute(SpanAttributes.MESSAGING_OPERATION, operation.value)
+        attributes[SpanAttributes.MESSAGING_OPERATION] = operation.value
     else:
-        span.set_attribute(SpanAttributes.MESSAGING_TEMP_DESTINATION, True)
-    span.set_attribute(SpanAttributes.MESSAGING_DESTINATION, task_destination)
+        attributes[SpanAttributes.MESSAGING_TEMP_DESTINATION] = True
     if properties.message_id:
-        span.set_attribute(
-            messaging_attributes.MESSAGING_MESSAGE_ID,
-            properties.message_id,
-        )
+        attributes[messaging_attributes.MESSAGING_MESSAGE_ID] = properties.message_id
     if properties.correlation_id:
-        span.set_attribute(SpanAttributes.MESSAGING_CONVERSATION_ID, properties.correlation_id)
+        attributes[SpanAttributes.MESSAGING_CONVERSATION_ID] = properties.correlation_id
     if not channel:
-        return
+        return attributes
     if not hasattr(channel.connection, "params"):
-        span.set_attribute(NET_PEER_NAME, channel.connection._impl.params.host)
-        span.set_attribute(NET_PEER_PORT, channel.connection._impl.params.port)
+        attributes[NET_PEER_NAME] = channel.connection._impl.params.host
+        attributes[NET_PEER_PORT] = channel.connection._impl.params.port
     else:
-        span.set_attribute(NET_PEER_NAME, channel.connection.params.host)
-        span.set_attribute(NET_PEER_PORT, channel.connection.params.port)
+        attributes[NET_PEER_NAME] = channel.connection.params.host
+        attributes[NET_PEER_PORT] = channel.connection.params.port
+    return attributes
 
 
 # pylint:disable=abstract-method

instrumentation/opentelemetry-instrumentation-pika/tests/test_sampling.py

@@ -0,0 +1,112 @@
+# Copyright The OpenTelemetry Authors
+# SPDX-License-Identifier: Apache-2.0
+
+from collections import deque
+from types import SimpleNamespace
+from typing import Any
+from unittest import mock
+
+from pika.adapters.blocking_connection import _ConsumerDeliveryEvt
+from pika.spec import Basic, BasicProperties
+
+from opentelemetry import context
+from opentelemetry.instrumentation.pika import utils
+from opentelemetry.semconv._incubating.attributes import messaging_attributes
+from opentelemetry.semconv._incubating.attributes.net_attributes import (
+    NET_PEER_NAME,
+    NET_PEER_PORT,
+)
+from opentelemetry.semconv.trace import MessagingOperationValues, SpanAttributes
+from opentelemetry.test.samplertestutil import CapturingSampler
+from opentelemetry.test.test_base import TestBase
+from opentelemetry.trace import SpanKind, StatusCode
+
+
+class TestPikaSampling(TestBase):
+    def _configure_sampler(self, sampler: CapturingSampler) -> None:
+        self.tracer_provider, self.memory_exporter = self.create_tracer_provider(sampler=sampler)
+        self.addCleanup(self.tracer_provider.shutdown)
+        self.tracer = self.tracer_provider.get_tracer(__name__)
+        self.channel = SimpleNamespace(connection=SimpleNamespace(params=SimpleNamespace(host="localhost", port=5672)))
+
+    def _invoke(self, operation: str, callback: mock.Mock, properties: BasicProperties) -> Any:
+        if operation == "publish":
+            decorated = utils._decorate_basic_publish(callback, self.channel, self.tracer)
+            return decorated("exchange", "queue", b"body", properties)
+        method = Basic.Deliver(exchange="exchange", routing_key="queue")
+        if operation == "callback":
+            decorated = utils._decorate_callback(callback, self.tracer, "task")
+            return decorated(self.channel, method, properties, b"body")
+        event = _ConsumerDeliveryEvt(method, properties, b"body")
+        generator_info = SimpleNamespace(consumer_tag="consumer")
+        proxy = utils.ReadyMessagesDequeProxy(deque([event]), generator_info, self.tracer, callback)
+        try:
+            return proxy.popleft()
+        finally:
+            if proxy._self_active_token:
+                context.detach(proxy._self_active_token)
+
+    def test_attributes_passed_to_sampler(self) -> None:
+        for operation in ("publish", "callback", "generator"):
+            with self.subTest(operation=operation):
+                sampler = CapturingSampler()
+                self._configure_sampler(sampler)
+                properties = BasicProperties(headers={}, message_id="message-1", correlation_id="conversation-1")
+                callback = mock.Mock()
+
+                self._invoke(operation, callback, properties)
+
+                expected = {
+                    messaging_attributes.MESSAGING_SYSTEM: messaging_attributes.MessagingSystemValues.RABBITMQ.value,
+                    SpanAttributes.MESSAGING_DESTINATION: "exchange",
+                    messaging_attributes.MESSAGING_MESSAGE_ID: "message-1",
+                    SpanAttributes.MESSAGING_CONVERSATION_ID: "conversation-1",
+                }
+                if operation == "publish":
+                    expected[SpanAttributes.MESSAGING_TEMP_DESTINATION] = True
+                else:
+                    expected[SpanAttributes.MESSAGING_OPERATION] = MessagingOperationValues.RECEIVE.value
+                if operation != "generator":
+                    expected[NET_PEER_NAME] = "localhost"
+                    expected[NET_PEER_PORT] = 5672
+                self.assertEqual(sampler.kind, SpanKind.PRODUCER if operation == "publish" else 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)
+                self.assertEqual(spans[0].attributes, expected)
+                callback.assert_called_once()
+
+    def test_attributes_control_sampling(self) -> None:
+        for operation in ("publish", "callback", "generator"):
+            for message_id, expected_count in (("sample", 1), ("drop", 0)):
+                with self.subTest(operation=operation, message_id=message_id):
+                    sampler = CapturingSampler(
+                        required_attributes={messaging_attributes.MESSAGING_MESSAGE_ID: "sample"}
+                    )
+                    self._configure_sampler(sampler)
+                    properties = BasicProperties(headers={}, message_id=message_id)
+                    callback = mock.Mock()
+
+                    self._invoke(operation, callback, properties)
+
+                    callback.assert_called_once()
+                    self.assertEqual(len(self.memory_exporter.get_finished_spans()), expected_count)
+
+    def test_errors_propagate(self) -> None:
+        for operation in ("publish", "callback"):
+            with self.subTest(operation=operation):
+                self._configure_sampler(CapturingSampler())
+                error = RuntimeError("test error")
+                callback = mock.Mock(side_effect=error)
+                properties = BasicProperties(headers={}, message_id="message-1")
+
+                with self.assertRaises(RuntimeError) as raised:
+                    self._invoke(operation, callback, properties)
+
+                self.assertIs(raised.exception, error)
+                spans = self.memory_exporter.get_finished_spans()
+                self.assertEqual(len(spans), 1)
+                self.assertEqual(spans[0].status.status_code, StatusCode.ERROR)
+                self.assertEqual(spans[0].attributes[messaging_attributes.MESSAGING_MESSAGE_ID], "message-1")

instrumentation/opentelemetry-instrumentation-pika/tests/test_utils.py

@@ -21,16 +21,16 @@
     MessagingOperationValues,
     SpanAttributes,
 )
-from opentelemetry.trace import Span, SpanKind, Tracer
+from opentelemetry.trace import SpanKind, Tracer
 
 
 class TestUtils(TestCase):
     @staticmethod
     @mock.patch("opentelemetry.context.get_value")
     @mock.patch("opentelemetry.instrumentation.pika.utils._generate_span_name")
-    @mock.patch("opentelemetry.instrumentation.pika.utils._enrich_span")
+    @mock.patch("opentelemetry.instrumentation.pika.utils._get_attributes")
     def test_get_span(
-        enrich_span: mock.MagicMock,
+        get_attributes: mock.MagicMock,
         generate_span_name: mock.MagicMock,
         get_value: mock.MagicMock,
     ) -> None:
@@ -43,9 +43,10 @@ def test_get_span(
         get_value.return_value = None
         _ = utils._get_span(tracer, channel, properties, task_name, destination, span_kind)
         generate_span_name.assert_called_once()
-        tracer.start_span.assert_called_once_with(name=generate_span_name.return_value, kind=span_kind)
-        enrich_span.assert_called_once_with(
-            tracer.start_span.return_value,
+        tracer.start_span.assert_called_once_with(
+            name=generate_span_name.return_value, kind=span_kind, attributes=get_attributes.return_value
+        )
+        get_attributes.assert_called_once_with(
             channel,
             properties,
             destination,
@@ -54,10 +55,10 @@ def test_get_span(
 
     @mock.patch("opentelemetry.context.get_value")
     @mock.patch("opentelemetry.instrumentation.pika.utils._generate_span_name")
-    @mock.patch("opentelemetry.instrumentation.pika.utils._enrich_span")
+    @mock.patch("opentelemetry.instrumentation.pika.utils._get_attributes")
     def test_get_span_suppressed(
         self,
-        enrich_span: mock.MagicMock,
+        get_attributes: mock.MagicMock,
         generate_span_name: mock.MagicMock,
         get_value: mock.MagicMock,
     ) -> None:
@@ -71,7 +72,7 @@ def test_get_span_suppressed(
         span = utils._get_span(tracer, channel, properties, task_name, span_kind, ctx)
         self.assertEqual(span, None)
         generate_span_name.assert_not_called()
-        enrich_span.assert_not_called()
+        get_attributes.assert_not_called()
 
     def test_generate_span_name_no_operation(self) -> None:
         task_name = "test.test"
@@ -86,85 +87,47 @@ def test_generate_span_name_with_operation(self) -> None:
         span_name = utils._generate_span_name(task_name, operation)
         self.assertEqual(span_name, f"{task_name} {operation.value}")
 
-    @staticmethod
-    def test_enrich_span_basic_values() -> None:
+    def test_get_attributes_basic_values(self) -> None:
         channel = mock.MagicMock()
         properties = mock.MagicMock()
         task_destination = "test.test"
-        span = mock.MagicMock(spec=Span)
-        utils._enrich_span(span, channel, properties, task_destination)
-        span.set_attribute.assert_has_calls(
-            any_order=True,
-            calls=[
-                mock.call(messaging_attributes.MESSAGING_SYSTEM, "rabbitmq"),
-                mock.call(SpanAttributes.MESSAGING_TEMP_DESTINATION, True),
-                mock.call(SpanAttributes.MESSAGING_DESTINATION, task_destination),
-                mock.call(
-                    messaging_attributes.MESSAGING_MESSAGE_ID,
-                    properties.message_id,
-                ),
-                mock.call(
-                    SpanAttributes.MESSAGING_CONVERSATION_ID,
-                    properties.correlation_id,
-                ),
-                mock.call(
-                    NET_PEER_NAME,
-                    channel.connection.params.host,
-                ),
-                mock.call(
-                    NET_PEER_PORT,
-                    channel.connection.params.port,
-                ),
-            ],
+        attributes = utils._get_attributes(channel, properties, task_destination)
+        self.assertEqual(
+            attributes,
+            {
+                messaging_attributes.MESSAGING_SYSTEM: messaging_attributes.MessagingSystemValues.RABBITMQ.value,
+                SpanAttributes.MESSAGING_TEMP_DESTINATION: True,
+                SpanAttributes.MESSAGING_DESTINATION: task_destination,
+                messaging_attributes.MESSAGING_MESSAGE_ID: properties.message_id,
+                SpanAttributes.MESSAGING_CONVERSATION_ID: properties.correlation_id,
+                NET_PEER_NAME: channel.connection.params.host,
+                NET_PEER_PORT: channel.connection.params.port,
+            },
         )
 
-    @staticmethod
-    def test_enrich_span_with_operation() -> None:
-        channel = mock.MagicMock()
-        properties = mock.MagicMock()
-        task_destination = "test.test"
-        operation = mock.MagicMock()
-        span = mock.MagicMock(spec=Span)
-        utils._enrich_span(span, channel, properties, task_destination, operation)
-        span.set_attribute.assert_has_calls(
-            any_order=True,
-            calls=[mock.call(SpanAttributes.MESSAGING_OPERATION, operation.value)],
+    def test_get_attributes_with_operation(self) -> None:
+        attributes = utils._get_attributes(None, BasicProperties(), "test.test", MessagingOperationValues.RECEIVE)
+        self.assertEqual(attributes[SpanAttributes.MESSAGING_OPERATION], MessagingOperationValues.RECEIVE.value)
+        self.assertNotIn(SpanAttributes.MESSAGING_TEMP_DESTINATION, attributes)
+
+    def test_get_attributes_without_optional_properties(self) -> None:
+        attributes = utils._get_attributes(None, BasicProperties(), "test.test")
+        self.assertEqual(
+            attributes,
+            {
+                messaging_attributes.MESSAGING_SYSTEM: messaging_attributes.MessagingSystemValues.RABBITMQ.value,
+                SpanAttributes.MESSAGING_TEMP_DESTINATION: True,
+                SpanAttributes.MESSAGING_DESTINATION: "test.test",
+            },
         )
 
-    @staticmethod
-    def test_enrich_span_without_operation() -> None:
+    def test_get_attributes_unique_connection(self) -> None:
         channel = mock.MagicMock()
-        properties = mock.MagicMock()
-        task_destination = "test.test"
-        span = mock.MagicMock(spec=Span)
-        utils._enrich_span(span, channel, properties, task_destination)
-        span.set_attribute.assert_has_calls(
-            any_order=True,
-            calls=[mock.call(SpanAttributes.MESSAGING_TEMP_DESTINATION, True)],
-        )
-
-    @staticmethod
-    def test_enrich_span_unique_connection() -> None:
-        channel = mock.MagicMock()
-        properties = mock.MagicMock()
-        task_destination = "test.test"
-        span = mock.MagicMock(spec=Span)
-        # We do this to create the behaviour of hasattr(channel.connection, "params") == False
+        # Blocking connections expose parameters through their implementation.
         del channel.connection.params
-        utils._enrich_span(span, channel, properties, task_destination)
-        span.set_attribute.assert_has_calls(
-            any_order=True,
-            calls=[
-                mock.call(
-                    NET_PEER_NAME,
-                    channel.connection._impl.params.host,
-                ),
-                mock.call(
-                    NET_PEER_PORT,
-                    channel.connection._impl.params.port,
-                ),
-            ],
-        )
+        attributes = utils._get_attributes(channel, BasicProperties(), "test.test")
+        self.assertEqual(attributes[NET_PEER_NAME], channel.connection._impl.params.host)
+        self.assertEqual(attributes[NET_PEER_PORT], channel.connection._impl.params.port)
 
     @mock.patch("opentelemetry.instrumentation.pika.utils._get_span")
     @mock.patch("opentelemetry.propagate.extract")