Kafka 2.4 + KIP-455 for Xinfra Monitor Multie Cluster Topic Management Service (#317)

Kafka 2.4 + KIP-455 for Xinfra Monitor Multi Cluster Topic Management Service Total extends CumulativeSum : A non-sampled cumulative total maintained over all time. This is a non-sampled version of {@link WindowedSum}. Total is deprecated since 2.4. Recommended to use {@link CumulativeSum} instead. (New) CumulativeSum: A non-sampled cumulative total maintained over all time. This is a non-sampled version of {@link WindowedSum}.

Author: Andrew Choi

build.gradle

@@ -38,7 +38,7 @@ allprojects {
  compile 'net.savantly:graphite-client:1.1.0-RELEASE'
  compile 'com.timgroup:java-statsd-client:3.0.1'
  compile 'com.signalfx.public:signalfx-codahale:0.0.47'
- compile group: 'org.apache.kafka', name: 'kafka_2.12', version: '2.3.1'
+ compile group: 'org.apache.kafka', name: 'kafka_2.12', version: '2.4.0'
  compile group: 'org.apache.kafka', name: 'kafka-clients', version: '2.3.1'
  testCompile 'org.mockito:mockito-core:2.24.0'
  testCompile 'org.testng:testng:6.8.8'

src/main/java/com/linkedin/xinfra/monitor/common/Utils.java

@@ -25,6 +25,8 @@
 import java.util.Random;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import javax.management.MBeanAttributeInfo;
 import javax.management.MBeanInfo;
 import javax.management.MBeanServer;
@@ -37,7 +39,10 @@
 import org.apache.avro.io.JsonEncoder;
 import org.apache.kafka.clients.admin.AdminClient;
 import org.apache.kafka.clients.admin.CreateTopicsResult;
+import org.apache.kafka.clients.admin.ListPartitionReassignmentsResult;
 import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.admin.PartitionReassignment;
+import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.TopicExistsException;
 import org.json.JSONObject;
 import org.slf4j.Logger;
@@ -51,6 +56,10 @@ public class Utils {
  private static final Logger LOG = LoggerFactory.getLogger(Utils.class);
  public static final int ZK_CONNECTION_TIMEOUT_MS = 30_000;
  public static final int ZK_SESSION_TIMEOUT_MS = 30_000;
+ private static final long LIST_PARTITION_REASSIGNMENTS_TIMEOUT_MS = 60000L;
+ private static final int LIST_PARTITION_REASSIGNMENTS_MAX_ATTEMPTS = 3;
+ private static final String LIST_PARTITION_REASSIGNMENTS_TIMEOUT_MS_CONFIG = "list.partition.reassignment.timeout.ms";
+ private static final int DEFAULT_RETRY_BACKOFF_BASE = 2;
 
  public static String prettyPrint(Object value) throws JsonProcessingException {
  ObjectMapper objectMapper = new ObjectMapper();
@@ -61,6 +70,41 @@ public static String prettyPrint(Object value) throws JsonProcessingException {
  return written;
  }
 
+ /**
+ * Retrieve the map of {@link PartitionReassignment reassignment} by {@link TopicPartition partitions}.
+ *
+ * If the response times out, the method retries up to {@link #LIST_PARTITION_REASSIGNMENTS_MAX_ATTEMPTS} times.
+ * The max time to wait for the {@link AdminClient adminClient} response is computed.
+ *
+ * @param adminClient The {@link AdminClient adminClient} to ask for ongoing partition reassignments
+ * @return The map of {@link PartitionReassignment reassignment} by {@link TopicPartition partitions}
+ */
+ public static Map<TopicPartition, PartitionReassignment> ongoingPartitionReassignments(AdminClient adminClient)
+ throws InterruptedException, ExecutionException, TimeoutException {
+ Map<TopicPartition, PartitionReassignment> partitionReassignments = null;
+ int attempts = 0;
+ long timeoutMs = LIST_PARTITION_REASSIGNMENTS_TIMEOUT_MS;
+ do {
+ ListPartitionReassignmentsResult responseResult = adminClient.listPartitionReassignments();
+ try {
+ // A successful response is expected to be non-null.
+ partitionReassignments = responseResult.reassignments().get(timeoutMs, TimeUnit.MILLISECONDS);
+ } catch (TimeoutException timeoutException) {
+ LOG.info(
+ "Xinfra Monitor has failed to list partition reassignments in {}ms (attempt={}). "
+ + "Please consider increasing the value of {} config.",
+ timeoutMs, 1 + attempts, LIST_PARTITION_REASSIGNMENTS_TIMEOUT_MS_CONFIG);
+ attempts++;
+ if (attempts == LIST_PARTITION_REASSIGNMENTS_MAX_ATTEMPTS) {
+ throw timeoutException;
+ }
+ timeoutMs *= DEFAULT_RETRY_BACKOFF_BASE;
+ }
+ } while (partitionReassignments == null);
+
+ return partitionReassignments;
+ }
+
  public static List<Integer> replicaIdentifiers(Set<BrokerMetadata> brokers) {
  if (brokers == null || brokers.size() == 0) {
  throw new IllegalArgumentException("brokers are either null or empty.");

src/main/java/com/linkedin/xinfra/monitor/services/ConsumeService.java

@@ -11,12 +11,12 @@
 package com.linkedin.xinfra.monitor.services;
 
 import com.linkedin.xinfra.monitor.common.DefaultTopicSchema;
-import com.linkedin.xinfra.monitor.services.metrics.CommitAvailabilityMetrics;
-import com.linkedin.xinfra.monitor.services.metrics.CommitLatencyMetrics;
-import com.linkedin.xinfra.monitor.services.metrics.ConsumeMetrics;
 import com.linkedin.xinfra.monitor.common.Utils;
 import com.linkedin.xinfra.monitor.consumer.BaseConsumerRecord;
 import com.linkedin.xinfra.monitor.consumer.KMBaseConsumer;
+import com.linkedin.xinfra.monitor.services.metrics.CommitAvailabilityMetrics;
+import com.linkedin.xinfra.monitor.services.metrics.CommitLatencyMetrics;
+import com.linkedin.xinfra.monitor.services.metrics.ConsumeMetrics;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
@@ -40,7 +40,7 @@
 import org.apache.kafka.common.metrics.Metrics;
 import org.apache.kafka.common.metrics.MetricsReporter;
 import org.apache.kafka.common.metrics.Sensor;
-import org.apache.kafka.common.metrics.stats.Total;
+import org.apache.kafka.common.metrics.stats.CumulativeSum;
 import org.apache.kafka.common.utils.SystemTime;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -244,7 +244,8 @@ public synchronized void start() {
  @SuppressWarnings("ConstantConditions")
  double partitionCount = topicDescription.partitions().size();
  topicPartitionCount.add(
- new MetricName("topic-partitions-count", METRIC_GROUP_NAME, "The total number of partitions for the topic.", tags), new Total(partitionCount));
+ new MetricName("topic-partitions-count", METRIC_GROUP_NAME, "The total number of partitions for the topic.",
+ tags), new CumulativeSum(partitionCount));
  }
  }
 

src/main/java/com/linkedin/xinfra/monitor/services/MultiClusterTopicManagementService.java

@@ -11,10 +11,10 @@
 package com.linkedin.xinfra.monitor.services;
 
 import com.linkedin.xinfra.monitor.common.Utils;
-import com.linkedin.xinfra.monitor.services.configs.MultiClusterTopicManagementServiceConfig;
-import com.linkedin.xinfra.monitor.topicfactory.TopicFactory;
 import com.linkedin.xinfra.monitor.services.configs.CommonServiceConfig;
+import com.linkedin.xinfra.monitor.services.configs.MultiClusterTopicManagementServiceConfig;
 import com.linkedin.xinfra.monitor.services.configs.TopicManagementServiceConfig;
+import com.linkedin.xinfra.monitor.topicfactory.TopicFactory;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -23,6 +23,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Properties;
 import java.util.Random;
 import java.util.Set;
@@ -31,15 +32,18 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
 import kafka.admin.AdminUtils;
 import kafka.admin.BrokerMetadata;
 import kafka.server.ConfigType;
 import kafka.zk.KafkaZkClient;
 import org.apache.kafka.clients.admin.AdminClient;
 import org.apache.kafka.clients.admin.AdminClientConfig;
+import org.apache.kafka.clients.admin.AlterPartitionReassignmentsResult;
 import org.apache.kafka.clients.admin.CreatePartitionsResult;
 import org.apache.kafka.clients.admin.ElectPreferredLeadersResult;
+import org.apache.kafka.clients.admin.NewPartitionReassignment;
 import org.apache.kafka.clients.admin.NewPartitions;
 import org.apache.kafka.clients.admin.NewTopic;
 import org.apache.kafka.clients.admin.TopicDescription;
@@ -54,6 +58,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import scala.Option$;
+import scala.collection.JavaConverters;
 import scala.collection.Seq;
 
 
@@ -184,7 +189,7 @@ public void run() {
  TopicManagementHelper helper = entry.getValue();
  try {
  helper.maybeReassignPartitionAndElectLeader();
- } catch (IOException | KafkaException e) {
+ } catch (KafkaException e) {
  LOGGER.warn(_serviceName + "/MultiClusterTopicManagementService will retry later in cluster " + clusterName, e);
  }
  }
@@ -405,7 +410,7 @@ private Set<Node> getAvailableBrokers() throws ExecutionException, InterruptedEx
  return brokers;
  }
 
- void maybeReassignPartitionAndElectLeader() throws Exception {
+ void maybeReassignPartitionAndElectLeader() throws ExecutionException, InterruptedException, TimeoutException {
  try (KafkaZkClient zkClient = KafkaZkClient.apply(_zkConnect, JaasUtils.isZkSecurityEnabled(),
  Utils.ZK_SESSION_TIMEOUT_MS, Utils.ZK_CONNECTION_TIMEOUT_MS,
  Integer.MAX_VALUE, Time.SYSTEM, METRIC_GROUP_NAME, "SessionExpireListener", null)) {
@@ -426,13 +431,12 @@ void maybeReassignPartitionAndElectLeader() throws Exception {
  "Configured replication factor {} is smaller than the current replication factor {} of the topic {} in cluster.",
  _replicationFactor, currentReplicationFactor, _topic);
 
- if (expectedReplicationFactor > currentReplicationFactor && !zkClient
- .reassignPartitionsInProgress()) {
+ if (expectedReplicationFactor > currentReplicationFactor
+ && Utils.ongoingPartitionReassignments(_adminClient).isEmpty()) {
  LOGGER.info(
  "MultiClusterTopicManagementService will increase the replication factor of the topic {} in cluster"
  + "from {} to {}", _topic, currentReplicationFactor, expectedReplicationFactor);
- reassignPartitions(zkClient, brokers, _topic, partitionInfoList.size(),
- expectedReplicationFactor);
+ reassignPartitions(_adminClient, brokers, _topic, partitionInfoList.size(), expectedReplicationFactor);
  partitionReassigned = true;
  }
 
@@ -450,19 +454,16 @@ void maybeReassignPartitionAndElectLeader() throws Exception {
  zkClient.setOrCreateEntityConfigs(ConfigType.Topic(), _topic, expectedProperties);
  }
 
- if (partitionInfoList.size() >= brokers.size() &&
- someBrokerNotPreferredLeader(partitionInfoList, brokers) && !zkClient
- .reassignPartitionsInProgress()) {
- LOGGER.info("{} will reassign partitions of the topic {} in cluster.",
- this.getClass().toString(), _topic);
- reassignPartitions(zkClient, brokers, _topic, partitionInfoList.size(),
- expectedReplicationFactor);
+ if (partitionInfoList.size() >= brokers.size() && someBrokerNotPreferredLeader(partitionInfoList, brokers)
+ && Utils.ongoingPartitionReassignments(_adminClient).isEmpty()) {
+ LOGGER.info("{} will reassign partitions of the topic {} in cluster.", this.getClass().toString(), _topic);
+ reassignPartitions(_adminClient, brokers, _topic, partitionInfoList.size(), expectedReplicationFactor);
  partitionReassigned = true;
  }
 
  if (partitionInfoList.size() >= brokers.size() &&
  someBrokerNotElectedLeader(partitionInfoList, brokers)) {
- if (!partitionReassigned || !zkClient.reassignPartitionsInProgress()) {
+ if (!partitionReassigned || Utils.ongoingPartitionReassignments(_adminClient).isEmpty()) {
  LOGGER.info(
  "MultiClusterTopicManagementService will trigger preferred leader election for the topic {} in "
  + "cluster.", _topic
@@ -483,7 +484,7 @@ void maybeElectLeader() throws Exception {
 
  try (KafkaZkClient zkClient = KafkaZkClient.apply(_zkConnect, JaasUtils.isZkSecurityEnabled(), Utils.ZK_SESSION_TIMEOUT_MS,
  Utils.ZK_CONNECTION_TIMEOUT_MS, Integer.MAX_VALUE, Time.SYSTEM, METRIC_GROUP_NAME, "SessionExpireListener", null)) {
- if (!zkClient.reassignPartitionsInProgress()) {
+ if (Utils.ongoingPartitionReassignments(_adminClient).isEmpty()) {
  List<TopicPartitionInfo> partitionInfoList = _adminClient
  .describeTopics(Collections.singleton(_topic)).all().get().get(_topic).partitions();
  LOGGER.info(
@@ -495,8 +496,7 @@ void maybeElectLeader() throws Exception {
  }
  }
 
- private void triggerPreferredLeaderElection(List<TopicPartitionInfo> partitionInfoList, String partitionTopic)
- throws ExecutionException, InterruptedException {
+ private void triggerPreferredLeaderElection(List<TopicPartitionInfo> partitionInfoList, String partitionTopic) {
  Collection<TopicPartition> partitions = new HashSet<>();
  for (TopicPartitionInfo javaPartitionInfo : partitionInfoList) {
  partitions.add(new TopicPartition(partitionTopic, javaPartitionInfo.partition()));
@@ -506,37 +506,54 @@ private void triggerPreferredLeaderElection(List<TopicPartitionInfo> partitionIn
  LOGGER.info("{}: triggerPreferredLeaderElection - {}", this.getClass().toString(), electPreferredLeadersResult.all());
  }
 
- private static void reassignPartitions(KafkaZkClient zkClient, Collection<Node> brokers, String topic, int partitionCount, int replicationFactor) {
- scala.collection.mutable.ArrayBuffer<BrokerMetadata> brokersMetadata = new scala.collection.mutable.ArrayBuffer<>(brokers.size());
+ private static void reassignPartitions(AdminClient adminClient, Collection<Node> brokers, String topic,
+ int partitionCount, int replicationFactor) {
+
+ scala.collection.mutable.ArrayBuffer<BrokerMetadata> brokersMetadata =
+ new scala.collection.mutable.ArrayBuffer<>(brokers.size());
  for (Node broker : brokers) {
  brokersMetadata.$plus$eq(new BrokerMetadata(broker.id(), Option$.MODULE$.apply(broker.rack())));
  }
  scala.collection.Map<Object, Seq<Object>> assignedReplicas =
  AdminUtils.assignReplicasToBrokers(brokersMetadata, partitionCount, replicationFactor, 0, 0);
-
- scala.collection.immutable.Map<TopicPartition, Seq<Object>> newAssignment = new scala.collection.immutable.HashMap<>();
+ scala.collection.immutable.Map<TopicPartition, Seq<Object>> newAssignment =
+   new scala.collection.immutable.HashMap<>();
  scala.collection.Iterator<scala.Tuple2<Object, scala.collection.Seq<Object>>> it = assignedReplicas.iterator();
  while (it.hasNext()) {
  scala.Tuple2<Object, scala.collection.Seq<Object>> scalaTuple = it.next();
  TopicPartition tp = new TopicPartition(topic, (Integer) scalaTuple._1);
  newAssignment = newAssignment.$plus(new scala.Tuple2<>(tp, scalaTuple._2));
  }
 
- scala.collection.immutable.Set<String> topicList = new scala.collection.immutable.Set.Set1<>(topic);
- scala.collection.Map<Object, scala.collection.Seq<Object>> currentAssignment =
- zkClient.getPartitionAssignmentForTopics(topicList).apply(topic);
- String currentAssignmentJson = formatAsNewReassignmentJson(topic, currentAssignment);
  String newAssignmentJson = formatAsNewReassignmentJson(topic, assignedReplicas);
-
  LOGGER.info("Reassign partitions for topic " + topic);
- LOGGER.info("Current partition replica assignment " + currentAssignmentJson);
- LOGGER.info("New partition replica assignment " + newAssignmentJson);
- zkClient.createPartitionReassignment(newAssignment);
+ LOGGER.info("New topic partition replica assignments: {}", newAssignmentJson);
+
+ Set<Map.Entry<TopicPartition, Seq<Object>>> newAssignmentMap =
+ scala.collection.JavaConverters.mapAsJavaMap(newAssignment).entrySet();
+ Map<TopicPartition, Optional<NewPartitionReassignment>> reassignments = new HashMap<>();
+ for (Map.Entry<TopicPartition, Seq<Object>> topicPartitionSeqEntry : newAssignmentMap) {
+ List<Integer> targetReplicas = new ArrayList<>();
+ List<Object> replicas = JavaConverters.seqAsJavaList(topicPartitionSeqEntry.getValue());
+ for (Object replica : replicas) {
+ targetReplicas.add((int) replica);
+ }
+ NewPartitionReassignment newPartitionReassignment = new NewPartitionReassignment(targetReplicas);
+ reassignments.put(topicPartitionSeqEntry.getKey(), Optional.of(newPartitionReassignment));
+ }
+ AlterPartitionReassignmentsResult alterPartitionReassignmentsResult =
+ adminClient.alterPartitionReassignments(reassignments);
+ try {
+ alterPartitionReassignmentsResult.all().get();
+ } catch (InterruptedException | ExecutionException e) {
+ LOGGER.error("An exception occurred while altering the partition reassignments for {}", topic, e);
+ }
  }
 
  static int getReplicationFactor(List<TopicPartitionInfo> partitionInfoList) {
- if (partitionInfoList.isEmpty())
+ if (partitionInfoList.isEmpty()) {
  throw new RuntimeException("Partition list is empty.");
+ }
 
  int replicationFactor = partitionInfoList.get(0).replicas().size();
  for (TopicPartitionInfo partitionInfo : partitionInfoList) {