diff --git a/.changelog/4858.fixed b/.changelog/4858.fixed new file mode 100644 index 0000000000..96a101d933 --- /dev/null +++ b/.changelog/4858.fixed @@ -0,0 +1 @@ +`opentelemetry-instrumentation-kafka-python`: record the actual `messaging.kafka.partition` read back from the producer `send()` future, instead of a pre-computed estimate that did not match the delivered partition for randomly-partitioned messages diff --git a/instrumentation/opentelemetry-instrumentation-kafka-python/src/opentelemetry/instrumentation/kafka/utils.py b/instrumentation/opentelemetry-instrumentation-kafka-python/src/opentelemetry/instrumentation/kafka/utils.py index 00da544325..7bf751faf7 100644 --- a/instrumentation/opentelemetry-instrumentation-kafka-python/src/opentelemetry/instrumentation/kafka/utils.py +++ b/instrumentation/opentelemetry-instrumentation-kafka-python/src/opentelemetry/instrumentation/kafka/utils.py @@ -56,36 +56,25 @@ def extract_send_headers(args, kwargs): ) @staticmethod - def extract_send_partition(instance, args, kwargs): - """extract partition `send` method arguments, using the `_partition` method in KafkaProducer class""" + def extract_send_partition(future) -> int | None: + """Return the partition a message was actually assigned to. + + ``KafkaProducer.send()`` resolves the destination partition internally + (running the partitioner exactly once) and exposes it on the returned + ``FutureRecordMetadata`` via ``_produce_future.topic_partition``. Reading + it back from the future is accurate for every case — explicit partition, + key-based, and the random keyless case. + + The previous implementation estimated the partition *before* ``send()`` + by calling ``instance._partition()`` itself; with the default + partitioner and no key that runs a random choice, which ``send()`` then + redoes, so the recorded value frequently did not match where the message + actually landed. See + https://github.com/open-telemetry/opentelemetry-python-contrib/issues/4625 + """ try: - topic = KafkaPropertiesExtractor.extract_send_topic(args, kwargs) - key = KafkaPropertiesExtractor.extract_send_key(args, kwargs) - value = KafkaPropertiesExtractor.extract_send_value(args, kwargs) - partition = KafkaPropertiesExtractor._extract_argument( - "partition", 4, None, args, kwargs - ) - key_bytes = instance._serialize( - instance.config["key_serializer"], topic, key - ) - value_bytes = instance._serialize( - instance.config["value_serializer"], topic, value - ) - valid_types = (bytes, bytearray, memoryview, type(None)) - if ( - type(key_bytes) not in valid_types - or type(value_bytes) not in valid_types - ): - return None - - instance._wait_on_metadata( - topic, instance.config["max_block_ms"] / 1000.0 - ) - - return instance._partition( - topic, partition, key, value, key_bytes, value_bytes - ) - except Exception as exception: # pylint: disable=W0703 + return future._produce_future.topic_partition[1] + except (AttributeError, IndexError, TypeError) as exception: _LOG.debug("Unable to extract partition: %s", exception) return None @@ -126,12 +115,18 @@ def set(self, carrier: textmap.CarrierT, key: str, value: str) -> None: def _enrich_span( - span, bootstrap_servers: List[str], topic: str, partition: int + span, + bootstrap_servers: list[str], + topic: str, + partition: int | None, ): if span.is_recording(): span.set_attribute(SpanAttributes.MESSAGING_SYSTEM, "kafka") span.set_attribute(SpanAttributes.MESSAGING_DESTINATION, topic) - span.set_attribute(SpanAttributes.MESSAGING_KAFKA_PARTITION, partition) + if partition is not None: + span.set_attribute( + SpanAttributes.MESSAGING_KAFKA_PARTITION, partition + ) span.set_attribute( SpanAttributes.MESSAGING_URL, json.dumps(bootstrap_servers) ) @@ -152,14 +147,10 @@ def _traced_send(func, instance, args, kwargs): bootstrap_servers = KafkaPropertiesExtractor.extract_bootstrap_servers( instance ) - partition = KafkaPropertiesExtractor.extract_send_partition( - instance, args, kwargs - ) span_name = _get_span_name("send", topic) with tracer.start_as_current_span( span_name, kind=trace.SpanKind.PRODUCER ) as span: - _enrich_span(span, bootstrap_servers, topic, partition) propagate.inject( headers, context=trace.set_span_in_context(span), @@ -171,7 +162,10 @@ def _traced_send(func, instance, args, kwargs): except Exception as hook_exception: # pylint: disable=W0703 _LOG.exception(hook_exception) - return func(*args, **kwargs) + future = func(*args, **kwargs) + partition = KafkaPropertiesExtractor.extract_send_partition(future) + _enrich_span(span, bootstrap_servers, topic, partition) + return future return _traced_send diff --git a/instrumentation/opentelemetry-instrumentation-kafka-python/tests/test_utils.py b/instrumentation/opentelemetry-instrumentation-kafka-python/tests/test_utils.py index 3d5935a25f..bb3427d8e5 100644 --- a/instrumentation/opentelemetry-instrumentation-kafka-python/tests/test_utils.py +++ b/instrumentation/opentelemetry-instrumentation-kafka-python/tests/test_utils.py @@ -4,15 +4,19 @@ from unittest import TestCase, mock +from kafka.producer.future import FutureProduceResult, FutureRecordMetadata + from opentelemetry.instrumentation.kafka.utils import ( KafkaPropertiesExtractor, _create_consumer_span, + _enrich_span, _get_span_name, _kafka_getter, _kafka_setter, _wrap_next, _wrap_send, ) +from opentelemetry.semconv.trace import SpanAttributes from opentelemetry.trace import SpanKind @@ -96,8 +100,10 @@ def wrap_send_helper( ) extract_bootstrap_servers.assert_called_once_with(kafka_producer) + # The partition is read back from the future returned by send(), not + # estimated from the call arguments. extract_send_partition.assert_called_once_with( - kafka_producer, self.args, self.kwargs + original_send_callback.return_value ) tracer.start_as_current_span.assert_called_once_with( expected_span_name, kind=SpanKind.PRODUCER @@ -214,27 +220,73 @@ def test_create_consumer_span( ) detach.assert_called_once_with(attach.return_value) - @mock.patch( - "opentelemetry.instrumentation.kafka.utils.KafkaPropertiesExtractor" - ) - def test_kafka_properties_extractor( - self, - kafka_properties_extractor: mock.MagicMock, - ): - kafka_properties_extractor._serialize.return_value = None - kafka_properties_extractor._partition.return_value = "partition" - assert ( - KafkaPropertiesExtractor.extract_send_partition( - kafka_properties_extractor, self.args, self.kwargs - ) - == "partition" - ) - kafka_properties_extractor._wait_on_metadata.side_effect = Exception( - "mocked error" - ) - assert ( - KafkaPropertiesExtractor.extract_send_partition( - kafka_properties_extractor, self.args, self.kwargs - ) - is None + def test_extract_send_partition_reads_actual_partition_from_future(self): + """The partition is read back from the future returned by ``send()``, + which reflects where the message was actually routed. + + Regression test for + https://github.com/open-telemetry/opentelemetry-python-contrib/issues/4625 + """ + future = mock.MagicMock() + future._produce_future.topic_partition = ("test_topic", 7) + + partition = KafkaPropertiesExtractor.extract_send_partition(future) + + self.assertEqual(partition, 7) + + def test_extract_send_partition_zero(self): + """Partition ``0`` is falsy but valid and must not be dropped.""" + future = mock.MagicMock() + future._produce_future.topic_partition = ("test_topic", 0) + + partition = KafkaPropertiesExtractor.extract_send_partition(future) + + self.assertEqual(partition, 0) + + def test_extract_send_partition_matches_real_kafka_future(self): + """Guards against kafka-python changing the internal attribute the + extractor relies on.""" + produce_future = FutureProduceResult(("test_topic", 4)) + future = FutureRecordMetadata( + produce_future, 0, None, None, -1, -1, -1 + ) + + partition = KafkaPropertiesExtractor.extract_send_partition(future) + + self.assertEqual(partition, 4) + + def test_extract_send_partition_returns_none_when_unavailable(self): + """If the future does not expose the expected internals, the attribute + is omitted rather than raising.""" + + class _NoInternals: + pass + + partition = KafkaPropertiesExtractor.extract_send_partition( + _NoInternals() + ) + + self.assertIsNone(partition) + + def test_enrich_span_records_partition_when_present(self): + span = mock.MagicMock() + span.is_recording.return_value = True + + _enrich_span(span, ["localhost:9092"], self.topic_name, 2) + + span.set_attribute.assert_any_call( + SpanAttributes.MESSAGING_KAFKA_PARTITION, 2 + ) + + def test_enrich_span_omits_partition_when_none(self): + span = mock.MagicMock() + span.is_recording.return_value = True + + _enrich_span(span, ["localhost:9092"], self.topic_name, None) + + recorded_attributes = { + call.args[0] for call in span.set_attribute.call_args_list + } + self.assertNotIn( + SpanAttributes.MESSAGING_KAFKA_PARTITION, recorded_attributes )