code: /*
   * 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.analyzer.agent.kafka;
  
  import com.google.common.collect.ImmutableMap;
  
  import java.time.Duration;
  import java.util.Collection;
  import java.util.List;
  import java.util.Objects;
  import java.util.Properties;
  import java.util.Set;
  import java.util.concurrent.ArrayBlockingQueue;
  import java.util.concurrent.ExecutionException;
  import java.util.concurrent.ThreadPoolExecutor;
  import java.util.concurrent.TimeUnit;
  import java.util.stream.Collectors;
  
  import com.google.common.collect.Lists;
  import lombok.extern.slf4j.Slf4j;
  import org.apache.kafka.clients.admin.AdminClient;
  import org.apache.kafka.clients.admin.NewTopic;
  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.BytesDeserializer;
  import org.apache.kafka.common.serialization.StringDeserializer;
  import org.apache.kafka.common.utils.Bytes;
  import org.apache.skywalking.oap.server.analyzer.agent.kafka.module.KafkaFetcherConfig;
  import org.apache.skywalking.oap.server.analyzer.agent.kafka.provider.handler.KafkaHandler;
  import org.apache.skywalking.oap.server.library.module.ModuleStartException;
  import org.apache.skywalking.oap.server.library.server.pool.CustomThreadFactory;
  
  /**
   * Configuring and initializing a KafkaConsumer client as a dispatcher to delivery Kafka Message to registered handler
   * by topic.
   */
  @Slf4j
  public class KafkaFetcherHandlerRegister {
  
      private ImmutableMap.Builder<String, KafkaHandler> builder = ImmutableMap.builder();
      private ImmutableMap<String, KafkaHandler> handlerMap;
  
      private final KafkaFetcherConfig config;
      private final Properties properties;
  
      private final ThreadPoolExecutor executor;
      private final boolean enableKafkaMessageAutoCommit;
      private final List<KafkaConsumer<String, Bytes>> consumers = Lists.newArrayList();
  
      public KafkaFetcherHandlerRegister(KafkaFetcherConfig config) {
          this.config = config;
  
          properties = new Properties();
          properties.setProperty(ConsumerConfig.GROUP_ID_CONFIG, config.getGroupId());
          properties.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, config.getBootstrapServers());
          properties.putAll(config.getKafkaConsumerConfig());
  
          int threadPoolSize = Runtime.getRuntime().availableProcessors() * 2;
          if (config.getKafkaHandlerThreadPoolSize() > 0) {
              threadPoolSize = config.getKafkaHandlerThreadPoolSize();
          }
          int threadPoolQueueSize = 10000;
          if (config.getKafkaHandlerThreadPoolQueueSize() > 0) {
              threadPoolQueueSize = config.getKafkaHandlerThreadPoolQueueSize();
          }
  
          enableKafkaMessageAutoCommit = (boolean) properties.getOrDefault(
              ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, true);
          for (int i = 0; i < config.getConsumers(); i++) {
              KafkaConsumer<String, Bytes> consumer = new KafkaConsumer<>(
                  properties, new StringDeserializer(), new BytesDeserializer());
  
              consumers.add(consumer);
          }
          executor = new ThreadPoolExecutor(threadPoolSize, threadPoolSize, 60, TimeUnit.SECONDS,
                                            new ArrayBlockingQueue<>(threadPoolQueueSize),
                                            new CustomThreadFactory("KafkaConsumer"),
                                            new ThreadPoolExecutor.CallerRunsPolicy()
          );
      }
  
      public void register(KafkaHandler handler) {
          builder.put(handler.getTopic(), handler);
      }
  
      public void start() throws ModuleStartException {
          handlerMap = builder.build();
          builder = null;
  
          createTopicIfNeeded(handlerMap.keySet(), properties);
          for (KafkaConsumer<String, Bytes> consumer : consumers) {
              consumer.subscribe(handlerMap.keySet());
              consumer.seekToEnd(consumer.assignment());
              executor.submit(() -> runTask(consumer));
          }
      }
  
      private void runTask(final KafkaConsumer<String, Bytes> consumer) {
          while (true) {
              try {
                  ConsumerRecords<String, Bytes> consumerRecords = consumer.poll(Duration.ofMillis(500L));
                  if (!consumerRecords.isEmpty()) {
                      for (final ConsumerRecord<String, Bytes> record : consumerRecords) {
                          executor.submit(() -> Objects.requireNonNull(handlerMap.get(record.topic())).handle(record));
                      }
                      if (!enableKafkaMessageAutoCommit) {
                          consumer.commitAsync();
                      }
                  }
              } catch (Exception e) {
                  log.error("Kafka handle message error.", e);
              }
          }
      }
  
      private void createTopicIfNeeded(Collection<String> topics, Properties properties) throws ModuleStartException {
          Properties adminProps = new Properties();
          adminProps.putAll(properties);
          // remove 'group.id' to avoid unknown configure warning.
          adminProps.remove(ConsumerConfig.GROUP_ID_CONFIG);
  
          AdminClient adminClient = AdminClient.create(adminProps);
          Set<String> missedTopics = adminClient.describeTopics(topics)
                                                .values()
                                                .entrySet()
                                                .stream()
                                                .map(entry -> {
                                                    try {
                                                        entry.getValue().get();
                                                        return null;
                                                    } catch (InterruptedException | ExecutionException ignore) {
                                                    }
                                                    return entry.getKey();
                                                })
                                                .filter(Objects::nonNull)
                                                .collect(Collectors.toSet());
  
          if (!missedTopics.isEmpty()) {
              log.info("Topics " + missedTopics + " not exist.");
              List<NewTopic> newTopicList = missedTopics.stream()
                                                        .map(topic -> new NewTopic(
                                                            topic,
                                                            config.getPartitions(),
                                                            (short) config.getReplicationFactor()
                                                        )).collect(Collectors.toList());
  
              try {
                  adminClient.createTopics(newTopicList).all().get();
              } catch (Exception e) {
                  throw new ModuleStartException("Failed to create Kafka Topics" + missedTopics + ".", e);
              }
          }
      }
  }
  