Skip to content

Commit 0327842

Browse files
committed
aio-pika instrumentation: Removed check for non-sampled span when inject message headers. Reason to change is that sampled flag can be propagate https://www.w3.org/TR/trace-context/#sampled-flag and be useful when trace is not sampled.
1 parent e318c94 commit 0327842

File tree

5 files changed

+119
-11
lines changed

5 files changed

+119
-11
lines changed

CHANGELOG.md

+5
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
77

88
## Unreleased
99

10+
### Fixed
11+
12+
- `opentelemetry-instrumentation-aio-pika` and `opentelemetry-instrumentation-pika` Fix missing trace context propagation when trace not recording.
13+
([#1969](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1969))
14+
1015
## Version 1.20.0/0.41b0 (2023-09-01)
1116

1217
### Fixed

instrumentation/opentelemetry-instrumentation-aio-pika/src/opentelemetry/instrumentation/aio_pika/publish_decorator.py

+1-2
Original file line numberDiff line numberDiff line change
@@ -45,8 +45,7 @@ async def decorated_publish(
4545
if not span:
4646
return await publish(message, routing_key, **kwargs)
4747
with trace.use_span(span, end_on_exit=True):
48-
if span.is_recording():
49-
propagate.inject(message.properties.headers)
48+
propagate.inject(message.properties.headers)
5049
return_value = await publish(message, routing_key, **kwargs)
5150
return return_value
5251

instrumentation/opentelemetry-instrumentation-aio-pika/tests/test_publish_decorator.py

+59
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import asyncio
1515
from typing import Type
1616
from unittest import TestCase, mock, skipIf
17+
from unittest.mock import MagicMock
1718

1819
from aio_pika import Exchange, RobustExchange
1920

@@ -92,6 +93,35 @@ def test_publish(self):
9293
def test_robust_publish(self):
9394
self._test_publish(RobustExchange)
9495

96+
def _test_publish_works_with_not_recording_span(self, exchange_type):
97+
exchange = exchange_type(CONNECTION_7, CHANNEL_7, EXCHANGE_NAME)
98+
with mock.patch.object(
99+
PublishDecorator, "_get_publish_span"
100+
) as mock_get_publish_span:
101+
mocked_not_recording_span = MagicMock()
102+
mocked_not_recording_span.is_recording.return_value = False
103+
mock_get_publish_span.return_value = mocked_not_recording_span
104+
with mock.patch.object(
105+
Exchange, "publish", return_value=asyncio.sleep(0)
106+
) as mock_publish:
107+
with mock.patch(
108+
"opentelemetry.instrumentation.aio_pika.publish_decorator.propagate.inject") as mock_inject:
109+
decorated_publish = PublishDecorator(
110+
self.tracer, exchange
111+
).decorate(mock_publish)
112+
self.loop.run_until_complete(
113+
decorated_publish(MESSAGE, ROUTING_KEY)
114+
)
115+
mock_publish.assert_called_once()
116+
mock_get_publish_span.assert_called_once()
117+
mock_inject.assert_called_once()
118+
119+
def test_publish_works_with_not_recording_span(self):
120+
self._test_publish_works_with_not_recording_span(Exchange)
121+
122+
def test_publish_works_with_not_recording_span_robust(self):
123+
self._test_publish_works_with_not_recording_span(RobustExchange)
124+
95125

96126
@skipIf(AIOPIKA_VERSION_INFO <= (8, 0), "Only for aio_pika 8")
97127
class TestInstrumentedExchangeAioRmq8(TestCase):
@@ -144,3 +174,32 @@ def test_publish(self):
144174

145175
def test_robust_publish(self):
146176
self._test_publish(RobustExchange)
177+
178+
def _test_publish_works_with_not_recording_span(self, exchange_type):
179+
exchange = exchange_type(CONNECTION_7, CHANNEL_7, EXCHANGE_NAME)
180+
with mock.patch.object(
181+
PublishDecorator, "_get_publish_span"
182+
) as mock_get_publish_span:
183+
mocked_not_recording_span = MagicMock()
184+
mocked_not_recording_span.is_recording.return_value = False
185+
mock_get_publish_span.return_value = mocked_not_recording_span
186+
with mock.patch.object(
187+
Exchange, "publish", return_value=asyncio.sleep(0)
188+
) as mock_publish:
189+
with mock.patch(
190+
"opentelemetry.instrumentation.aio_pika.publish_decorator.propagate.inject") as mock_inject:
191+
decorated_publish = PublishDecorator(
192+
self.tracer, exchange
193+
).decorate(mock_publish)
194+
self.loop.run_until_complete(
195+
decorated_publish(MESSAGE, ROUTING_KEY)
196+
)
197+
mock_publish.assert_called_once()
198+
mock_get_publish_span.assert_called_once()
199+
mock_inject.assert_called_once()
200+
201+
def test_publish_works_with_not_recording_span(self):
202+
self._test_publish_works_with_not_recording_span(Exchange)
203+
204+
def test_publish_works_with_not_recording_span_robust(self):
205+
self._test_publish_works_with_not_recording_span(RobustExchange)

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

+5-6
Original file line numberDiff line numberDiff line change
@@ -113,12 +113,11 @@ def decorated_function(
113113
exchange, routing_key, body, properties, mandatory
114114
)
115115
with trace.use_span(span, end_on_exit=True):
116-
if span.is_recording():
117-
propagate.inject(properties.headers)
118-
try:
119-
publish_hook(span, body, properties)
120-
except Exception as hook_exception: # pylint: disable=W0703
121-
_LOG.exception(hook_exception)
116+
propagate.inject(properties.headers)
117+
try:
118+
publish_hook(span, body, properties)
119+
except Exception as hook_exception: # pylint: disable=W0703
120+
_LOG.exception(hook_exception)
122121
retval = original_function(
123122
exchange, routing_key, body, properties, mandatory
124123
)

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

+49-3
Original file line numberDiff line numberDiff line change
@@ -292,7 +292,6 @@ def test_decorate_basic_publish(
292292
use_span.assert_called_once_with(
293293
get_span.return_value, end_on_exit=True
294294
)
295-
get_span.return_value.is_recording.assert_called_once()
296295
inject.assert_called_once_with(properties.headers)
297296
callback.assert_called_once_with(
298297
exchange_name, routing_key, mock_body, properties, False
@@ -323,7 +322,6 @@ def test_decorate_basic_publish_no_properties(
323322
use_span.assert_called_once_with(
324323
get_span.return_value, end_on_exit=True
325324
)
326-
get_span.return_value.is_recording.assert_called_once()
327325
inject.assert_called_once_with(basic_properties.return_value.headers)
328326
self.assertEqual(retval, callback.return_value)
329327

@@ -393,7 +391,6 @@ def test_decorate_basic_publish_with_hook(
393391
use_span.assert_called_once_with(
394392
get_span.return_value, end_on_exit=True
395393
)
396-
get_span.return_value.is_recording.assert_called_once()
397394
inject.assert_called_once_with(properties.headers)
398395
publish_hook.assert_called_once_with(
399396
get_span.return_value, mock_body, properties
@@ -402,3 +399,52 @@ def test_decorate_basic_publish_with_hook(
402399
exchange_name, routing_key, mock_body, properties, False
403400
)
404401
self.assertEqual(retval, callback.return_value)
402+
403+
@mock.patch("opentelemetry.instrumentation.pika.utils._get_span")
404+
@mock.patch("opentelemetry.propagate.inject")
405+
@mock.patch("opentelemetry.trace.use_span")
406+
def test_decorate_basic_publish_when_span_is_not_recording(
407+
self,
408+
use_span: mock.MagicMock,
409+
inject: mock.MagicMock,
410+
get_span: mock.MagicMock,
411+
) -> None:
412+
callback = mock.MagicMock()
413+
tracer = mock.MagicMock()
414+
channel = mock.MagicMock(spec=Channel)
415+
exchange_name = "test-exchange"
416+
routing_key = "test-routing-key"
417+
properties = mock.MagicMock()
418+
mock_body = b"mock_body"
419+
publish_hook = mock.MagicMock()
420+
421+
mocked_span = mock.MagicMock()
422+
mocked_span.is_recording.return_value = False
423+
get_span.return_value = mocked_span
424+
425+
decorated_basic_publish = utils._decorate_basic_publish(
426+
callback, channel, tracer, publish_hook
427+
)
428+
retval = decorated_basic_publish(
429+
exchange_name, routing_key, mock_body, properties
430+
)
431+
get_span.assert_called_once_with(
432+
tracer,
433+
channel,
434+
properties,
435+
destination=exchange_name,
436+
span_kind=SpanKind.PRODUCER,
437+
task_name="(temporary)",
438+
operation=None,
439+
)
440+
use_span.assert_called_once_with(
441+
get_span.return_value, end_on_exit=True
442+
)
443+
inject.assert_called_once_with(properties.headers)
444+
publish_hook.assert_called_once_with(
445+
get_span.return_value, mock_body, properties
446+
)
447+
callback.assert_called_once_with(
448+
exchange_name, routing_key, mock_body, properties, False
449+
)
450+
self.assertEqual(retval, callback.return_value)

0 commit comments

Comments
 (0)