fix(kafka): trace using subclass of producer/consumer, not objectproxy [Backport #5545 to 1.12] (#5661)

Backport #5545 to 1.12.

Fixes #5494.

Previously, we patched `confluent-kafka` by wrapping the
`Producer/Consumer` classes with an ObjectProxy
`TracedProducer/TracedConsumer` class. However, since
`confluent-kafka`'s `Producer/Consumer` classes are implemented in C and
have strictly defined slots, our traced classes break for subclasses
such as `AvroProducer/AvroConsumer`.

This PR fixes the issue by basing our traced classes from
`confluent-kafka`'s base `Producer/Consumer` classes directly instead of
using an ObjectProxy. Local testing shows that the bug from #5494 is
resolved.

Before:
```
>>> from ddtrace import patch_all                             
>>> patch_all()                                               
>>> from confluent_kafka.avro import AvroConsumer
>>> AvroConsumer({"group.id": 12345}, "http://localhost:8080")
<stdin>:1: DeprecationWarning: AvroConsumer has been deprecated. Use AvroDeserializer instead.
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/yun.kim/go/src/github.com/DataDog/dd-trace-py/.riot/venv_py3105_mock_pytest_pytest-mock_coverage_pytest-cov_opentracing_hypothesis6451_confluent-kafka/lib/python3.10/site-packages/confluent_kafka/avro/__init__.py", line 161, in __init__
    self._serializer = MessageSerializer(schema_registry, reader_key_schema, reader_value_schema)
AttributeError: 'cimpl.Consumer' object has no attribute '_serializer'
```

After:
```
>>> from ddtrace import patch_all                             
>>> patch_all()                                               
>>> from confluent_kafka.avro import AvroConsumer             
>>> AvroConsumer({"group.id": 12345}, "http://localhost:8080")
<stdin>:1: DeprecationWarning: AvroConsumer has been deprecated. Use AvroDeserializer instead.
<confluent_kafka.avro.AvroConsumer object at 0x10479a680>
```

## Checklist

- [x] Change(s) are motivated and described in the PR description.
- [x] Testing strategy is described if automated tests are not included
in the PR.
- [x] Risk is outlined (performance impact, potential for breakage,
maintainability, etc).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/contributing.html#Release-Note-Guidelines)
are followed.
- [x] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/)).
- [x] PR description includes explicit acknowledgement/acceptance of the
performance implications of this PR as reported in the benchmarks PR
comment.

## Reviewer Checklist

- [x] Title is accurate.
- [x] No unnecessary changes are introduced.
- [x] Description motivates each change.
- [x] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes unless absolutely necessary.
- [x] Testing strategy adequately addresses listed risk(s).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] Release note makes sense to a user of the library.
- [x] Reviewer has explicitly acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment.

Author: Yun-Kim

ddtrace/contrib/kafka/patch.py

@@ -14,7 +14,6 @@
 from ddtrace.internal.utils import ArgumentError
 from ddtrace.internal.utils import get_argument_value
 from ddtrace.pin import Pin
-from ddtrace.vendor.wrapt import ObjectProxy
 
 
 _Producer = confluent_kafka.Producer
@@ -27,43 +26,9 @@
 )
 
 
-class TracedProducer(ObjectProxy):
-    def __init__(self, *args, **kwargs):
-        producer = _Producer(*args, **kwargs)
-        super(TracedProducer, self).__init__(producer)
-        Pin().onto(self)
-
-    def produce(self, *args, **kwargs):
-        func = self.__wrapped__.produce
-        topic = get_argument_value(args, kwargs, 0, "topic")
-        try:
-            value = get_argument_value(args, kwargs, 1, "value")
-        except ArgumentError:
-            value = None
-        message_key = kwargs.get("key", "")
-        partition = kwargs.get("partition", -1)
-
-        pin = Pin.get_from(self)
-        if not pin or not pin.enabled():
-            return func(*args, **kwargs)
-
-        with pin.tracer.trace(
-            kafkax.PRODUCE,
-            service=trace_utils.ext_service(pin, config.kafka),
-            span_type=SpanTypes.WORKER,
-        ) as span:
-            span.set_tag_str(MESSAGING_SYSTEM, kafkax.SERVICE)
-            span.set_tag_str(COMPONENT, config.kafka.integration_name)
-            span.set_tag_str(SPAN_KIND, SpanKind.PRODUCER)
-            span.set_tag_str(kafkax.TOPIC, topic)
-            span.set_tag_str(kafkax.MESSAGE_KEY, ensure_text(message_key))
-            span.set_tag(kafkax.PARTITION, partition)
-            span.set_tag_str(kafkax.TOMBSTONE, str(value is None))
-            span.set_tag(SPAN_MEASURED_KEY)
-            rate = config.kafka.get_analytics_sample_rate()
-            if rate is not None:
-                span.set_tag(ANALYTICS_SAMPLE_RATE_KEY, rate)
-            return func(*args, **kwargs)
+class TracedProducer(confluent_kafka.Producer):
+    def produce(self, topic, value=None, *args, **kwargs):
+        super(TracedProducer, self).produce(topic, value, *args, **kwargs)
 
     # in older versions of confluent_kafka, bool(Producer()) evaluates to False,
     # which makes the Pin functionality ignore it.
