[GitHub] [kafka] showuon commented on a change in pull request #10552: KAFKA-12675: improve the sticky general assignor scalability and performance

[GitBox]

showuon commented on a change in pull request #10552:
URL: https://github.com/apache/kafka/pull/10552#discussion_r638718719



##########
File path: 
clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractStickyAssignor.java
##########
@@ -637,19 +709,27 @@ private void assignPartition(TopicPartition partition,
         }
     }
 
-    private boolean canParticipateInReassignment(TopicPartition partition,
-                                                 Map<TopicPartition, 
List<String>> partition2AllPotentialConsumers) {
+    private boolean canParticipateInReassignment(String topic,
+                                                 Map<String, List<String>> 
topic2AllPotentialConsumers) {
         // if a partition has two or more potential consumers it is subject to 
reassignment.
-        return partition2AllPotentialConsumers.get(partition).size() >= 2;
+        return topic2AllPotentialConsumers.get(topic).size() >= 2;
     }
 
     private boolean canParticipateInReassignment(String consumer,
                                                  Map<String, 
List<TopicPartition>> currentAssignment,
-                                                 Map<String, 
List<TopicPartition>> consumer2AllPotentialPartitions,
-                                                 Map<TopicPartition, 
List<String>> partition2AllPotentialConsumers) {
+                                                 Map<String, List<String>> 
consumer2AllPotentialTopics,
+                                                 Map<String, List<String>> 
topic2AllPotentialConsumers,
+                                                 Map<String, Integer> 
partitionsPerTopic,
+                                                 int totalPartitionCount) {
         List<TopicPartition> currentPartitions = 
currentAssignment.get(consumer);
         int currentAssignmentSize = currentPartitions.size();
-        int maxAssignmentSize = 
consumer2AllPotentialPartitions.get(consumer).size();
+        List<String> allSubscribedTopics = 
consumer2AllPotentialTopics.get(consumer);
+        int maxAssignmentSize;
+        if (allSubscribedTopics.size() == partitionsPerTopic.size()) {
+            maxAssignmentSize = totalPartitionCount;
+        } else {
+            maxAssignmentSize = allSubscribedTopics.stream().map(topic -> 
partitionsPerTopic.get(topic)).reduce(0, Integer::sum);
+        }

Review comment:
       Good suggestion! Updated. Thanks.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org