Kafka support

.ci/.jenkins_framework.yml

@@ -53,3 +53,4 @@ FRAMEWORK:
   - sanic-newest
   - aiomysql-newest
   - aiobotocore-newest
+  - kafka-python-newest

.ci/.jenkins_framework_full.yml

@@ -85,3 +85,4 @@ FRAMEWORK:
   - sanic-20.12
   - sanic-newest
   - aiobotocore-newest
+  - kafka-python-newest

docs/supported-technologies.asciidoc

@@ -466,6 +466,23 @@ Collected trace data:
 
 * Destination (address and port)
 
+[float]
+[[automatic-instrumentation-db-kafka-python]]
+==== kafka-python
+
+Library: `kafka-python` (`>=2.0`)
+
+Instrumented methods:
+
+ * `kafka.KafkaProducer.send`,
+ * `kafka.KafkaConsumer.poll`,
+ * `kafka.KafkaConsumer.\\__next__`
+
+Collected trace data:
+
+ * Destination (address and port)
+ * topic (if applicable)
+
 
 [float]
 [[automatic-instrumentation-http]]

elasticapm/base.py

@@ -632,12 +632,21 @@ def load_processors(self):
         return [seen.setdefault(path, import_string(path)) for path in processors if path not in seen]
 
     def should_ignore_url(self, url):
+        """Checks if URL should be ignored based on the transaction_ignore_urls setting"""
         if self.config.transaction_ignore_urls:
             for pattern in self.config.transaction_ignore_urls:
                 if pattern.match(url):
                     return True
         return False
 
+    def should_ignore_topic(self, topic: str) -> bool:
+        """Checks if messaging topic should be ignored based on the ignore_message_queues setting"""
+        if self.config.ignore_message_queues:
+            for pattern in self.config.ignore_message_queues:
+                if pattern.match(topic):
+                    return True
+        return False
+
     def check_python_version(self):
         v = tuple(map(int, platform.python_version_tuple()[:2]))
         if v == (2, 7):

elasticapm/conf/__init__.py

@@ -642,6 +642,7 @@ class Config(_ConfigBase):
     autoinsert_django_middleware = _BoolConfigValue("AUTOINSERT_DJANGO_MIDDLEWARE", default=True)
     transactions_ignore_patterns = _ListConfigValue("TRANSACTIONS_IGNORE_PATTERNS", default=[])
     transaction_ignore_urls = _ListConfigValue("TRANSACTION_IGNORE_URLS", type=starmatch_to_regex, default=[])
+    ignore_message_queues = _ListConfigValue("IGNORE_MESSAGE_QUEUES", type=starmatch_to_regex, default=[])
     service_version = _ConfigValue("SERVICE_VERSION")
     framework_name = _ConfigValue("FRAMEWORK_NAME")
     framework_version = _ConfigValue("FRAMEWORK_VERSION")

elasticapm/conf/constants.py

@@ -64,6 +64,7 @@ def _starmatch_to_regex(pattern):
 TRACE_CONTEXT_VERSION = 0
 TRACEPARENT_HEADER_NAME = "traceparent"
 TRACEPARENT_LEGACY_HEADER_NAME = "elastic-apm-traceparent"
+TRACEPARENT_BINARY_HEADER_NAME = "elasticapmtraceparent"
 TRACESTATE_HEADER_NAME = "tracestate"
 
 TIMESTAMP_FORMAT = "%Y-%m-%dT%H:%M:%S.%fZ"

elasticapm/instrumentation/packages/kafka.py

