SAMZA-1888: Kafka consumer improvements

Author: Boris S <bshkolnik@linkedin.com>
Author: Boris S <boryas@apache.org>
Author: Boris Shkolnik <bshkolni@linkedin.com>

Reviewers: bharathkk <codin.martial@gmail.com>

Closes apache#738 from sborya/KafkaConsumerImprovements

Author: Boris S

samza-kafka/src/main/java/org/apache/samza/system/kafka/KafkaSystemAdmin.java

@@ -80,9 +80,9 @@ public class KafkaSystemAdmin implements ExtendedSystemAdmin {
 
   protected final String systemName;
   protected final Consumer metadataConsumer;
+  protected final Config  config;
 
-  // get ZkUtils object to connect to Kafka's ZK.
-  private final Supplier<ZkUtils> getZkConnection;
+  protected AdminClient adminClient = null;
 
   // Custom properties to create a new coordinator stream.
   private final Properties coordinatorStreamProperties;
@@ -96,52 +96,21 @@ public class KafkaSystemAdmin implements ExtendedSystemAdmin {
   // Kafka properties for intermediate topics creation
   private final Map<String, Properties> intermediateStreamProperties;
 
-  // adminClient is required for deleteCommittedMessages operation
-  private final AdminClient adminClient;
-
   // used for intermediate streams
-  private final boolean deleteCommittedMessages;
+  protected final boolean deleteCommittedMessages;
 
   private final AtomicBoolean stopped = new AtomicBoolean(false);
 
   public KafkaSystemAdmin(String systemName, Config config, Consumer metadataConsumer) {
     this.systemName = systemName;
+    this.config = config;
 
     if (metadataConsumer == null) {
       throw new SamzaException(
           "Cannot construct KafkaSystemAdmin for system " + systemName + " with null metadataConsumer");
     }
     this.metadataConsumer = metadataConsumer;
 
-    // populate brokerList from either consumer or producer configs
-    Properties props = new Properties();
-    String brokerList = config.get(
-        String.format(KafkaConfig.CONSUMER_CONFIGS_CONFIG_KEY(), systemName, ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG));
-    if (brokerList == null) {
-      brokerList = config.get(String.format(KafkaConfig.PRODUCER_CONFIGS_CONFIG_KEY(), systemName,
-          ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG));
-    }
-    if (brokerList == null) {
-      throw new SamzaException(
-          ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG + " is required for systemAdmin for system " + systemName);
-    }
-    props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList);
-
-    // kafka.admin.AdminUtils requires zkConnect
-    // this will change after we move to the new org.apache..AdminClient
-    String zkConnect =
-        config.get(String.format(KafkaConfig.CONSUMER_CONFIGS_CONFIG_KEY(), systemName, ZOOKEEPER_CONNECT));
-    if (StringUtils.isBlank(zkConnect)) {
-      throw new SamzaException("Missing zookeeper.connect config for admin for system " + systemName);
-    }
-    props.put(ZOOKEEPER_CONNECT, zkConnect);
-
-    adminClient = AdminClient.create(props);
-
-    getZkConnection = () -> {
-      return ZkUtils.apply(zkConnect, 6000, 6000, false);
-    };
-
     KafkaConfig kafkaConfig = new KafkaConfig(config);
     coordinatorStreamReplicationFactor = Integer.valueOf(kafkaConfig.getCoordinatorReplicationFactor());
     coordinatorStreamProperties = KafkaSystemAdminUtilsScala.getCoordinatorTopicProperties(kafkaConfig);
@@ -197,6 +166,8 @@ public void stop() {
       } catch (Exception e) {
         LOG.warn("metadataConsumer.close for system " + systemName + " failed with exception.", e);
       }
+    }
+    if (adminClient != null) {
       try {
         adminClient.close();
       } catch (Exception e) {
@@ -546,14 +517,14 @@ public Integer offsetComparator(String offset1, String offset2) {
   public boolean createStream(StreamSpec streamSpec) {
     LOG.info("Creating Kafka topic: {} on system: {}", streamSpec.getPhysicalName(), streamSpec.getSystemName());
 
-    return KafkaSystemAdminUtilsScala.createStream(toKafkaSpec(streamSpec), getZkConnection);
+    return KafkaSystemAdminUtilsScala.createStream(toKafkaSpec(streamSpec), getZkConnection());
   }
 
   @Override
   public boolean clearStream(StreamSpec streamSpec) {
     LOG.info("Creating Kafka topic: {} on system: {}", streamSpec.getPhysicalName(), streamSpec.getSystemName());
 
-    KafkaSystemAdminUtilsScala.clearStream(streamSpec, getZkConnection);
+    KafkaSystemAdminUtilsScala.clearStream(streamSpec, getZkConnection());
 
     Map<String, List<PartitionInfo>> topicsMetadata = getTopicMetadata(ImmutableSet.of(streamSpec.getPhysicalName()));
     return topicsMetadata.get(streamSpec.getPhysicalName()).isEmpty();
@@ -630,11 +601,56 @@ Map<String, List<PartitionInfo>> getTopicMetadata(Set<String> topics) {
   @Override
   public void deleteMessages(Map<SystemStreamPartition, String> offsets) {
     if (deleteCommittedMessages) {
+      if (adminClient == null) {
+        adminClient = AdminClient.create(createAdminClientProperties());
+      }
       KafkaSystemAdminUtilsScala.deleteMessages(adminClient, offsets);
       deleteMessageCalled = true;
     }
   }
 
+  protected Properties createAdminClientProperties() {
+    // populate brokerList from either consumer or producer configs
+    Properties props = new Properties();
+    // included SSL settings if needed
+
+    props.putAll(config.subset(String.format("systems.%s.consumer.", systemName), true));
+
+    //validate brokerList
+    String brokerList = config.get(
+        String.format(KafkaConfig.CONSUMER_CONFIGS_CONFIG_KEY(), systemName, ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG));
+    if (brokerList == null) {
+      brokerList = config.get(String.format(KafkaConfig.PRODUCER_CONFIGS_CONFIG_KEY(), systemName,
+          ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG));
+    }
+    if (brokerList == null) {
+      throw new SamzaException(
+          ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG + " is required for systemAdmin for system " + systemName);
+    }
+    props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList);
+
+
+    // kafka.admin.AdminUtils requires zkConnect
+    // this will change after we move to the new org.apache..AdminClient
+    String zkConnect =
+        config.get(String.format(KafkaConfig.CONSUMER_CONFIGS_CONFIG_KEY(), systemName, ZOOKEEPER_CONNECT));
+    if (StringUtils.isBlank(zkConnect)) {
+      throw new SamzaException("Missing zookeeper.connect config for admin for system " + systemName);
+    }
+    props.put(ZOOKEEPER_CONNECT, zkConnect);
+
+    return props;
+  }
+
+  private Supplier<ZkUtils> getZkConnection() {
+    String zkConnect =
+        config.get(String.format(KafkaConfig.CONSUMER_CONFIGS_CONFIG_KEY(), systemName, ZOOKEEPER_CONNECT));
+    if (StringUtils.isBlank(zkConnect)) {
+      throw new SamzaException("Missing zookeeper.connect config for admin for system " + systemName);
+    }
+    return () -> ZkUtils.apply(zkConnect, 6000, 6000, false);
+  }
+
   /**
    * Container for metadata about offsets.
    */

samza-kafka/src/main/java/org/apache/samza/system/kafka/KafkaSystemConsumer.java

@@ -77,8 +77,10 @@ public class KafkaSystemConsumer<K, V> extends BlockingEnvelopeMap implements Sy
 
   /**
    * Create a KafkaSystemConsumer for the provided {@code systemName}
+   * @param kafkaConsumer kafka Consumer object to be used by this system consumer
    * @param systemName system name for which we create the consumer
    * @param config application config
+   * @param clientId clientId from the kafka consumer to be used in the KafkaConsumerProxy
    * @param metrics metrics for this KafkaSystemConsumer
    * @param clock system clock
    */
@@ -106,12 +108,13 @@ public KafkaSystemConsumer(Consumer<K, V> kafkaConsumer, String systemName, Conf
 
   /**
    * Create internal kafka consumer object, which will be used in the Proxy.
+   * @param <K> key type for the consumer
+   * @param <V> value type for the consumer
    * @param systemName system name for which we create the consumer
    * @param kafkaConsumerConfig config object for Kafka's KafkaConsumer
-   * @return KafkaConsumer object
+   * @return KafkaConsumer newly created kafka consumer object
    */
-  public static <K,V> KafkaConsumer<K, V> createKafkaConsumerImpl(String systemName,
-      HashMap<String, Object> kafkaConsumerConfig) {
+  public static <K, V> KafkaConsumer<K, V> createKafkaConsumerImpl(String systemName, HashMap<String, Object> kafkaConsumerConfig) {
 
     LOG.info("Instantiating KafkaConsumer for systemName {} with properties {}", systemName, kafkaConsumerConfig);
     return new KafkaConsumer<>(kafkaConsumerConfig);
@@ -176,7 +179,7 @@ void startConsumer() {
         throw new SamzaException(msg, e);
       }
 
-      LOG.info("{}: Changing consumer's starting offset for tp = %s to %s", this, tp, startingOffsetString);
+      LOG.info("{}: Changing consumer's starting offset for tp = {} to {}", this, tp, startingOffsetString);
 
       // add the partition to the proxy
       proxy.addTopicPartition(topicPartitionsToSSP.get(tp), startingOffset);
@@ -310,16 +313,10 @@ public Map<SystemStreamPartition, List<IncomingMessageEnvelope>> poll(
     return super.poll(systemStreamPartitions, timeout);
   }
 
-  /**
-   * convert from TopicPartition to TopicAndPartition
-   */
   public static TopicAndPartition toTopicAndPartition(TopicPartition tp) {
     return new TopicAndPartition(tp.topic(), tp.partition());
   }
 
-  /**
-   * convert to TopicPartition from SystemStreamPartition
-   */
   public static TopicPartition toTopicPartition(SystemStreamPartition ssp) {
     return new TopicPartition(ssp.getStream(), ssp.getPartition().getPartitionId());
   }