Add suppress_instrumentation functionality

Author: oxeye-nikolay

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

@@ -4,6 +4,7 @@
 from pika.spec import Basic, BasicProperties
 
 from opentelemetry import propagate, trace
+from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
 from opentelemetry.propagators.textmap import CarrierT, Getter
 from opentelemetry.semconv.trace import (
     MessagingOperationValues,
@@ -76,8 +77,13 @@ def decorated_function(
             task_name="(temporary)",
             operation=None,
         )
+        if not span:
+            return original_function(
+                exchange, routing_key, body, properties, mandatory
+            )
         with trace.use_span(span, end_on_exit=True):
-            propagate.inject(properties.headers)
+            if span.is_recording():
+                propagate.inject(properties.headers)
             retval = original_function(
                 exchange, routing_key, body, properties, mandatory
             )
@@ -92,10 +98,14 @@ def get_span(
     properties: BasicProperties,
     task_name: str,
     operation: Optional[MessagingOperationValues] = None,
-) -> Span:
+) -> Optional[Span]:
     if properties.headers is None:
         properties.headers = {}
     ctx = propagate.extract(properties.headers, getter=pika_getter)
+    if ctx.get_value("suppress_instrumentation") or ctx.get_value(
+        _SUPPRESS_INSTRUMENTATION_KEY
+    ):
+        return None
     task_name = properties.type if properties.type else task_name
     span = tracer.start_span(
         context=ctx, name=generate_span_name(task_name, operation)

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

@@ -28,10 +28,10 @@ def setUp(self) -> None:
         self.channel._impl._consumers = {"mock_key": self.mock_callback}
 
     @mock.patch(
-        "opentelemetry.instrumentation.pika.PikaInstrumentation.instrument_channel"
+        "opentelemetry.instrumentation.pika.PikaInstrumentor.instrument_channel"
     )
     @mock.patch(
-        "opentelemetry.instrumentation.pika.PikaInstrumentation._uninstrument_channel_functions"
+        "opentelemetry.instrumentation.pika.PikaInstrumentor._uninstrument_channel_functions"
     )
     def test_instrument_api(
         self,
@@ -54,10 +54,10 @@ def test_instrument_api(
         )
 
     @mock.patch(
-        "opentelemetry.instrumentation.pika.PikaInstrumentation._instrument_channel_functions"
+        "opentelemetry.instrumentation.pika.PikaInstrumentor._instrument_channel_functions"
     )
     @mock.patch(
-        "opentelemetry.instrumentation.pika.PikaInstrumentation._instrument_consumers"
+        "opentelemetry.instrumentation.pika.PikaInstrumentor._instrument_consumers"
     )
     def test_instrument(
         self,

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

@@ -32,6 +32,9 @@ def test_get_span(
         channel = mock.MagicMock()
         properties = mock.MagicMock()
         task_name = "test.test"
+        context = mock.MagicMock()
+        context.get_value.return_value = None
+        extract.return_value = context
         span = utils.get_span(tracer, channel, properties, task_name)
         extract.assert_called_once()
         generate_span_name.assert_called_once()
@@ -44,6 +47,24 @@ def test_get_span(
             for call in enrich_span.call_args_list
         ), "The returned span was not enriched using enrich_span!"
 
+    @mock.patch("opentelemetry.instrumentation.pika.utils.generate_span_name")
+    @mock.patch("opentelemetry.instrumentation.pika.utils.enrich_span")
+    @mock.patch("opentelemetry.propagate.extract")
+    def test_get_span_suppressed(
+        self,
+        extract: mock.MagicMock,
+        enrich_span: mock.MagicMock,
+        generate_span_name: mock.MagicMock,
+    ) -> None:
+        tracer = mock.MagicMock(spec=Tracer)
+        channel = mock.MagicMock()
+        properties = mock.MagicMock()
+        task_name = "test.test"
+        span = utils.get_span(tracer, channel, properties, task_name)
+        self.assertEqual(span, None)
+        extract.assert_called_once()
+        generate_span_name.assert_not_called()
+
     def test_generate_span_name_no_operation(self) -> None:
         task_name = "test.test"
         operation = None