@@ -73,60 +38,100 @@ def __bool__(self):
     __nonzero__ = __bool__
 
 
-class TracedConsumer(ObjectProxy):
-
-    __slots__ = "_group_id"
-
-    def __init__(self, *args, **kwargs):
-        consumer = _Consumer(*args, **kwargs)
-        super(TracedConsumer, self).__init__(consumer)
-        self._group_id = get_argument_value(args, kwargs, 0, "config")["group.id"]
-        Pin().onto(self)
-
-    def poll(self, *args, **kwargs):
-        func = self.__wrapped__.poll
-        pin = Pin.get_from(self)
-        if not pin or not pin.enabled():
-            return func(*args, **kwargs)
-
-        with pin.tracer.trace(
-            kafkax.CONSUME,
-            service=trace_utils.ext_service(pin, config.kafka),
-            span_type=SpanTypes.WORKER,
-        ) as span:
-            message = func(*args, **kwargs)
-            span.set_tag_str(MESSAGING_SYSTEM, kafkax.SERVICE)
-            span.set_tag_str(COMPONENT, config.kafka.integration_name)
-            span.set_tag_str(SPAN_KIND, SpanKind.CONSUMER)
-            span.set_tag_str(kafkax.RECEIVED_MESSAGE, str(message is not None))
-            span.set_tag_str(kafkax.GROUP_ID, self._group_id)
-            if message is not None:
-                message_key = message.key() or ""
-                message_offset = message.offset() or -1
-                span.set_tag_str(kafkax.TOPIC, message.topic())
-                span.set_tag_str(kafkax.MESSAGE_KEY, ensure_text(message_key))
-                span.set_tag(kafkax.PARTITION, message.partition())
-                span.set_tag_str(kafkax.TOMBSTONE, str(len(message) == 0))
-                span.set_tag(kafkax.MESSAGE_OFFSET, message_offset)
-            span.set_tag(SPAN_MEASURED_KEY)
-            rate = config.kafka.get_analytics_sample_rate()
-            if rate is not None:
-                span.set_tag(ANALYTICS_SAMPLE_RATE_KEY, rate)
-            return message
+class TracedConsumer(confluent_kafka.Consumer):
+    def __init__(self, config):
+        super(TracedConsumer, self).__init__(config)
+        self._group_id = config["group.id"]
+
+    def poll(self, timeout=1):
+        return super(TracedConsumer, self).poll(timeout)
 
 
 def patch():
     if getattr(confluent_kafka, "_datadog_patch", False):
         return
     setattr(confluent_kafka, "_datadog_patch", True)
 
-    setattr(confluent_kafka, "Producer", TracedProducer)
-    setattr(confluent_kafka, "Consumer", TracedConsumer)
+    confluent_kafka.Producer = TracedProducer
+    confluent_kafka.Consumer = TracedConsumer
+
+    trace_utils.wrap(TracedProducer, "produce", traced_produce)
+    trace_utils.wrap(TracedConsumer, "poll", traced_poll)
+    Pin().onto(confluent_kafka.Producer)
+    Pin().onto(confluent_kafka.Consumer)
 
 
 def unpatch():
     if getattr(confluent_kafka, "_datadog_patch", False):
         setattr(confluent_kafka, "_datadog_patch", False)
 
