fix(kafka-python): record actual producer partition from send() future

.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

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) -> Optional[int]:
+        """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: Optional[int],
 ):
     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
 

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
         )