feat(kafka): add confluent_kafka integration (#5303)

This PR adds automatic instrumentation for
[confluent_kafka](https://github.com/confluentinc/confluent-kafka-python),
which is the Python client for Apache Kafka v0.8 and above.

Currently, this integration provides a wrapper for the
`confluent_kafka.Producer` and `confluent_kafka.Consumer` classes,
specifically providing tracing for `confluent_kafka.Producer.produce()`
and `confluent_kafka.Consumer.poll()`.

This integration supports `confluent_kafka>=1.7.0`.

Notes: 
- consumer group_id as well as consumer/producer client_id are not
publicly available attributes in the `confluent_kafka` API. While the
`message_key, partition, and tombstone` attributes are set in our
instrumentation, the former two attributes are unavailable until the
`confluent_kafka` API changes to publicly include them in the future.
- When a traced producer creates a message with
`Producer.produce(topic)`, but no partition is explicitly set, the
resulting span will have a default partition of `-1` to signify that the
partition will be set by the built-in partitioner. `Consumer poll()`
spans with messages will have the correct partition set as the tag as
that information is available directly from the message.
- The traced `Consumer.poll()` span will have a boolean tag
`kafka.received_message` to indicate if the `poll()` method resulted in
a received message or not. If a message was received, then the span will
include message-specific tags including `topic,partition,tombstone`.
- The traced `Consumer.poll()` (if a message was received) and
`Producer.produce()` spans will include a boolean `tombstone` tag to
indicate whether or not the message was a tombstone (empty/null message
value).


## 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] Author is aware 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 is aware of, and discussed the performance implications
of this PR as reported in the benchmarks PR comment.

---------

Co-authored-by: Yun Kim <[email protected]>
Co-authored-by: Yun Kim <[email protected]>
Co-authored-by: Munir Abdinur <[email protected]>

Author: 3 people

.circleci/config.yml

@@ -1084,6 +1084,15 @@ jobs:
           pattern: "wsgi"
           snapshot: true
 
+  kafka:
+    <<: *machine_executor
+    parallelism: 2
+    steps:
+      - run_test:
+          pattern: 'kafka'
+          snapshot: true
+          docker_services: 'kafka'
+
   kombu:
     <<: *contrib_job_small
     docker:
@@ -1184,6 +1193,7 @@ requires_tests: &requires_tests
     - vendor
     - profile
     - jinja2
+    - kafka
     - kombu
     - mako
     - mariadb
@@ -1288,6 +1298,7 @@ workflows:
       - vendor: *requires_base_venvs
       - profile: *requires_base_venvs
       - jinja2: *requires_base_venvs
+      - kafka: *requires_base_venvs
       - kombu: *requires_base_venvs
       - mako: *requires_base_venvs
       - mariadb: *requires_base_venvs

.riot/requirements/e376a11.txt

@@ -0,0 +1,21 @@
+#
+# This file is autogenerated by pip-compile with Python 3.10
+# by the following command:
+#
+#    pip-compile --no-annotate --resolver=backtracking .riot/requirements/e376a11.in
+#
+attrs==22.2.0
+confluent-kafka==1.9.2
+coverage[toml]==6.5.0
+exceptiongroup==1.1.1
+hypothesis==6.45.0
+iniconfig==2.0.0
+mock==4.0.3
+opentracing==2.4.0
+packaging==23.0
+pluggy==1.0.0
+pytest==7.2.0
+pytest-cov==4.0.0
+pytest-mock==3.10.0
+sortedcontainers==2.4.0
+tomli==2.0.1

ddtrace/_monkey.py

@@ -40,6 +40,7 @@
     "graphql": True,
     "grpc": True,
     "httpx": True,
+    "kafka": True,
     "mongoengine": True,
     "mysql": True,
     "mysqldb": True,
@@ -119,6 +120,7 @@
     "vertica": ("vertica_python",),
     "aws_lambda": ("datadog_lambda",),
     "httplib": ("httplib" if PY2 else "http.client",),
+    "kafka": ("confluent_kafka",),
 }
 
 IAST_PATCH = {

ddtrace/contrib/kafka/__init__.py

@@ -0,0 +1,54 @@
+"""
+This integration instruments the ``confluent-kafka<https://github.com/confluentinc/confluent-kafka-python>``
+library to trace event streaming.
+
+Enabling
+~~~~~~~~
+
+The kafka integration is enabled automatically when using
+:ref:`ddtrace-run <ddtracerun>` or :func:`patch_all() <ddtrace.patch_all>`.
+
+Or use :func:`patch() <ddtrace.patch>` to manually enable the integration::
+
+    from ddtrace import patch
+    patch(kafka=True)
+    import confluent_kafka
+    ...
+
+Global Configuration
+~~~~~~~~~~~~~~~~~~~~
+
+.. py:data:: ddtrace.config.kafka["service"]
+
+   The service name reported by default for your kafka spans.
+
+   This option can also be set with the ``DD_KAFKA_SERVICE`` environment
+   variable.
+
+   Default: ``"kafka"``
+
+
+To configure the kafka integration using the
+``Pin`` API::
+
+    from ddtrace import Pin
+    from ddtrace import patch
+
+    # Make sure to patch before importing confluent_kafka
+    patch(kafka=True)
+
+    import confluent_kafka
+
+    Pin.override(confluent_kafka, service="custom-service-name")
+"""
+from ...internal.utils.importlib import require_modules
+
+
+required_modules = ["confluent_kafka"]
+
+with require_modules(required_modules) as missing_modules:
+    if not missing_modules:
+        from .patch import patch
+        from .patch import unpatch
+
+        __all__ = ["patch", "unpatch"]

ddtrace/contrib/kafka/patch.py

@@ -0,0 +1,132 @@
+import confluent_kafka
+
+from ddtrace import config
+from ddtrace.constants import ANALYTICS_SAMPLE_RATE_KEY
+from ddtrace.constants import SPAN_KIND
+from ddtrace.constants import SPAN_MEASURED_KEY
+from ddtrace.contrib import trace_utils
+from ddtrace.ext import SpanKind
+from ddtrace.ext import SpanTypes
+from ddtrace.ext import kafka as kafkax
+from ddtrace.internal.compat import ensure_text
+from ddtrace.internal.constants import COMPONENT
+from ddtrace.internal.constants import MESSAGING_SYSTEM
+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
+_Consumer = confluent_kafka.Consumer
+
+
+config._add(
+    "kafka",
+    dict(_default_service="kafka"),
+)
+
+
+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)
+
+    # in older versions of confluent_kafka, bool(Producer()) evaluates to False,
+    # which makes the Pin functionality ignore it.
+    def __bool__(self):
+        return True
+
+    __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
+
+
+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)
+
+
+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)

