Skip to content

Commit b29682b

Browse files
nesb1shalevr
andauthored
aio-pika instrumentation: Removed check for non-sampled span when inject message header. (#1969)
* 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. * black formting --------- Co-authored-by: Shalev Roda <[email protected]>
1 parent 1b1c38d commit b29682b

File tree

5 files changed

+118
-11
lines changed

5 files changed

+118
-11
lines changed

CHANGELOG.md

+2
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
3434

3535
### Fixed
3636

37+
- `opentelemetry-instrumentation-aio-pika` and `opentelemetry-instrumentation-pika` Fix missing trace context propagation when trace not recording.
38+
([#1969](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1969))
3739
- Fix version of Flask dependency `werkzeug`
3840
([#1980](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1980))
3941
- `opentelemetry-resource-detector-azure` Using new Cloud Resource ID attribute.

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

+61
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,36 @@ 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"
109+
) as mock_inject:
110+
decorated_publish = PublishDecorator(
111+
self.tracer, exchange
112+
).decorate(mock_publish)
113+
self.loop.run_until_complete(
114+
decorated_publish(MESSAGE, ROUTING_KEY)
115+
)
116+
mock_publish.assert_called_once()
117+
mock_get_publish_span.assert_called_once()
118+
mock_inject.assert_called_once()
119+
120+
def test_publish_works_with_not_recording_span(self):
121+
self._test_publish_works_with_not_recording_span(Exchange)
122+
123+
def test_publish_works_with_not_recording_span_robust(self):
124+
self._test_publish_works_with_not_recording_span(RobustExchange)
125+
95126

96127
@skipIf(AIOPIKA_VERSION_INFO <= (8, 0), "Only for aio_pika 8")
97128
class TestInstrumentedExchangeAioRmq8(TestCase):
@@ -144,3 +175,33 @@ def test_publish(self):
144175

145176
def test_robust_publish(self):
146177
self._test_publish(RobustExchange)
178+
179+
def _test_publish_works_with_not_recording_span(self, exchange_type):
180+
exchange = exchange_type(CONNECTION_7, CHANNEL_7, EXCHANGE_NAME)
181+
with mock.patch.object(
182+
PublishDecorator, "_get_publish_span"
183+
) as mock_get_publish_span:
184+
mocked_not_recording_span = MagicMock()
185+
mocked_not_recording_span.is_recording.return_value = False
186+
mock_get_publish_span.return_value = mocked_not_recording_span
187+
with mock.patch.object(
188+
Exchange, "publish", return_value=asyncio.sleep(0)
189+
) as mock_publish:
190+
with mock.patch(
191+
"opentelemetry.instrumentation.aio_pika.publish_decorator.propagate.inject"
192+
) as mock_inject:
193+
decorated_publish = PublishDecorator(
194+
self.tracer, exchange
195+
).decorate(mock_publish)
196+
self.loop.run_until_complete(
197+
decorated_publish(MESSAGE, ROUTING_KEY)
198+
)
199+
mock_publish.assert_called_once()
200+
mock_get_publish_span.assert_called_once()
201+
mock_inject.assert_called_once()
202+
203+
def test_publish_works_with_not_recording_span(self):
204+
self._test_publish_works_with_not_recording_span(Exchange)
205+
206+
def test_publish_works_with_not_recording_span_robust(self):
207+
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,55 @@ 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()
394+
inject.assert_called_once_with(properties.headers)
395+
publish_hook.assert_called_once_with(
396+
get_span.return_value, mock_body, properties
397+
)
398+
callback.assert_called_once_with(
399+
exchange_name, routing_key, mock_body, properties, False
400+
)
401+
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+
)
397443
inject.assert_called_once_with(properties.headers)
398444
publish_hook.assert_called_once_with(
399445
get_span.return_value, mock_body, properties

0 commit comments

Comments
 (0)