oap-server/server-receiver-plugin/zipkin-receiver-plugin/src/main/java/org/apache/skywalking/oap/server/receiver/zipkin/kafka/KafkaHandler.java

/*
 * 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.
 *
 */

package org.apache.skywalking.oap.server.receiver.zipkin.kafka;

import com.google.gson.Gson;
import java.time.Duration;
import java.util.Arrays;
import java.util.List;
import java.util.Properties;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.common.serialization.ByteArrayDeserializer;
import org.apache.skywalking.oap.server.library.module.ModuleManager;
import org.apache.skywalking.oap.server.library.module.ModuleStartException;
import org.apache.skywalking.oap.server.library.server.pool.CustomThreadFactory;
import org.apache.skywalking.oap.server.receiver.zipkin.ZipkinReceiverConfig;
import org.apache.skywalking.oap.server.receiver.zipkin.trace.SpanForward;
import org.apache.skywalking.oap.server.telemetry.TelemetryModule;
import org.apache.skywalking.oap.server.telemetry.api.CounterMetrics;
import org.apache.skywalking.oap.server.telemetry.api.HistogramMetrics;
import org.apache.skywalking.oap.server.telemetry.api.MetricsCreator;
import org.apache.skywalking.oap.server.telemetry.api.MetricsTag;
import zipkin2.Span;
import zipkin2.SpanBytesDecoderDetector;
import zipkin2.codec.BytesDecoder;

@Slf4j
public class KafkaHandler {
    private final ZipkinReceiverConfig config;
    private final SpanForward spanForward;
    private final Properties properties;
    private final ThreadPoolExecutor executor;
    private final boolean enableKafkaMessageAutoCommit;
    private final List<String> topics;
    private final CounterMetrics msgDroppedIncr;
    private final CounterMetrics errorCounter;
    private final HistogramMetrics histogram;

    public KafkaHandler(final ZipkinReceiverConfig config, ModuleManager manager) {
        this.config = config;
        this.spanForward = new SpanForward(config, manager);

        properties = new Properties();
        properties.setProperty(ConsumerConfig.GROUP_ID_CONFIG, config.getKafkaGroupId());
        properties.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, config.getKafkaBootstrapServers());
        properties.setProperty(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName());
        properties.setProperty(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName());
        Gson gson = new Gson();
        Properties override = gson.fromJson(config.getKafkaConsumerConfig(), Properties.class);
        properties.putAll(override);

        int threadPoolSize = Runtime.getRuntime().availableProcessors() * 2;
        if (config.getKafkaHandlerThreadPoolSize() > 0) {
            threadPoolSize = config.getKafkaHandlerThreadPoolSize();
        }
        int threadPoolQueueSize = 10000;
        if (config.getKafkaHandlerThreadPoolQueueSize() > 0) {
            threadPoolQueueSize = config.getKafkaHandlerThreadPoolQueueSize();
        }
        enableKafkaMessageAutoCommit = new ConsumerConfig(properties).getBoolean(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG);
        executor = new ThreadPoolExecutor(threadPoolSize, threadPoolSize, 60, TimeUnit.SECONDS,
                                          new ArrayBlockingQueue<>(threadPoolQueueSize),
                                          new CustomThreadFactory("Zipkin-Kafka-Consumer"),
                                          new ThreadPoolExecutor.CallerRunsPolicy()
        );

        topics = Arrays.asList(config.getKafkaTopic().split(","));
        MetricsCreator metricsCreator = manager.find(TelemetryModule.NAME)
                                               .provider()
                                               .getService(MetricsCreator.class);
        histogram = metricsCreator.createHistogramMetric(
            "trace_in_latency",
            "The process latency of trace data",
            new MetricsTag.Keys("protocol"),
            new MetricsTag.Values("zipkin-kafka")
        );
        msgDroppedIncr = metricsCreator.createCounter(
            "trace_dropped_count", "The dropped number of traces",
            new MetricsTag.Keys("protocol"), new MetricsTag.Values("zipkin-kafka"));
        errorCounter = metricsCreator.createCounter(
            "trace_analysis_error_count", "The error number of trace analysis",
            new MetricsTag.Keys("protocol"), new MetricsTag.Values("zipkin-kafka")
        );
    }

    public void start() throws ModuleStartException {
        for (int i = 0; i < config.getKafkaConsumers(); i++) {
            KafkaConsumer<byte[], byte[]> consumer = new KafkaConsumer<>(properties);
            consumer.subscribe(topics);
            consumer.seekToEnd(consumer.assignment());
            executor.submit(() -> runTask(consumer));
        }
    }

    private void runTask(final KafkaConsumer<byte[], byte[]> consumer) {
        if (log.isDebugEnabled()) {
            log.debug("Start Consume zipkin trace records from kafka.");
        }
        while (true) {
            try {
                ConsumerRecords<byte[], byte[]> consumerRecords = consumer.poll(Duration.ofMillis(1000L));
                if (log.isDebugEnabled()) {
                    log.debug(
                        "Consume zipkin trace records from kafka, records count:[{}].",
                        consumerRecords.count()
                    );
                }
                if (!consumerRecords.isEmpty()) {
                    for (final ConsumerRecord<byte[], byte[]> record : consumerRecords) {
                        final byte[] bytes = record.value();
                        //empty or illegal message
                        if (bytes.length < 2) {
                            msgDroppedIncr.inc();
                            continue;
                        }
                        executor.submit(() -> handleRecord(bytes));
                    }
                    if (!enableKafkaMessageAutoCommit) {
                        consumer.commitAsync();
                    }
                }
            } catch (Exception e) {
                log.error("Kafka handle message error.", e);
                errorCounter.inc();
            }
        }
    }

    private void handleRecord(byte[] bytes) {
        try (HistogramMetrics.Timer ignored = histogram.createTimer()) {
            BytesDecoder<Span> decoder = SpanBytesDecoderDetector.decoderForListMessage(bytes);
            final List<Span> spanList = decoder.decodeList(bytes);
            spanForward.send(spanList);
        }
    }
}