ddtrace/ext/kafka.py

@@ -0,0 +1,12 @@
+SERVICE = "kafka"
+
+TOPIC = "kafka.topic"
+PARTITION = "kafka.partition"
+MESSAGE_KEY = "kafka.message_key"
+MESSAGE_OFFSET = "kafka.message_offset"
+GROUP_ID = "kafka.group_id"
+TOMBSTONE = "kafka.tombstone"
+RECEIVED_MESSAGE = "kafka.received_message"
+
+PRODUCE = "kafka.produce"
+CONSUME = "kafka.consume"

ddtrace/internal/constants.py

@@ -42,3 +42,5 @@
 {"errors": [{"title": "You've been blocked", "detail": "Sorry, you cannot access this page.
 Please contact the customer service team. Security provided by Datadog."}]}
 """
+
+MESSAGING_SYSTEM = "messaging.system"

docker-compose.yml

@@ -56,6 +56,26 @@ services:
         image: redis:4.0-alpine
         ports:
             - "127.0.0.1:6379:6379"
+    zookeeper:
+      image: confluentinc/cp-zookeeper:7.3.1
+      environment:
+        - ZOOKEEPER_CLIENT_PORT=2181
+        - ZOOKEEPER_TICK_TIME=2000
+      ports:
+        - "127.0.0.1:22181:2181"
+    kafka:
+      image: confluentinc/cp-kafka:7.3.2
+      ports:
+        - "127.0.0.1:29092:29092"
+      environment:
+        - KAFKA_BROKER_ID=1
+        - KAFKA_ZOOKEEPER_CONNECT=zookeeper:2181
+        - KAFKA_ADVERTISED_LISTENERS=PLAINTEXT://kafka:9092,PLAINTEXT_HOST://localhost:29092
+        - KAFKA_LISTENER_SECURITY_PROTOCOL_MAP=PLAINTEXT:PLAINTEXT,PLAINTEXT_HOST:PLAINTEXT
+        - KAFKA_INTER_BROKER_LISTENER_NAME=PLAINTEXT
+        - KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR=1
+      depends_on:
+        - zookeeper
     rediscluster:
         image: grokzen/redis-cluster:6.2.0
         environment:

docs/index.rst

@@ -75,6 +75,8 @@ contacting support.
 +--------------------------------------------------+---------------+----------------+
 | :ref:`cherrypy`                                  | >= 11.2.0     | No             |
 +--------------------------------------------------+---------------+----------------+
+| :ref:`confluent-kafka <kafka>`                   | >= 1.7.0      | Yes            |
++--------------------------------------------------+---------------+----------------+
 | :ref:`consul`                                    | >= 0.7        | Yes [2]_       |
 +--------------------------------------------------+---------------+----------------+
 | :ref:`django`                                    | >= 1.8        | Yes            |

docs/integrations.rst

@@ -236,6 +236,12 @@ Jinja2
 ^^^^^^
 .. automodule:: ddtrace.contrib.jinja2
 
+.. _kafka:
+
+Kafka
+^^^^^
+
+.. automodule:: ddtrace.contrib.kafka
 
 .. _kombu: