feat: Add support to Kafka spans.

This commit adds the essential support to handle Kafka spans independently of which supported Python package is used.

Signed-off-by: Paulo Vital <[email protected]>

Author: pvital

src/instana/span/kind.py

@@ -31,6 +31,7 @@
     "tornado-server",
     "gcps-consumer",
     "asgi",
+    "kafka-consumer",
 )
 
 EXIT_SPANS = (
@@ -53,6 +54,7 @@
     "pymongo",
     "gcs",
     "gcps-producer",
+    "kafka-producer",
 )
 
 REGISTERED_SPANS = LOCAL_SPANS + ENTRY_SPANS + EXIT_SPANS

src/instana/span/registered_span.py

@@ -1,15 +1,31 @@
 # (c) Copyright IBM Corp. 2024
 
+from typing import TYPE_CHECKING, Any, Dict
+
+from opentelemetry.semconv.trace import SpanAttributes
+from opentelemetry.trace import SpanKind
+
 from instana.log import logger
 from instana.span.base_span import BaseSpan
-from instana.span.kind import ENTRY_SPANS, EXIT_SPANS, HTTP_SPANS, LOCAL_SPANS
+from instana.span.kind import (
+    ENTRY_SPANS,
+    EXIT_SPANS,
+    HTTP_SPANS,
+    LOCAL_SPANS,
+)
 
-from opentelemetry.trace import SpanKind
-from opentelemetry.semconv.trace import SpanAttributes
+if TYPE_CHECKING:
+    from instana.span.span import InstanaSpan
 
 
 class RegisteredSpan(BaseSpan):
-    def __init__(self, span, source, service_name, **kwargs) -> None:
+    def __init__(
+        self,
+        span: "InstanaSpan",
+        source: Dict[str, Any],
+        service_name: str,
+        **kwargs: Dict[str, Any],
+    ) -> None:
         # pylint: disable=invalid-name
         super(RegisteredSpan, self).__init__(span, source, **kwargs)
         self.n = span.name
@@ -34,13 +50,17 @@ def __init__(self, span, source, service_name, **kwargs) -> None:
         if "gcps" in span.name:
             self.n = "gcps"
 
+        # unify the span name for kafka-producer and kafka-consumer
+        if "kafka" in span.name:
+            self.n = "kafka"
+
         # Logic to store custom attributes for registered spans (not used yet)
         if len(span.attributes) > 0:
             self.data["sdk"]["custom"]["tags"] = self._validate_attributes(
                 span.attributes
             )
 
-    def _populate_entry_span_data(self, span) -> None:
+    def _populate_entry_span_data(self, span: "InstanaSpan") -> None:
         if span.name in HTTP_SPANS:
             self._collect_http_attributes(span)
 
@@ -127,10 +147,14 @@ def _populate_entry_span_data(self, span) -> None:
             self.data["rpc"]["params"] = span.attributes.pop("rpc.params", None)
             # self.data["rpc"]["baggage"] = span.attributes.pop("rpc.baggage", None)
             self.data["rpc"]["error"] = span.attributes.pop("rpc.error", None)
+
+        elif span.name.startswith("kafka"):
+            self._collect_kafka_attributes(span)
+
         else:
             logger.debug("SpanRecorder: Unknown entry span: %s" % span.name)
 
-    def _populate_local_span_data(self, span) -> None:
+    def _populate_local_span_data(self, span: "InstanaSpan") -> None:
         if span.name == "render":
             self.data["render"]["name"] = span.attributes.pop("name", None)
             self.data["render"]["type"] = span.attributes.pop("type", None)
@@ -139,7 +163,7 @@ def _populate_local_span_data(self, span) -> None:
         else:
             logger.debug("SpanRecorder: Unknown local span: %s" % span.name)
 
-    def _populate_exit_span_data(self, span) -> None:
+    def _populate_exit_span_data(self, span: "InstanaSpan") -> None:
         if span.name in HTTP_SPANS:
             self._collect_http_attributes(span)
 
@@ -239,8 +263,12 @@ def _populate_exit_span_data(self, span) -> None:
             self.data["mysql"]["host"] = span.attributes.pop("host", None)
             self.data["mysql"]["port"] = span.attributes.pop("port", None)
             self.data["mysql"]["db"] = span.attributes.pop(SpanAttributes.DB_NAME, None)