@@ -0,0 +1,189 @@
+#  BSD 3-Clause License
+#
+#  Copyright (c) 2019, Elasticsearch BV
+#  All rights reserved.
+#
+#  Redistribution and use in source and binary forms, with or without
+#  modification, are permitted provided that the following conditions are met:
+#
+#  * Redistributions of source code must retain the above copyright notice, this
+#    list of conditions and the following disclaimer.
+#
+#  * Redistributions in binary form must reproduce the above copyright notice,
+#    this list of conditions and the following disclaimer in the documentation
+#    and/or other materials provided with the distribution.
+#
+#  * Neither the name of the copyright holder nor the names of its
+#    contributors may be used to endorse or promote products derived from
+#    this software without specific prior written permission.
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+#  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+#  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+#  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+#  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+#  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+#  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+#  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import time
+from typing import Optional
+
+import elasticapm
+from elasticapm import get_client
+from elasticapm.conf import constants
+from elasticapm.instrumentation.packages.base import AbstractInstrumentedModule
+from elasticapm.traces import DroppedSpan, capture_span, execution_context
+from elasticapm.utils.disttracing import TraceParent
+
+
+class KafkaInstrumentation(AbstractInstrumentedModule):
+
+    instrument_list = [
+        ("kafka", "KafkaProducer.send"),
+        ("kafka", "KafkaConsumer.poll"),
+        ("kafka", "KafkaConsumer.__next__"),
+    ]
+    provider_name = "kafka"
+    name = "kafka"
+
+    def _trace_send(self, instance, wrapped, *args, destination_info=None, **kwargs):
+        topic = args[0] if args else kwargs["topic"]
+        headers = args[4] if len(args) > 4 else kwargs.get("headers", None)
+
+        span_name = f"Kafka SEND to {topic}"
+        destination_info["service"]["resource"] += topic
+        with capture_span(
+            name=span_name,
+            span_type="messaging",
+            span_subtype=self.provider_name,
+            span_action="send",
+            leaf=True,
+            extra={
+                "message": {"queue": {"name": topic}},
+                "destination": destination_info,
+            },
+        ) as span:
+            transaction = execution_context.get_transaction()
+            if transaction:
+                tp = transaction.trace_parent.copy_from(span_id=span.id)
+            if headers:
+                headers.append((constants.TRACEPARENT_BINARY_HEADER_NAME, tp.to_binary()))
+            else:
+                headers = [(constants.TRACEPARENT_BINARY_HEADER_NAME, tp.to_binary())]
+                if len(args) > 4:
+                    args = list(args)
+                    args[4] = headers
+                else:
+                    kwargs["headers"] = headers
+            result = wrapped(*args, **kwargs)
+            if instance and instance._metadata.controller and not isinstance(span, DroppedSpan):
+                address = instance._metadata.controller[1]
+                port = instance._metadata.controller[2]
+                span.context["destination"]["address"] = address
+                span.context["destination"]["port"] = port
+            return result
+
+    def call_if_sampling(self, module, method, wrapped, instance, args, kwargs):
+        # Contrasting to the superclass implementation, we *always* want to
+        # return a proxied connection, even if there is no ongoing elasticapm
+        # transaction yet. This ensures that we instrument the cursor once
+        # the transaction started.
+        return self.call(module, method, wrapped, instance, args, kwargs)
+
+    def call(self, module, method, wrapped, instance, args, kwargs):
+        client = get_client()
+        destination_info = {
+            "service": {"name": "kafka", "resource": "kafka/", "type": "messaging"},
+        }
+
+        if method == "KafkaProducer.send":
+            topic = args[0] if args else kwargs["topic"]
+            if client.should_ignore_topic(topic) or not execution_context.get_transaction():
+                return wrapped(*args, **kwargs)
+            return self._trace_send(instance, wrapped, destination_info=destination_info, *args, **kwargs)
+
+        elif method == "KafkaConsumer.poll":
+            transaction = execution_context.get_transaction()
+            if transaction:
+                with capture_span(
+                    name="Kafka POLL",
+                    span_type="messaging",
+                    span_subtype=self.provider_name,
+                    span_action="poll",
+                    leaf=True,
+                    extra={
+                        "destination": destination_info,
+                    },
+                ) as span:
+                    if not isinstance(span, DroppedSpan) and instance._subscription.subscription:
+                        span.name += " from " + ", ".join(sorted(instance._subscription.subscription))
+                    results = wrapped(*args, **kwargs)
+                    return results
+            else:
+                return wrapped(*args, **kwargs)
+
+        elif method == "KafkaConsumer.__next__":
+            transaction = execution_context.get_transaction()
+            if transaction and transaction.transaction_type != "messaging":
+                # somebody started a transaction outside of the consumer,
+                # so we capture it as a span, and record the causal trace as a link
+                with capture_span(
+                    name="consumer",
+                    span_type="messaging",
+                    span_subtype=self.provider_name,
+                    span_action="receive",
+                    leaf=True,
+                    extra={
+                        "message": {"queue": {"name": ""}},
+                        "destination": destination_info,
+                    },
+                ) as span:
+                    try:
+                        result = wrapped(*args, **kwargs)
+                    except StopIteration:
+                        span.cancel()
+                        raise
+                    if not isinstance(span, DroppedSpan):
+                        topic = result[0]
+                        if client.should_ignore_topic(topic):
+                            span.cancel()
+                            return result
+                        trace_parent = self.get_traceparent_from_result(result)
+                        if trace_parent:
+                            span.add_link(trace_parent)
+                        destination_info["service"]["resource"] += topic
+                        span.context["message"]["queue"]["name"] = topic
+                        span.name = "Kafka RECEIVE from " + topic
+                    return result
+            else:
+                # No transaction running, or this is a transaction started by us,
+                # so let's end it and start the next,
+                # unless a StopIteration is raised, at which point we do nothing.
+                if transaction:
+                    client.end_transaction()
+                result = wrapped(*args, **kwargs)
+                topic = result[0]
+                if client.should_ignore_topic(topic):
+                    return result
+                trace_parent = self.get_traceparent_from_result(result)
+                transaction = client.begin_transaction("messaging", trace_parent=trace_parent)
+                if result.timestamp_type == 0:
+                    current_time_millis = int(round(time.time() * 1000))
+                    age = current_time_millis - result.timestamp
+                    transaction.context = {
+                        "message": {"age": {"ms": age}, "queue": {"name": topic}},
+                        "service": {"framework": {"name": "Kafka"}},
+                    }
+                transaction_name = "Kafka RECEIVE from " + topic
+                elasticapm.set_transaction_name(transaction_name, override=True)
+                res = constants.OUTCOME.SUCCESS
+                elasticapm.set_transaction_result(res, override=False)
+                return result
+
+    def get_traceparent_from_result(self, result) -> Optional[TraceParent]:
+        for k, v in result.headers:
+            if k == constants.TRACEPARENT_BINARY_HEADER_NAME:
+                return TraceParent.from_binary(v)