-    setattr(confluent_kafka, "Producer", _Producer)
-    setattr(confluent_kafka, "Consumer", _Consumer)
+    if trace_utils.iswrapped(TracedProducer.produce):
+        trace_utils.unwrap(TracedProducer, "produce")
+    if trace_utils.iswrapped(TracedConsumer.poll):
+        trace_utils.unwrap(TracedConsumer, "poll")
+
+    confluent_kafka.Producer = _Producer
+    confluent_kafka.Consumer = _Consumer
+
+
+def traced_produce(func, instance, args, kwargs):
+    pin = Pin.get_from(instance)
+    if not pin or not pin.enabled():
+        return func(*args, **kwargs)
+
+    topic = get_argument_value(args, kwargs, 0, "topic")
+    try:
+        value = get_argument_value(args, kwargs, 1, "value")
+    except ArgumentError:
+        value = None
+    message_key = kwargs.get("key", "")
+    partition = kwargs.get("partition", -1)
+
+    with pin.tracer.trace(
+        kafkax.PRODUCE,
+        service=trace_utils.ext_service(pin, config.kafka),
+        span_type=SpanTypes.WORKER,
+    ) as span:
+        span.set_tag_str(MESSAGING_SYSTEM, kafkax.SERVICE)
+        span.set_tag_str(COMPONENT, config.kafka.integration_name)
+        span.set_tag_str(SPAN_KIND, SpanKind.PRODUCER)
+        span.set_tag_str(kafkax.TOPIC, topic)
+        span.set_tag_str(kafkax.MESSAGE_KEY, ensure_text(message_key))
+        span.set_tag(kafkax.PARTITION, partition)
+        span.set_tag_str(kafkax.TOMBSTONE, str(value is None))
+        span.set_tag(SPAN_MEASURED_KEY)
+        rate = config.kafka.get_analytics_sample_rate()
+        if rate is not None:
+            span.set_tag(ANALYTICS_SAMPLE_RATE_KEY, rate)
+        return func(*args, **kwargs)
+
+
+def traced_poll(func, instance, args, kwargs):
+    pin = Pin.get_from(instance)
+    if not pin or not pin.enabled():
+        return func(*args, **kwargs)
+
+    with pin.tracer.trace(
+        kafkax.CONSUME,
+        service=trace_utils.ext_service(pin, config.kafka),
+        span_type=SpanTypes.WORKER,
+    ) as span:
+        message = func(*args, **kwargs)
+        span.set_tag_str(MESSAGING_SYSTEM, kafkax.SERVICE)
+        span.set_tag_str(COMPONENT, config.kafka.integration_name)
+        span.set_tag_str(SPAN_KIND, SpanKind.CONSUMER)
+        span.set_tag_str(kafkax.RECEIVED_MESSAGE, str(message is not None))
+        span.set_tag_str(kafkax.GROUP_ID, instance._group_id)
+        if message is not None:
+            message_key = message.key() or ""
+            message_offset = message.offset() or -1
+            span.set_tag_str(kafkax.TOPIC, message.topic())
+            span.set_tag_str(kafkax.MESSAGE_KEY, ensure_text(message_key))
+            span.set_tag(kafkax.PARTITION, message.partition())
+            span.set_tag_str(kafkax.TOMBSTONE, str(len(message) == 0))
+            span.set_tag(kafkax.MESSAGE_OFFSET, message_offset)
+        span.set_tag(SPAN_MEASURED_KEY)
+        rate = config.kafka.get_analytics_sample_rate()
+        if rate is not None:
+            span.set_tag(ANALYTICS_SAMPLE_RATE_KEY, rate)
+        return message

docs/spelling_wordlist.txt

@@ -136,6 +136,7 @@ opensearch
 opentracer
 opentracing
 otel
+ObjectProxy
 parameterized
 perf
 pid
@@ -182,6 +183,7 @@ sqlite
 stacktrace
 starlette
 stringable
+subclass
 subdomains
 submodule
 submodules

releasenotes/notes/fix-kafka-subclass-b63931a9fd22ddaa.yaml

@@ -0,0 +1,6 @@
+---
+fixes:
+  - |
+    kafka: Previously instantiating a subclass of kafka's Producer/Consumer classes would result in attribute errors
+    due to patching the Producer/Consumer classes with an ObjectProxy. This fix resolves this issue by making the
+    traced classes directly inherit from kafka's base Producer/Consumer classes.

tests/contrib/kafka/test_kafka_patch.py

@@ -9,17 +9,14 @@ class TestKafkaPatch(PatchTestCase.Base):
     __patch_func__ = patch
     __unpatch_func__ = unpatch
 
-    # DEV: normally, we directly patch methods, but since confluent-kafka's methods are implemented in C and
-    # directly imported, we have to patch the Producer/Consumer classes via proxy Traced Producer/Consumer classes.
-    # Because of this, we need to create instances of each proxy class to make wrapping status assertions.
     def assert_module_patched(self, confluent_kafka):
-        self.assert_wrapped(confluent_kafka.Producer({}))
-        self.assert_wrapped(confluent_kafka.Consumer({"group.id": "group_id"}))
+        self.assert_wrapped(confluent_kafka.Producer({}).produce)
+        self.assert_wrapped(confluent_kafka.Consumer({"group.id": "group_id"}).poll)
 
     def assert_not_module_patched(self, confluent_kafka):
-        self.assert_not_wrapped(confluent_kafka.Producer({}))
-        self.assert_not_wrapped(confluent_kafka.Consumer({"group.id": "group_id"}))
+        self.assert_not_wrapped(confluent_kafka.Producer({}).produce)
+        self.assert_not_wrapped(confluent_kafka.Consumer({"group.id": "group_id"}).poll)
 
     def assert_not_module_double_patched(self, confluent_kafka):
-        self.assert_not_double_wrapped(confluent_kafka.Producer({}))
-        self.assert_not_double_wrapped(confluent_kafka.Consumer({"group.id": "group_id"}))
+        self.assert_not_double_wrapped(confluent_kafka.Producer({}).produce)
+        self.assert_not_double_wrapped(confluent_kafka.Consumer({"group.id": "group_id"}).poll)