Refactor and clean up (#339)

1. Add logging to the MultiClusterTopicManagementService class
2. Add a timeout to some admin client requests
3. Rename some variables with non-inclusive language

Author: Lincong Li

src/main/java/com/linkedin/xinfra/monitor/XinfraMonitor.java

@@ -65,22 +65,22 @@ public XinfraMonitor(Map<String, Map> allClusterProps) throws Exception {
     _services = new ConcurrentHashMap<>();
 
     for (Map.Entry<String, Map> clusterProperty : allClusterProps.entrySet()) {
-      String name = clusterProperty.getKey();
+      String clusterName = clusterProperty.getKey();
       Map props = clusterProperty.getValue();
       if (!props.containsKey(XinfraMonitorConstants.CLASS_NAME_CONFIG))
-        throw new IllegalArgumentException(name + " is not configured with " + XinfraMonitorConstants.CLASS_NAME_CONFIG);
+        throw new IllegalArgumentException(clusterName + " is not configured with " + XinfraMonitorConstants.CLASS_NAME_CONFIG);
       String className = (String) props.get(XinfraMonitorConstants.CLASS_NAME_CONFIG);
 
       Class<?> aClass = Class.forName(className);
       if (App.class.isAssignableFrom(aClass)) {
-        App clusterApp = (App) Class.forName(className).getConstructor(Map.class, String.class).newInstance(props, name);
-        _apps.put(name, clusterApp);
+        App clusterApp = (App) Class.forName(className).getConstructor(Map.class, String.class).newInstance(props, clusterName);
+        _apps.put(clusterName, clusterApp);
       } else if (Service.class.isAssignableFrom(aClass)) {
         ServiceFactory serviceFactory = (ServiceFactory) Class.forName(className + XinfraMonitorConstants.FACTORY)
             .getConstructor(Map.class, String.class)
-            .newInstance(props, name);
+            .newInstance(props, clusterName);
         Service service = serviceFactory.createService();
-        _services.put(name, service);
+        _services.put(clusterName, service);
       } else {
         throw new IllegalArgumentException(className + " should implement either " + App.class.getSimpleName() + " or " + Service.class.getSimpleName());
       }

src/main/java/com/linkedin/xinfra/monitor/apps/SingleClusterMonitor.java

@@ -55,21 +55,22 @@ public class SingleClusterMonitor implements App {
 
   private static final int SERVICES_INITIAL_CAPACITY = 4;
   private final TopicManagementService _topicManagementService;
-  private final String _name;
+  private final String _clusterName;
   private final List<Service> _allServices;
   private final boolean _isTopicManagementServiceEnabled;
 
-  public SingleClusterMonitor(Map<String, Object> props, String name) throws Exception {
+  public SingleClusterMonitor(Map<String, Object> props, String clusterName) throws Exception {
     ConsumerFactory consumerFactory = new ConsumerFactoryImpl(props);
-    _name = name;
+    _clusterName = clusterName;
     LOG.info("SingleClusterMonitor properties: {}", prettyPrint(props));
     TopicManagementServiceConfig config = new TopicManagementServiceConfig(props);
     _isTopicManagementServiceEnabled =
         config.getBoolean(TopicManagementServiceConfig.TOPIC_MANAGEMENT_ENABLED_CONFIG);
     _allServices = new ArrayList<>(SERVICES_INITIAL_CAPACITY);
     CompletableFuture<Void> topicPartitionResult;
     if (_isTopicManagementServiceEnabled) {
-      _topicManagementService = new TopicManagementService(props, name);
+      String topicManagementServiceName = String.format("Topic-management-service-for-%s", clusterName);
+      _topicManagementService = new TopicManagementService(props, topicManagementServiceName);
       topicPartitionResult = _topicManagementService.topicPartitionResult();
 
       // block on the MultiClusterTopicManagementService to complete.
@@ -80,10 +81,9 @@ public SingleClusterMonitor(Map<String, Object> props, String name) throws Excep
       _topicManagementService = null;
       topicPartitionResult = new CompletableFuture<>();
       topicPartitionResult.complete(null);
-
     }
-    ProduceService produceService = new ProduceService(props, name);
-    ConsumeService consumeService = new ConsumeService(name, topicPartitionResult, consumerFactory);
+    ProduceService produceService = new ProduceService(props, clusterName);
+    ConsumeService consumeService = new ConsumeService(clusterName, topicPartitionResult, consumerFactory);
     _allServices.add(produceService);
     _allServices.add(consumeService);
   }
@@ -126,15 +126,15 @@ public void start() throws Exception {
       }
     }
 
-    LOG.info(_name + "/SingleClusterMonitor started!");
+    LOG.info(_clusterName + "/SingleClusterMonitor started!");
   }
 
   @Override
   public void stop() {
     for (Service service : _allServices) {
       service.stop();
     }
-    LOG.info(_name + "/SingleClusterMonitor stopped.");
+    LOG.info(_clusterName + "/SingleClusterMonitor stopped.");
   }
 
   @Override

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

@@ -164,14 +164,14 @@ private void createDeleteClusterTopic() {
       try {
         int brokerCount = _adminClient.describeCluster().nodes().get().size();
 
-        Set<Integer> blackListedBrokers = _topicFactory.getBlackListedBrokers(_adminClient);
         Set<BrokerMetadata> brokers = new HashSet<>();
         for (Node broker : _adminClient.describeCluster().nodes().get()) {
           BrokerMetadata brokerMetadata = new BrokerMetadata(broker.id(), null);
           brokers.add(brokerMetadata);
         }
-        if (!blackListedBrokers.isEmpty()) {
-          brokers.removeIf(broker -> blackListedBrokers.contains(broker.id()));
+        Set<Integer> excludedBrokers = _topicFactory.getExcludedBrokers(_adminClient);
+        if (!excludedBrokers.isEmpty()) {
+          brokers.removeIf(broker -> excludedBrokers.contains(broker.id()));
         }
 
         // map from partition id to replica ids (i.e. broker ids).

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

@@ -15,6 +15,7 @@
 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.time.Duration;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -26,6 +27,7 @@
 import java.util.Properties;
 import java.util.Random;
 import java.util.Set;
+import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executors;
@@ -41,7 +43,6 @@
 import org.apache.kafka.clients.admin.AlterPartitionReassignmentsResult;
 import org.apache.kafka.clients.admin.Config;
 import org.apache.kafka.clients.admin.ConfigEntry;
-import org.apache.kafka.clients.admin.CreatePartitionsResult;
 import org.apache.kafka.clients.admin.ElectLeadersResult;
 import org.apache.kafka.clients.admin.NewPartitionReassignment;
 import org.apache.kafka.clients.admin.NewPartitions;
@@ -248,7 +249,7 @@ static class TopicManagementHelper {
     private final int _minPartitionNum;
     private final Properties _topicProperties;
     private boolean _preferredLeaderElectionRequested;
-    private final int _requestTimeoutMs;
+    private final Duration _requestTimeout;
     private final List<String> _bootstrapServers;
 
     // package private for unit testing
@@ -261,6 +262,7 @@ static class TopicManagementHelper {
 
     @SuppressWarnings("unchecked")
     TopicManagementHelper(Map<String, Object> props) throws Exception {
+
       TopicManagementServiceConfig config = new TopicManagementServiceConfig(props);
       AdminClientConfig adminClientConfig = new AdminClientConfig(props);
       String topicFactoryClassName = config.getString(TopicManagementServiceConfig.TOPIC_FACTORY_CLASS_CONFIG);
@@ -273,7 +275,7 @@ static class TopicManagementHelper {
       _minPartitionsToBrokersRatio = config.getDouble(TopicManagementServiceConfig.PARTITIONS_TO_BROKERS_RATIO_CONFIG);
       _minPartitionNum = config.getInt(TopicManagementServiceConfig.MIN_PARTITION_NUM_CONFIG);
       _preferredLeaderElectionRequested = false;
-      _requestTimeoutMs = adminClientConfig.getInt(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG);
+      _requestTimeout = Duration.ofMillis(adminClientConfig.getInt(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG));
       _bootstrapServers = adminClientConfig.getList(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG);
       _topicProperties = new Properties();
       if (props.containsKey(TopicManagementServiceConfig.TOPIC_PROPS_CONFIG)) {
@@ -288,9 +290,17 @@ static class TopicManagementHelper {
               TopicManagementServiceConfig.TOPIC_FACTORY_PROPS_CONFIG) : new HashMap();
       _topicFactory =
           (TopicFactory) Class.forName(topicFactoryClassName).getConstructor(Map.class).newInstance(topicFactoryConfig);
-
       _adminClient = constructAdminClient(props);
       LOGGER.info("{} configs: {}", _adminClient.getClass().getSimpleName(), props);
+      logConfigurationValues();
+    }
+
+    private void logConfigurationValues() {
+      LOGGER.info("TopicManagementHelper for cluster with Zookeeper connect {} is configured with " +
+              "[topic={}, topicCreationEnabled={}, topicAddPartitionEnabled={}, " +
+              "topicReassignPartitionAndElectLeaderEnabled={}, minPartitionsToBrokersRatio={}, " +
+              "minPartitionNum={}]", _zkConnect, _topic, _topicCreationEnabled, _topicAddPartitionEnabled,
+              _topicReassignPartitionAndElectLeaderEnabled, _minPartitionsToBrokersRatio, _minPartitionNum);
     }
 
     @SuppressWarnings("unchecked")
@@ -317,7 +327,8 @@ int minPartitionNum() throws InterruptedException, ExecutionException {
       return Math.max((int) Math.ceil(_minPartitionsToBrokersRatio * brokerCount), _minPartitionNum);
     }
 
-    void maybeAddPartitions(int minPartitionNum) throws ExecutionException, InterruptedException {
+    void maybeAddPartitions(final int requiredMinPartitionNum)
+            throws ExecutionException, InterruptedException, CancellationException, TimeoutException {
       if (!_topicAddPartitionEnabled) {
         LOGGER.info("Adding partition to {} topic is not enabled in a cluster with Zookeeper URL {}. " +
                 "Refer to config: {}", _topic, _zkConnect, TopicManagementServiceConfig.TOPIC_ADD_PARTITION_ENABLED_CONFIG);
@@ -326,32 +337,36 @@ void maybeAddPartitions(int minPartitionNum) throws ExecutionException, Interrup
       Map<String, KafkaFuture<TopicDescription>> kafkaFutureMap =
           _adminClient.describeTopics(Collections.singleton(_topic)).values();
       KafkaFuture<TopicDescription> topicDescriptions = kafkaFutureMap.get(_topic);
-      List<TopicPartitionInfo> partitions = topicDescriptions.get().partitions();
-
-      int partitionNum = partitions.size();
-      if (partitionNum < minPartitionNum) {
-        LOGGER.info("{} will increase partition of the topic {} in the cluster from {}" + " to {}.",
-            this.getClass().toString(), _topic, partitionNum, minPartitionNum);
-        Set<Integer> blackListedBrokers = _topicFactory.getBlackListedBrokers(_adminClient);
-        Set<BrokerMetadata> brokers = new HashSet<>();
-        for (Node broker : _adminClient.describeCluster().nodes().get()) {
-          BrokerMetadata brokerMetadata = new BrokerMetadata(broker.id(), null);
-          brokers.add(brokerMetadata);
-        }
+      List<TopicPartitionInfo> partitions = topicDescriptions.get(_requestTimeout.toMillis(), TimeUnit.MILLISECONDS).partitions();
 
-        if (!blackListedBrokers.isEmpty()) {
-          brokers.removeIf(broker -> blackListedBrokers.contains(broker.id()));
-        }
+      final int currPartitionNum = partitions.size();
+      if (currPartitionNum >= requiredMinPartitionNum) {
+        LOGGER.debug("{} will not increase partition of the topic {} in the cluster. Current partition count {} and '" +
+                "minimum required partition count is {}.", this.getClass().toString(), _topic, currPartitionNum, requiredMinPartitionNum);
+        return;
+      }
+      LOGGER.info("{} will increase partition of the topic {} in the cluster from {}" + " to {}.",
+              this.getClass().toString(), _topic, currPartitionNum,  requiredMinPartitionNum);
+      Set<BrokerMetadata> brokers = new HashSet<>();
+      for (Node broker : _adminClient.describeCluster().nodes().get(_requestTimeout.toMillis(), TimeUnit.MILLISECONDS)) {
+        BrokerMetadata brokerMetadata = new BrokerMetadata(broker.id(), null);
+        brokers.add(brokerMetadata);
+      }
+      Set<Integer> excludedBrokers = _topicFactory.getExcludedBrokers(_adminClient);
+      if (!excludedBrokers.isEmpty()) {
+        brokers.removeIf(broker ->  excludedBrokers.contains(broker.id()));
+      }
 
-        List<List<Integer>> newPartitionAssignments =
-            newPartitionAssignments(minPartitionNum, partitionNum, brokers, _replicationFactor);
+      List<List<Integer>> newPartitionAssignments =
+              newPartitionAssignments(requiredMinPartitionNum, currPartitionNum, brokers, _replicationFactor);
 
-        NewPartitions newPartitions = NewPartitions.increaseTo(minPartitionNum, newPartitionAssignments);
+      NewPartitions newPartitions = NewPartitions.increaseTo(requiredMinPartitionNum, newPartitionAssignments);
 
-        Map<String, NewPartitions> newPartitionsMap = new HashMap<>();
-        newPartitionsMap.put(_topic, newPartitions);
-        CreatePartitionsResult createPartitionsResult = _adminClient.createPartitions(newPartitionsMap);
-      }
+      Map<String, NewPartitions> newPartitionsMap = new HashMap<>();
+      newPartitionsMap.put(_topic, newPartitions);
+      _adminClient.createPartitions(newPartitionsMap).all().get(_requestTimeout.toMillis(), TimeUnit.MILLISECONDS);
+      LOGGER.info("{} finished increasing partition of the topic {} in the cluster from {} to {}",
+              this.getClass().toString(), _topic, currPartitionNum,  requiredMinPartitionNum);
     }
 
     static List<List<Integer>> newPartitionAssignments(int minPartitionNum, int partitionNum,
@@ -427,8 +442,8 @@ int numPartitions() throws InterruptedException, ExecutionException {
 
     private Set<Node> getAvailableBrokers() throws ExecutionException, InterruptedException {
       Set<Node> brokers = new HashSet<>(_adminClient.describeCluster().nodes().get());
-      Set<Integer> blackListedBrokers = _topicFactory.getBlackListedBrokers(_adminClient);
-      brokers.removeIf(broker -> blackListedBrokers.contains(broker.id()));
+      Set<Integer> excludedBrokers = _topicFactory.getExcludedBrokers(_adminClient);
+      brokers.removeIf(broker -> excludedBrokers.contains(broker.id()));
       return brokers;
     }
 

src/main/java/com/linkedin/xinfra/monitor/topicfactory/DefaultTopicFactory.java

@@ -32,7 +32,7 @@ public int createTopicIfNotExist(String topic, short replicationFactor, double p
   }
 
   @Override
-  public Set<Integer> getBlackListedBrokers(AdminClient adminClient) {
+  public Set<Integer> getExcludedBrokers(AdminClient adminClient) {
     return Collections.emptySet();
   }
 }

src/main/java/com/linkedin/xinfra/monitor/topicfactory/TopicFactory.java

@@ -42,6 +42,6 @@ int createTopicIfNotExist(String topic, short replicationFactor, double partitio
    * @param adminClient AdminClient object
    * @return A set of brokers that don't take new partitions or reassigned partitions for topics.
    */
-  Set<Integer> getBlackListedBrokers(AdminClient adminClient);
+  Set<Integer> getExcludedBrokers(AdminClient adminClient);
 
 }