elasticapm/instrumentation/register.py

@@ -66,6 +66,7 @@
     "elasticapm.instrumentation.packages.httpx.sync.httpcore.HTTPCoreInstrumentation",
     "elasticapm.instrumentation.packages.httplib2.Httplib2Instrumentation",
     "elasticapm.instrumentation.packages.azure.AzureInstrumentation",
+    "elasticapm.instrumentation.packages.kafka.KafkaInstrumentation",
 }
 
 if sys.version_info >= (3, 7):

elasticapm/traces.py

@@ -38,7 +38,7 @@
 from collections import defaultdict
 from datetime import timedelta
 from types import TracebackType
-from typing import Any, Callable, Dict, Optional, Tuple, Type, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
 
 import elasticapm
 from elasticapm.conf import constants
@@ -100,6 +100,7 @@ def __init__(self, labels=None, start=None):
         self.start_time: float = time_to_perf_counter(start) if start is not None else _time_func()
         self.ended_time: Optional[float] = None
         self.duration: Optional[timedelta] = None
+        self.links: List[Dict[str, str]] = []
         if labels:
             self.label(**labels)
 
@@ -144,6 +145,12 @@ def label(self, **labels):
         labels = encoding.enforce_label_format(labels)
         self.labels.update(labels)
 
+    def add_link(self, trace_parent: TraceParent) -> None:
+        """
+        Causally link this span/transaction to another span/transaction
+        """
+        self.links.append({"trace_id": trace_parent.trace_id, "span_id": trace_parent.span_id})
+
     def set_success(self):
         self.outcome = constants.OUTCOME.SUCCESS
 
@@ -394,6 +401,8 @@ def to_dict(self) -> dict:
             # only set parent_id if this transaction isn't the root
             if self.trace_parent.span_id and self.trace_parent.span_id != self.id:
                 result["parent_id"] = self.trace_parent.span_id
+        if self.links:
+            result["links"] = self.links
         # faas context belongs top-level on the transaction
         if "faas" in self.context:
             result["faas"] = self.context.pop("faas")
@@ -477,6 +486,7 @@ class Span(BaseSpan):
         "sync",
         "outcome",
         "_child_durations",
+        "_cancelled",
     )
 
     def __init__(
@@ -527,6 +537,7 @@ def __init__(
         self.action = span_action
         self.dist_tracing_propagated = False
         self.composite: Dict[str, Any] = {}
+        self._cancelled: bool = False
         super(Span, self).__init__(labels=labels, start=start)
         self.timestamp = transaction.timestamp + (self.start_time - transaction.start_time)
         if self.transaction._breakdown:
@@ -564,6 +575,8 @@ def to_dict(self) -> dict:
             if self.context is None:
                 self.context = {}
             self.context["tags"] = self.labels
+        if self.links:
+            result["links"] = self.links
         if self.context:
             self.autofill_resource_context()
             # otel attributes and spankind need to be top-level
@@ -666,6 +679,8 @@ def report(self) -> None:
         if self.discardable and self.duration < self.transaction.config_exit_span_min_duration:
             self.transaction.track_dropped_span(self)
             self.transaction.dropped_spans += 1
+        elif self._cancelled:
+            self.transaction._span_counter -= 1
         else:
             self.tracer.queue_func(SPAN, self.to_dict())
 
@@ -757,6 +772,15 @@ def autofill_resource_context(self):
                 if "type" not in self.context["destination"]["service"]:
                     self.context["destination"]["service"]["type"] = ""
 
+    def cancel(self) -> None:
+        """
+        Mark span as cancelled. Cancelled spans don't count towards started spans nor dropped spans.
+
+        No checks are made to ensure that spans which already propagated distributed context are not
+        cancelled.
+        """
+        self._cancelled = True
+
     def __str__(self):
         return "{}/{}/{}".format(self.name, self.type, self.subtype)