NIFI-15226 Review suggestions

Author: tpalfy

nifi-extension-bundles/nifi-kafka-bundle/nifi-kafka-processors/src/main/java/org/apache/nifi/kafka/processors/ConsumeKafka.java

@@ -73,13 +73,13 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Queue;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -343,7 +343,7 @@ public class ConsumeKafka extends AbstractProcessor implements VerifiableProcess
     private final AtomicInteger activeConsumerCount = new AtomicInteger();
 
     private final Queue<PollingContext> availablePartitionedPollingContexts = new LinkedBlockingQueue<>();
-    private final Map<KafkaConsumerService, PollingContext> consumerServiceToPartitionedPollingContext = new HashMap<>();
+    private final Map<KafkaConsumerService, PollingContext> consumerServiceToPartitionedPollingContext = new ConcurrentHashMap<>();
 
     @Override
     public List<PropertyDescriptor> getSupportedPropertyDescriptors() {
@@ -445,11 +445,6 @@ public void onScheduled(final ProcessContext context) {
             if (partitionCount != numAssignedPartitions) {
                 context.yield();
 
-                KafkaConsumerService service;
-                while ((service = consumerServices.poll()) != null) {
-                    close(service, "Not all partitions are assigned");
-                }
-
                 throw new ProcessException("Illegal Partition Assignment: There are "
                         + numAssignedPartitions + " partitions statically assigned using the " + PARTITIONS_PROPERTY_PREFIX + ".* property names,"
                         + " but the Kafka topic(s) have " + partitionCount + " partitions");
@@ -463,9 +458,9 @@ public void onScheduled(final ProcessContext context) {
             }
 
             for (int partition : assignedPartitions) {
-                PollingContext partitionedPollingContext = createPollingContext(context, partition);
-
-                availablePartitionedPollingContexts.add(partitionedPollingContext);
+                final PollingContext partitionedPollingContext = createPollingContext(context, partition);
+                final KafkaConsumerService partitionedConsumerService = obtainPartitionedConsumerService(partitionedPollingContext);
+                consumerServices.add(partitionedConsumerService);
             }
         }
     }
@@ -658,9 +653,7 @@ private KafkaConsumerService getConsumerService(final ProcessContext context) {
 
             getLogger().info("No Partitioned Kafka Consumer Service available; creating a new one.");
 
-            final KafkaConsumerService partitionedConsumerService = connectionService.getConsumerService(partitionedPollingContext);
-
-            consumerServiceToPartitionedPollingContext.put(partitionedConsumerService, partitionedPollingContext);
+            final KafkaConsumerService partitionedConsumerService = obtainPartitionedConsumerService(partitionedPollingContext);
 
             return partitionedConsumerService;
         } else {
@@ -701,6 +694,14 @@ private int getPartitionCount(final KafkaConnectionService connectionService) {
         return partitionsEachTopic;
     }
 
+    private KafkaConsumerService obtainPartitionedConsumerService(PollingContext partitionedPollingContext) {
+        final KafkaConsumerService partitionedConsumerService = connectionService.getConsumerService(partitionedPollingContext);
+
+        consumerServiceToPartitionedPollingContext.put(partitionedConsumerService, partitionedPollingContext);
+
+        return partitionedConsumerService;
+    }
+
     private int getMaxConsumerCount() {
         return maxConsumerCount;
     }

nifi-extension-bundles/nifi-kafka-bundle/nifi-kafka-service-shared/src/main/java/org/apache/nifi/kafka/service/Kafka3ConnectionService.java

@@ -240,14 +240,6 @@ public KafkaConsumerService getConsumerService(final PollingContext pollingConte
         final ByteArrayDeserializer deserializer = new ByteArrayDeserializer();
         final Consumer<byte[], byte[]> consumer = new KafkaConsumer<>(properties, deserializer, deserializer);
 
-        Integer partition = pollingContext.getPartition();
-        if (partition != null) {
-            List<TopicPartition> topicPartitions = pollingContext.getTopics().stream()
-                    .map(topic -> new TopicPartition(topic, partition))
-                    .collect(Collectors.toList());
-            consumer.assign(topicPartitions);
-        }
-
         return new Kafka3ConsumerService(getLogger(), consumer, subscription);
     }
 
@@ -258,7 +250,7 @@ private Subscription createSubscription(final PollingContext pollingContext) {
 
         return topicPatternFound
             .map(pattern -> new Subscription(groupId, pattern, autoOffsetReset))
-            .orElseGet(() -> new Subscription(groupId, pollingContext.getTopics(), autoOffsetReset));
+            .orElseGet(() -> new Subscription(groupId, pollingContext.getPartition(), pollingContext.getTopics(), autoOffsetReset));
     }
 
     @Override

nifi-extension-bundles/nifi-kafka-bundle/nifi-kafka-service-shared/src/main/java/org/apache/nifi/kafka/service/consumer/Kafka3ConsumerService.java

@@ -62,7 +62,8 @@ public Kafka3ConsumerService(final ComponentLog componentLog, final Consumer<byt
         this.consumer = consumer;
         this.subscription = subscription;
 
-        if (consumer.assignment().isEmpty()) {
+        final Integer partition = subscription.getPartition();
+        if (partition == null) {
             final Optional<Pattern> topicPatternFound = subscription.getTopicPattern();
             if (topicPatternFound.isPresent()) {
                 final Pattern topicPattern = topicPatternFound.get();
@@ -71,6 +72,11 @@ public Kafka3ConsumerService(final ComponentLog componentLog, final Consumer<byt
                 final Collection<String> topics = subscription.getTopics();
                 consumer.subscribe(topics, this);
             }
+        } else {
+            List<TopicPartition> topicPartitions = subscription.getTopics().stream()
+                    .map(topic -> new TopicPartition(topic, partition))
+                    .collect(Collectors.toList());
+            consumer.assign(topicPartitions);
         }
     }
 

nifi-extension-bundles/nifi-kafka-bundle/nifi-kafka-service-shared/src/main/java/org/apache/nifi/kafka/service/consumer/Subscription.java

@@ -31,19 +31,22 @@ public class Subscription {
 
     private final String groupId;
     private final Collection<String> topics;
+    private final Integer partition;
     private final Pattern topicPattern;
     private final AutoOffsetReset autoOffsetReset;
 
-    public Subscription(final String groupId, final Collection<String> topics, final AutoOffsetReset autoOffsetReset) {
+    public Subscription(final String groupId, final Integer partition, final Collection<String> topics, final AutoOffsetReset autoOffsetReset) {
         this.groupId = Objects.requireNonNull(groupId, "Group ID required");
         this.topics = Collections.unmodifiableCollection(Objects.requireNonNull(topics, "Topics required"));
+        this.partition = partition;
         this.topicPattern = null;
         this.autoOffsetReset = Objects.requireNonNull(autoOffsetReset, "Auto Offset Reset required");
     }
 
     public Subscription(final String groupId, final Pattern topicPattern, final AutoOffsetReset autoOffsetReset) {
         this.groupId = Objects.requireNonNull(groupId, "Group ID required");
         this.topics = Collections.emptyList();
+        this.partition = null;
         this.topicPattern = Objects.requireNonNull(topicPattern, "Topic Pattern required");
         this.autoOffsetReset = Objects.requireNonNull(autoOffsetReset, "Auto Offset Reset required");
     }
@@ -56,6 +59,10 @@ public Collection<String> getTopics() {
         return topics;
     }
 
+    public Integer getPartition() {
+        return partition;
+    }
+
     public Optional<Pattern> getTopicPattern() {
         return Optional.ofNullable(topicPattern);
     }
@@ -94,7 +101,11 @@ public String toString() {
     }
 
     private boolean isTopicSubscriptionMatched(final Subscription subscription) {
-        if (topics.size() == subscription.topics.size() && topics.containsAll(subscription.topics)) {
+        if (
+                topics.size() == subscription.topics.size()
+                        && topics.containsAll(subscription.topics)
+                        && Objects.equals(this.partition, subscription.partition)
+        ) {
             final String regexLeft = (topicPattern == null ? null : topicPattern.pattern());
             final String regexRight = (subscription.topicPattern == null ? null : subscription.topicPattern.pattern());
             return Objects.equals(regexLeft, regexRight);