Support Kafka reporter protocol (#74)

Humbertzhang · web-flow · commit 50388c55428d · 2020-10-10T00:19:04.000+08:00
diff --git a/docs/EnvVars.md b/docs/EnvVars.md
@@ -22,3 +22,7 @@ Environment Variable | Description | Default
 | `SW_TRACE_IGNORE`| This config item controls that whether the trace should be ignore | `false` |
 | `SW_TRACE_IGNORE_PATH`| You can setup multiple URL path patterns, The endpoints match these patterns wouldn't be traced. the current matching rules follow Ant Path match style , like /path/*, /path/**, /path/?.| `''` |
 | `SW_ELASTICSEARCH_TRACE_DSL`| If true, trace all the DSL(Domain Specific Language) in ElasticSearch access, default is false | `false` |
+| `SW_KAFKA_REPORTER_BOOTSTRAP_SERVERS` | A list of host/port pairs to use for establishing the initial connection to the Kafka cluster. It is in the form host1:port1,host2:port2,... | `localhost:9092` |
+| `SW_KAFKA_REPORTER_TOPIC_MANAGEMENT` | Specifying Kafka topic name for service instance reporting and registering. | `skywalking-managements` | 
+| `SW_KAFKA_REPORTER_TOPIC_SEGMENT` | Specifying Kafka topic name for Tracing data. | `skywalking-segments` |
+| `SW_KAFKA_REPORTER_CONFIG_key` | The configs to init KafkaProducer. it support the basic arguments (whose type is either `str`, `bool`, or `int`) listed [here](https://kafka-python.readthedocs.io/en/master/apidoc/KafkaProducer.html#kafka.KafkaProducer) | unset |
diff --git a/skywalking/agent/__init__.py b/skywalking/agent/__init__.py
@@ -62,6 +62,9 @@ def __init():
     elif config.protocol == 'http':
         from skywalking.agent.protocol.http import HttpProtocol
         __protocol = HttpProtocol()
+    elif config.protocol == "kafka":
+        from skywalking.agent.protocol.kafka import KafkaProtocol
+        __protocol = KafkaProtocol()
 
     plugins.install()
 
diff --git a/skywalking/agent/protocol/kafka.py b/skywalking/agent/protocol/kafka.py
@@ -0,0 +1,94 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import logging
+from queue import Queue
+
+from skywalking import config
+from skywalking.agent import Protocol
+from skywalking.client.kafka import KafkaServiceManagementClient, KafkaTraceSegmentReportService
+from skywalking.protocol.common.Common_pb2 import KeyStringValuePair
+from skywalking.protocol.language_agent.Tracing_pb2 import SegmentObject, SpanObject, Log, SegmentReference
+from skywalking.trace.segment import Segment
+
+logger = logging.getLogger(__name__)
+
+# avoid too many kafka logs
+logger_kafka = logging.getLogger('kafka')
+logger_kafka.setLevel(logging.WARN)
+
+
+class KafkaProtocol(Protocol):
+    def __init__(self):
+        self.service_management = KafkaServiceManagementClient()
+        self.traces_reporter = KafkaTraceSegmentReportService()
+
+    def connected(self):
+        return True
+
+    def heartbeat(self):
+        self.service_management.send_heart_beat()
+
+    def report(self, queue: Queue):
+        def generator():
+            while True:
+                segment = queue.get()  # type: Segment
+
+                logger.debug('reporting segment %s', segment)
+
+                s = SegmentObject(
+                    traceId=str(segment.related_traces[0]),
+                    traceSegmentId=str(segment.segment_id),
+                    service=config.service_name,
+                    serviceInstance=config.service_instance,
+                    spans=[SpanObject(
+                        spanId=span.sid,
+                        parentSpanId=span.pid,
+                        startTime=span.start_time,
+                        endTime=span.end_time,
+                        operationName=span.op,
+                        peer=span.peer,
+                        spanType=span.kind.name,
+                        spanLayer=span.layer.name,
+                        componentId=span.component.value,
+                        isError=span.error_occurred,
+                        logs=[Log(
+                            time=int(log.timestamp * 1000),
+                            data=[KeyStringValuePair(key=item.key, value=item.val) for item in log.items],
+                        ) for log in span.logs],
+                        tags=[KeyStringValuePair(
+                            key=str(tag.key),
+                            value=str(tag.val),
+                        ) for tag in span.tags],
+                        refs=[SegmentReference(
+                            refType=0 if ref.ref_type == "CrossProcess" else 1,
+                            traceId=ref.trace_id,
+                            parentTraceSegmentId=ref.segment_id,
+                            parentSpanId=ref.span_id,
+                            parentService=ref.service,
+                            parentServiceInstance=ref.service_instance,
+                            parentEndpoint=ref.endpoint,
+                            networkAddressUsedAtPeer=ref.client_address,
+                        ) for ref in span.refs if ref.trace_id],
+                    ) for span in segment.spans],
+                )
+
+                yield s
+
+                queue.task_done()
+
+        self.traces_reporter.report(generator())
diff --git a/skywalking/client/kafka.py b/skywalking/client/kafka.py
@@ -0,0 +1,113 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import os
+import ast
+import logging
+
+from skywalking import config
+from skywalking.client import ServiceManagementClient, TraceSegmentReportService
+from skywalking.protocol.common.Common_pb2 import KeyStringValuePair
+from skywalking.protocol.management.Management_pb2 import InstancePingPkg, InstanceProperties
+
+from kafka import KafkaProducer
+
+logger = logging.getLogger(__name__)
+
+kafka_configs = {}
+
+
+def __init_kafka_configs():
+    kafka_configs["bootstrap_servers"] = config.kafka_bootstrap_servers.split(",")
+    # process all kafka configs in env
+    kafka_keys = [key for key in os.environ.keys() if key.startswith("SW_KAFKA_REPORTER_CONFIG_")]
+    for kafka_key in kafka_keys:
+        key = kafka_key[25:]
+        val = os.environ.get(kafka_key)
+
+        if val is not None:
+            if val.isnumeric():
+                val = int(val)
+            elif val in ["True", "False"]:
+                val = ast.literal_eval(val)
+        else:
+            continue
+
+        # check if the key was already set
+        if kafka_configs.get(key) is None:
+            kafka_configs[key] = val
+        else:
+            raise KafkaConfigDuplicated(key)
+
+
+__init_kafka_configs()
+
+
+class KafkaServiceManagementClient(ServiceManagementClient):
+    def __init__(self):
+        logger.debug("kafka reporter configs: %s", kafka_configs)
+        self.producer = KafkaProducer(**kafka_configs)
+        self.topic_key_register = "register-"
+        self.topic = config.kafka_topic_management
+
+        self.send_instance_props()
+
+    def send_instance_props(self):
+        instance = InstanceProperties(
+            service=config.service_name,
+            serviceInstance=config.service_instance,
+            properties=[KeyStringValuePair(key='language', value='Python')],
+        )
+
+        key = bytes(self.topic_key_register + instance.serviceInstance, encoding='utf-8')
+        value = bytes(instance.SerializeToString())
+        self.producer.send(topic=self.topic, key=key, value=value)
+
+    def send_heart_beat(self):
+        logger.debug(
+            'service heart beats, [%s], [%s]',
+            config.service_name,
+            config.service_instance,
+        )
+
+        instance_ping_pkg = InstancePingPkg(
+            service=config.service_name,
+            serviceInstance=config.service_instance,
+        )
+
+        key = bytes(instance_ping_pkg.serviceInstance, encoding="utf-8")
+        value = bytes(instance_ping_pkg.SerializeToString())
+        future = self.producer.send(topic=self.topic, key=key, value=value)
+        res = future.get(timeout=10)
+        logger.debug('heartbeat response: %s', res)
+
+
+class KafkaTraceSegmentReportService(TraceSegmentReportService):
+    def __init__(self):
+        self.producer = KafkaProducer(**kafka_configs)
+        self.topic = config.kafka_topic_segment
+
+    def report(self, generator):
+        for segment in generator:
+            key = bytes(segment.traceSegmentId, encoding="utf-8")
+            value = bytes(segment.SerializeToString())
+            self.producer.send(topic=self.topic, key=key, value=value)
+
+
+class KafkaConfigDuplicated(Exception):
+    def __init__(self, key):
+        self.key = key
diff --git a/skywalking/config.py b/skywalking/config.py
@@ -49,6 +49,9 @@
 trace_ignore_path = (os.getenv('SW_TRACE_IGNORE_PATH') or '').split(',')  # type: List[str]
 elasticsearch_trace_dsl = True if os.getenv('SW_ELASTICSEARCH_TRACE_DSL') and \
                                   os.getenv('SW_ELASTICSEARCH_TRACE_DSL') == 'True' else False  # type: bool
+kafka_bootstrap_servers = os.getenv('SW_KAFKA_REPORTER_BOOTSTRAP_SERVERS') or "localhost:9092"  # type: str
+kafka_topic_management = os.getenv('SW_KAFKA_REPORTER_TOPIC_MANAGEMENT') or "skywalking-managements"  # type: str
+kafka_topic_segment = os.getenv('SW_KAFKA_REPORTER_TOPIC_SEGMENT') or "skywalking-segments"  # type: str
 
 
 def init(
diff --git a/skywalking/plugins/sw_kafka.py b/skywalking/plugins/sw_kafka.py
@@ -16,6 +16,7 @@
 #
 import logging
 
+from skywalking import config
 from skywalking import Layer, Component
 from skywalking.trace import tags
 from skywalking.trace.carrier import Carrier
@@ -48,6 +49,7 @@ def _sw__poll_once(this, timeout_ms, max_records, update_offsets=True):
             context = get_context()
             topics = ";".join(this._subscription.subscription or
                               [t.topic for t in this._subscription._user_assignment])
+
             with context.new_entry_span(
                     op="Kafka/" + topics + "/Consumer/" + (this.config["group_id"] or "")) as span:
                 for consumerRecords in res.values():
@@ -71,6 +73,11 @@ def _sw__poll_once(this, timeout_ms, max_records, update_offsets=True):
 
 def _sw_send_func(_send):
     def _sw_send(this, topic, value=None, key=None, headers=None, partition=None, timestamp_ms=None):
+        # ignore trace skywalking self request
+        if config.protocol == 'kafka' and config.kafka_topic_segment == topic or config.kafka_topic_management == topic:
+            return _send(this, topic, value=value, key=key, headers=headers, partition=partition,
+                         timestamp_ms=timestamp_ms)
+
         peer = ";".join(this.config["bootstrap_servers"])
         context = get_context()
         carrier = Carrier()
