Automatically create a monitoring topic if one has not already been created. (linkedin#22)

smccauliff · lindong28 · commit 54483431cc45 · 2016-10-07T09:45:01.000-07:00
Automatically create a monitoring topic if one has not already been created
diff --git a/config/kafka-monitor.properties b/config/kafka-monitor.properties
@@ -39,6 +39,8 @@
     "zookeeper.connect": "localhost:2181",
     "bootstrap.servers": "localhost:9092",
     "produce.record.delay.ms": 100,
+    "produce.topic.autoTopicCreationEnabled": true,
+    "produce.topic.autoTopicReplicationFactor" : 1,
     "produce.producer.props": {
       "client.id": "kmf-client-id"
     },
diff --git a/services/src/main/java/com/linkedin/kmf/common/Utils.java b/services/src/main/java/com/linkedin/kmf/common/Utils.java
@@ -9,6 +9,8 @@
  */
 package com.linkedin.kmf.common;
 
+import java.util.Properties;
+import kafka.server.KafkaConfig;
 import org.apache.avro.generic.GenericData;
 import org.apache.avro.generic.GenericDatumWriter;
 import org.apache.avro.generic.GenericRecord;
@@ -19,14 +21,22 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import kafka.utils.ZkUtils;
+import kafka.admin.AdminUtils;
 import scala.collection.Seq;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.util.Arrays;
 
+
+/**
+ * Kafka monitoring utilities.
+ */
 public class Utils {
   private static final Logger LOG = LoggerFactory.getLogger(Utils.class);
 
+  private static final int ZK_CONNECTION_TIMEOUT_MS = 30_000;
+  private static final int ZK_SESSION_TIMEOUT_MS = 30_000;
+
   /**
    * Read number of partitions for the given topic on the specified zookeeper
    * @param zkUrl zookeeper connection url
@@ -35,11 +45,56 @@ public class Utils {
    * @return the number of partitions of the given topic
    */
   public static int getPartitionNumForTopic(String zkUrl, String topic) {
-    ZkUtils zkUtils = ZkUtils.apply(zkUrl, 30000, 30000, JaasUtils.isZkSecurityEnabled());
-    Seq<String> topics = scala.collection.JavaConversions.asScalaBuffer(Arrays.asList(topic));
-    int partition = zkUtils.getPartitionsForTopics(topics).apply(topic).size();
-    zkUtils.close();
-    return partition;
+    ZkUtils zkUtils = ZkUtils.apply(zkUrl, ZK_SESSION_TIMEOUT_MS, ZK_CONNECTION_TIMEOUT_MS, JaasUtils.isZkSecurityEnabled());
+    try {
+      Seq<String> topics = scala.collection.JavaConversions.asScalaBuffer(Arrays.asList(topic));
+      return zkUtils.getPartitionsForTopics(topics).apply(topic).size();
+    } finally {
+      zkUtils.close();
+    }
+  }
+
+  /**
+   * Create the topic that the monitor uses to monitor the cluster.  This method attempts to create a topic so that all
+   * the brokers in the cluster will have partitionFactor partitions.  If the topic exists, but has different parameters
+   * then this does nothing to update the parameters.
+   *
+   * TODO: Do we care about rack aware mode?  I would think no because we want to spread the topic over all brokers.
+   * @param zkUrl zookeeper connection url
+   * @param topic topic name
+   * @param replicationFactor the replication factor for the topic
+   * @param partitionFactor This is multiplied by the number brokers to compute the number of partitions in the topic.
+   * @return the number of partitions created
+   */
+  public static int createMonitoringTopicIfNotExists(String zkUrl, String topic, int replicationFactor,
+      int partitionFactor) {
+    ZkUtils zkUtils = ZkUtils.apply(zkUrl, ZK_SESSION_TIMEOUT_MS, ZK_CONNECTION_TIMEOUT_MS, JaasUtils.isZkSecurityEnabled());
+    try {
+      if (AdminUtils.topicExists(zkUtils, topic)) {
+        LOG.info("Monitoring topic \"" + topic + "\" already exists.");
+        return getPartitionNumForTopic(zkUrl, topic);
+      }
+
+      int brokerCount = zkUtils.getAllBrokersInCluster().size();
+
+      if (partitionFactor <= 0) {
+        throw new IllegalArgumentException("Partition factor must be greater than zero, but was configured for " +
+            partitionFactor + ".");
+      }
+      int partitionCount = brokerCount * partitionFactor;
+
+      int minIsr = Math.max(replicationFactor - 1, 1);
+      Properties topicConfig = new Properties();
+      topicConfig.setProperty(KafkaConfig.MinInSyncReplicasProp(), Integer.toString(minIsr));
+      AdminUtils.createTopic(zkUtils, topic, partitionCount, replicationFactor, topicConfig);
+
+      LOG.info("Created monitoring topic \"" + topic + "\" with " + partitionCount + " partitions, min ISR of " + minIsr
+          + " and replication factor of " + replicationFactor + ".");
+
+      return partitionCount;
+    } finally {
+      zkUtils.close();
+    }
   }
 
   /**
diff --git a/services/src/main/java/com/linkedin/kmf/services/ProduceService.java b/services/src/main/java/com/linkedin/kmf/services/ProduceService.java
@@ -14,6 +14,9 @@
 import com.linkedin.kmf.producer.KMBaseProducer;
 import com.linkedin.kmf.producer.BaseProducerRecord;
 import com.linkedin.kmf.producer.NewProducer;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.apache.kafka.clients.producer.ProducerConfig;
 import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.common.MetricName;
@@ -39,6 +42,10 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
 
+
+/**
+ * This sets up the producers used by Kafka Monitoring and some of the reported metrics.
+ */
 public class ProduceService implements Service {
   private static final Logger LOG = LoggerFactory.getLogger(ProduceService.class);
   private static final String METRIC_GROUP_NAME = "produce-service";
@@ -49,31 +56,57 @@ public class ProduceService implements Service {
   private final ScheduledExecutorService _executor;
   private final int _produceDelayMs;
   private final boolean _sync;
-  private final Map<Integer, AtomicLong> _nextIndexPerPartition;
-  private final int _partitionNum;
+  /** This can be updated while running when new partitions are added to the monitored topic. */
+  private final ConcurrentMap<Integer, AtomicLong> _nextIndexPerPartition;
+  /** This is the last thing that should be updated after adding new partitions. */
+  private final AtomicInteger _partitionNum;
   private final int _recordSize;
   private final String _topic;
   private final String _producerId;
   private final AtomicBoolean _running;
+  private final String _zkConnect;
+  private final String _brokerList;
+  private final double _rebalanceThreshold;
 
   public ProduceService(Map<String, Object> props, String name) throws Exception {
     _name = name;
     Map producerPropsOverride = (Map) props.get(ProduceServiceConfig.PRODUCER_PROPS_CONFIG);
     ProduceServiceConfig config = new ProduceServiceConfig(props);
-    String zkConnect = config.getString(ProduceServiceConfig.ZOOKEEPER_CONNECT_CONFIG);
-    String brokerList = config.getString(ProduceServiceConfig.BOOTSTRAP_SERVERS_CONFIG);
+    _zkConnect = config.getString(ProduceServiceConfig.ZOOKEEPER_CONNECT_CONFIG);
+    _brokerList = config.getString(ProduceServiceConfig.BOOTSTRAP_SERVERS_CONFIG);
+    _rebalanceThreshold = config.getDouble(ProduceServiceConfig.REBALANCE_THRESHOLD_CONFIG);
     String producerClass = config.getString(ProduceServiceConfig.PRODUCER_CLASS_CONFIG);
     int threadsNum = config.getInt(ProduceServiceConfig.PRODUCE_THREAD_NUM_CONFIG);
     _topic = config.getString(ProduceServiceConfig.TOPIC_CONFIG);
     _producerId = config.getString(ProduceServiceConfig.PRODUCER_ID_CONFIG);
     _produceDelayMs = config.getInt(ProduceServiceConfig.PRODUCE_RECORD_DELAY_MS_CONFIG);
     _recordSize = config.getInt(ProduceServiceConfig.PRODUCE_RECORD_SIZE_BYTE_CONFIG);
     _sync = config.getBoolean(ProduceServiceConfig.PRODUCE_SYNC_CONFIG);
-    _partitionNum = Utils.getPartitionNumForTopic(zkConnect, _topic);
+    _partitionNum = new AtomicInteger(0);
+
+    if (_rebalanceThreshold < 1) {
+      throw new IllegalArgumentException("Rebalance threshold must be greater than one but is set to " + _rebalanceThreshold + ".");
+    }
+
+    int existingPartitionCount = Utils.getPartitionNumForTopic(_zkConnect, _topic);
 
-    if (_partitionNum <= 0)
-      throw new RuntimeException("Can not find valid partition number for topic " + _topic +
-        ". Please verify that the topic has been created. Ideally the partition number should be a multiple of number of brokers in the cluster.");
+    if (existingPartitionCount <= 0) {
+      if (config.getBoolean(ProduceServiceConfig.AUTO_TOPIC_CREATION_ENABLED_CONFIG)) {
+        int autoTopicReplicationFactor = config.getInt(ProduceServiceConfig.AUTO_TOPIC_REPLICATION_FACTOR_CONFIG);
+        int autoTopicPartitionFactor = config.getInt(ProduceServiceConfig.REBALANCE_PARTITION_MULTIPLE_CONFIG);
+        _partitionNum.set(
+            Utils.createMonitoringTopicIfNotExists(_zkConnect, _topic, autoTopicReplicationFactor,
+                autoTopicPartitionFactor));
+      } else {
+        throw new RuntimeException("Can not find valid partition number for topic " + _topic +
+            ". Please verify that the topic \"" + _topic + "\" has been created. Ideally the partition number should be" +
+            " a multiple of number" +
+            " of brokers in the cluster.  Or else configure " + ProduceServiceConfig.AUTO_TOPIC_CREATION_ENABLED_CONFIG +
+            " to be true.");
+      }
+    } else {
+      _partitionNum.set(existingPartitionCount);
+    }
 
     Properties producerProps = new Properties();
 
@@ -92,7 +125,7 @@ public ProduceService(Map<String, Object> props, String name) throws Exception {
 
     // Assign config specified for ProduceService.
     producerProps.put(ProducerConfig.CLIENT_ID_CONFIG, _producerId);
-    producerProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList);
+    producerProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, _brokerList);
 
     // Assign config specified for producer. This has the highest priority.
     if (producerPropsOverride != null)
@@ -101,7 +134,7 @@ public ProduceService(Map<String, Object> props, String name) throws Exception {
     _producer = (KMBaseProducer) Class.forName(producerClass).getConstructor(Properties.class).newInstance(producerProps);
 
     _running = new AtomicBoolean(false);
-    _nextIndexPerPartition = new HashMap<>();
+    _nextIndexPerPartition = new ConcurrentHashMap<>();
     _executor = Executors.newScheduledThreadPool(threadsNum);
 
     MetricConfig metricConfig = new MetricConfig().samples(60).timeWindow(1000, TimeUnit.MILLISECONDS);
@@ -111,20 +144,26 @@ public ProduceService(Map<String, Object> props, String name) throws Exception {
     Map<String, String> tags = new HashMap<>();
     tags.put("name", _name);
     _sensors = new ProduceMetrics(metrics, tags);
+
   }
 
   @Override
   public void start() {
     if (_running.compareAndSet(false, true)) {
-      for (int partition = 0; partition < _partitionNum; partition++) {
-        _nextIndexPerPartition.put(partition, new AtomicLong(0));
-        _executor.scheduleWithFixedDelay(new ProduceRunnable(partition),
-          _produceDelayMs, _produceDelayMs, TimeUnit.MILLISECONDS);
+      int partitionNum = _partitionNum.get();
+      for (int partition = 0; partition < partitionNum; partition++) {
+        scheduleProduceRunnable(partition);
       }
       LOG.info(_name + "/ProduceService started");
     }
   }
 
+  private void scheduleProduceRunnable(int partition) {
+    _nextIndexPerPartition.put(partition, new AtomicLong(0));
+    _executor.scheduleWithFixedDelay(new ProduceRunnable(partition),
+        _produceDelayMs, _produceDelayMs, TimeUnit.MILLISECONDS);
+  }
+
   @Override
   public void stop() {
     if (_running.compareAndSet(true, false)) {
@@ -153,24 +192,20 @@ private class ProduceMetrics {
     public final Metrics metrics;
     private final Sensor _recordsProduced;
     private final Sensor _produceError;
-    private final Map<Integer, Sensor> _recordsProducedPerPartition;
-    private final Map<Integer, Sensor> _produceErrorPerPartition;
+    private final ConcurrentMap<Integer, Sensor> _recordsProducedPerPartition;
+    private final ConcurrentMap<Integer, Sensor> _produceErrorPerPartition;
+    private final Map<String, String> _tags;
 
     public ProduceMetrics(Metrics metrics, final Map<String, String> tags) {
       this.metrics = metrics;
+      this._tags = tags;
 
-      _recordsProducedPerPartition = new HashMap<>();
-      for (int partition = 0; partition < _partitionNum; partition++) {
-        Sensor sensor = metrics.sensor("records-produced-partition-" + partition);
-        sensor.add(new MetricName("records-produced-rate-partition-" + partition, METRIC_GROUP_NAME, "The average number of records per second that are produced to this partition", tags), new Rate());
-        _recordsProducedPerPartition.put(partition, sensor);
-      }
+      _recordsProducedPerPartition = new ConcurrentHashMap<>();
+      _produceErrorPerPartition = new ConcurrentHashMap<>();
 
-      _produceErrorPerPartition = new HashMap<>();
-      for (int partition = 0; partition < _partitionNum; partition++) {
-        Sensor sensor = metrics.sensor("produce-error-partition-" + partition);
-        sensor.add(new MetricName("produce-error-rate-partition-" + partition, METRIC_GROUP_NAME, "The average number of errors per second when producing to this partition", tags), new Rate());
-        _produceErrorPerPartition.put(partition, sensor);
+      int partitionNum = _partitionNum.get();
+      for (int partition = 0; partition < partitionNum; partition++) {
+        addPartitionSensors(partition);
       }
 
       _recordsProduced = metrics.sensor("records-produced");
@@ -186,11 +221,12 @@ public ProduceMetrics(Metrics metrics, final Map<String, String> tags) {
           @Override
           public double measure(MetricConfig config, long now) {
             double availabilitySum = 0.0;
-            for (int partition = 0; partition < _partitionNum; partition++) {
+            int partitionNum = _partitionNum.get();
+            for (int partition = 0; partition < partitionNum; partition++) {
               double recordsProduced = _sensors.metrics.metrics().get(new MetricName("records-produced-rate-partition-" + partition, METRIC_GROUP_NAME, tags)).value();
               double produceError = _sensors.metrics.metrics().get(new MetricName("produce-error-rate-partition-" + partition, METRIC_GROUP_NAME, tags)).value();
               // If there is no error, error rate sensor may expire and the value may be NaN. Treat NaN as 0 for error rate.
-              if (new Double(produceError).isNaN()) {
+              if (Double.isNaN(produceError) || Double.isInfinite(produceError)) {
                 produceError = 0;
               }
               // If there is either succeeded or failed produce to a partition, consider its availability as 0.
@@ -199,11 +235,23 @@ public double measure(MetricConfig config, long now) {
               }
             }
             // Assign equal weight to per-partition availability when calculating overall availability
-            return availabilitySum / _partitionNum;
+            return availabilitySum / partitionNum;
           }
         }
       );
     }
+
+    void addPartitionSensors(int partition) {
+      Sensor recordsProducedSensor = metrics.sensor("records-produced-partition-" + partition);
+      recordsProducedSensor.add(new MetricName("records-produced-rate-partition-" + partition, METRIC_GROUP_NAME,
+          "The average number of records per second that are produced to this partition", _tags), new Rate());
+      _recordsProducedPerPartition.put(partition, recordsProducedSensor);
+
+      Sensor errorsSensor = metrics.sensor("produce-error-partition-" + partition);
+      errorsSensor.add(new MetricName("produce-error-rate-partition-" + partition, METRIC_GROUP_NAME,
+          "The average number of errors per second when producing to this partition", _tags), new Rate());
+      _produceErrorPerPartition.put(partition, errorsSensor);
+    }
   }
 
   private class ProduceRunnable implements Runnable {
diff --git a/services/src/main/java/com/linkedin/kmf/services/configs/ProduceServiceConfig.java b/services/src/main/java/com/linkedin/kmf/services/configs/ProduceServiceConfig.java
@@ -49,6 +49,26 @@ public class ProduceServiceConfig extends AbstractConfig {
   public static final String PRODUCER_PROPS_CONFIG = "produce.producer.props";
   public static final String PRODUCER_PROPS_DOC = "The properties used to config producer in produce service.";
 
+  public static final String AUTO_TOPIC_REPLICATION_FACTOR_CONFIG = "produce.topic.autoTopicReplicationFactor";
+  public static final String AUTO_TOPIC_REPLICATION_FACTOR_DOC = "When a topic is created automatically this is the "
+      + "replication factor used.";
+
+  public static final String REBALANCE_ENABLED_CONFIG = "produce.topic.rebalanceEnabled";
+  public static final String REBALANCE_ENABLED_DOC = "Periodically move leaders and replica assignments so they are "
+      + "evenly spread amongst brokers.";
+
+  public static final String REBALANCE_PARTITION_MULTIPLE_CONFIG = "produce.topic.rebalancePartitionMultiple";
+  public static final String REBALANCE_PARTITION_MULTIPLE_DOC = "Determines the number of partitions per broker in the ideal case.";
+
+  public static final String REBALANCE_THRESHOLD_CONFIG = "produce.topic.rebalanceThreshold";
+  public static final String REBALANCE_THRESHOLD_DOC = "Determines the number of partitions per broker in the ideal case.";
+
+  public static final String AUTO_TOPIC_CREATION_ENABLED_CONFIG = "produce.topic.autoTopicCreationEnabled";
+  public static final String AUTO_TOPIC_CREATION_ENABLED_DOC = "When true this automatically creates the topic mentioned by \"" +
+      TOPIC_CONFIG + "\" with replication factor \"" + AUTO_TOPIC_REPLICATION_FACTOR_CONFIG + "and min ISR of max(" +
+      AUTO_TOPIC_REPLICATION_FACTOR_CONFIG + "-1, 1) with number of brokers * \"" + REBALANCE_PARTITION_MULTIPLE_CONFIG +
+      "\" partitions.";
+
   static {
     CONFIG = new ConfigDef().define(ZOOKEEPER_CONNECT_CONFIG,
                                     ConfigDef.Type.STRING,
@@ -63,6 +83,10 @@ public class ProduceServiceConfig extends AbstractConfig {
                                     "kafka-monitor-topic",
                                     ConfigDef.Importance.MEDIUM,
                                     TOPIC_DOC)
+                            .define(AUTO_TOPIC_CREATION_ENABLED_CONFIG,
+                                    ConfigDef.Type.BOOLEAN,
+                                    true,
+                                    ConfigDef.Importance.MEDIUM, AUTO_TOPIC_CREATION_ENABLED_DOC)
                             .define(PRODUCER_CLASS_CONFIG,
                                     ConfigDef.Type.STRING,
                                     NewProducer.class.getCanonicalName(),
@@ -92,7 +116,27 @@ public class ProduceServiceConfig extends AbstractConfig {
                                     ConfigDef.Type.INT,
                                     5,
                                     ConfigDef.Importance.LOW,
-                                    PRODUCE_THREAD_NUM_DOC);
+                                    PRODUCE_THREAD_NUM_DOC)
+                            .define(AUTO_TOPIC_REPLICATION_FACTOR_CONFIG,
+                                    ConfigDef.Type.INT,
+                                    1,
+                                    ConfigDef.Importance.LOW,
+                                    AUTO_TOPIC_REPLICATION_FACTOR_DOC)
+                            .define(REBALANCE_ENABLED_CONFIG,
+                                    ConfigDef.Type.BOOLEAN,
+                                    false,
+                                    ConfigDef.Importance.LOW,
+                                    REBALANCE_ENABLED_DOC)
+                            .define(REBALANCE_PARTITION_MULTIPLE_CONFIG,
+                                    ConfigDef.Type.INT,
+                                    2,
+                                    ConfigDef.Importance.LOW,
+                                    REBALANCE_PARTITION_MULTIPLE_DOC)
+                            .define(REBALANCE_THRESHOLD_CONFIG,
+                                    ConfigDef.Type.DOUBLE,
+                                    1.5,
+                                    ConfigDef.Importance.LOW,
+                                    REBALANCE_THRESHOLD_DOC);
 
   }
 
diff --git a/tests/src/main/java/com/linkedin/kmf/tests/BasicEndToEndTest.java b/tests/src/main/java/com/linkedin/kmf/tests/BasicEndToEndTest.java