-            self.data["mysql"]["user"] = span.attributes.pop(SpanAttributes.DB_USER, None)
-            self.data["mysql"]["stmt"] = span.attributes.pop(SpanAttributes.DB_STATEMENT, None)
+            self.data["mysql"]["user"] = span.attributes.pop(
+                SpanAttributes.DB_USER, None
+            )
+            self.data["mysql"]["stmt"] = span.attributes.pop(
+                SpanAttributes.DB_STATEMENT, None
+            )
             self.data["mysql"]["error"] = span.attributes.pop("mysql.error", None)
 
         elif span.name == "postgres":
@@ -303,10 +331,14 @@ def _populate_exit_span_data(self, span) -> None:
                     self.data["log"]["parameters"] = event.attributes.pop(
                         "parameters", None
                     )
+
+        elif span.name.startswith("kafka"):
+            self._collect_kafka_attributes(span)
+
         else:
             logger.debug("SpanRecorder: Unknown exit span: %s" % span.name)
 
-    def _collect_http_attributes(self, span) -> None:
+    def _collect_http_attributes(self, span: "InstanaSpan") -> None:
         self.data["http"]["host"] = span.attributes.pop("http.host", None)
         self.data["http"]["url"] = span.attributes.pop("http.url", None)
         self.data["http"]["path"] = span.attributes.pop("http.path", None)
@@ -325,3 +357,7 @@ def _collect_http_attributes(self, span) -> None:
             for key in custom_headers:
                 trimmed_key = key[12:]
                 self.data["http"]["header"][trimmed_key] = span.attributes.pop(key)
+
+    def _collect_kafka_attributes(self, span: "InstanaSpan") -> None:
+        self.data["kafka"]["service"] = span.attributes.pop("kafka.service", None)
+        self.data["kafka"]["access"] = span.attributes.pop("kafka.access", None)

src/instana/span/span.py

@@ -165,6 +165,8 @@ def record_exception(
                 self.set_attribute("sqlalchemy.err", message)
             elif self.name == "aws.lambda.entry":
                 self.set_attribute("lambda.error", message)
+            elif self.name.startswith("kafka"):
+                self.set_attribute("kafka.error", message)
             else:
                 _attributes = {"message": message}
                 if attributes:

tests/span/test_registered_span.py

@@ -146,6 +146,14 @@ def test_populate_local_span_data_with_other_name(
                     "rpc.port": 1234,
                 },
             ),
+            (
+                "kafka-consumer",
+                "kafka",
+                {
+                    "kafka.service": "my-topic",
+                    "kafka.access": "consume",
+                },
+            ),
         ],
     )
     def test_populate_entry_span_data(
@@ -350,6 +358,14 @@ def test_populate_entry_span_data_AWSlambda(
                     "gcps.top": "MY_SUBSCRIPTION_NAME",
                 },
             ),
+            (
+                "kafka-producer",
+                "kafka",
+                {
+                    "kafka.service": "my-topic",
+                    "kafka.access": "send",
+                },
+            ),
         ],
     )
     def test_populate_exit_span_data(
@@ -453,3 +469,29 @@ def test_populate_exit_span_data_log(
 
         while self.span._events:
             self.span._events.pop()
+
+    def test_collect_kafka_attributes(
+        self,
+        span_context: SpanContext,
+        span_processor: StanRecorder,
+    ) -> None:
+        span_name = "test-kafka-registered-span"
+        attributes = {
+            "kafka.service": "my-topic",
+            "kafka.access": "send",
+        }
+        service_name = "test-kafka-registered-service"
+        self.span = InstanaSpan(
+            span_name, span_context, span_processor, attributes=attributes
+        )
+        reg_span = RegisteredSpan(self.span, None, service_name)
+
+        excepted_result = {
+            "kafka.service": attributes["kafka.service"],
+            "kafka.access": attributes["kafka.access"],
+        }
+
+        reg_span._collect_kafka_attributes(self.span)
+
+        assert excepted_result["kafka.service"] == reg_span.data["kafka"]["service"]
+        assert excepted_result["kafka.access"] == reg_span.data["kafka"]["access"]

tests/span/test_span.py

@@ -569,6 +569,7 @@ def test_span_add_event(
             ("celery-worker", "error"),
             ("sqlalchemy", "sqlalchemy.err"),
             ("aws.lambda.entry", "lambda.error"),
+            ("kafka", "kafka.error"),
         ],
     )
     def test_span_record_exception_default(