Re: [PR] KAFKA-19193 Support rack-aware partitioning for Kafka producer [kafka]

["ivanyu (via GitHub)" <gi...@apache.org>]

ivanyu commented on code in PR #19850:
URL: https://github.com/apache/kafka/pull/19850#discussion_r3026079852


##########
clients/src/main/java/org/apache/kafka/clients/producer/internals/BuiltInPartitioner.java:
##########
@@ -293,18 +324,55 @@ public void updatePartitionLoadStats(int[] queueSizes, int[] partitionIds, int l
             // and we didn't exclude partitions that experience high latencies (greater than
             // partitioner.availability.timeout.ms).
             log.trace("All queue lengths are the same, not using adaptive for topic {}", topic);
-            partitionLoadStats = null;
+            partitionLoadStatsHolder = null;
             return;
         }
 
+        // Before inverting and folding, build fully the load stats for this rack, because this depends on the raw queue sizes.
+        PartitionLoadStats partitionLoadStatsInThisRack = createPartitionLoadStatsForThisRackIfNeeded(queueSizes, partitionIds, partitionLeaderRacks, length);
+
         // Invert and fold the queue size, so that they become separator values in the CFT.
+        invertAndFoldQueueSizeArray(queueSizes, maxSizePlus1, length);
+
+        log.trace("Partition load stats for topic {}: CFT={}, IDs={}, length={}",
+                topic, queueSizes, partitionIds, length);
+        partitionLoadStatsHolder = new PartitionLoadStatsHolder(
+            new PartitionLoadStats(queueSizes, partitionIds, length),
+            partitionLoadStatsInThisRack
+        );
+    }
+
+    private PartitionLoadStats createPartitionLoadStatsForThisRackIfNeeded(int[] queueSizes, int[] partitionIds, String[] partitionLeaderRacks, int length) {
+        if (!rackAware) {
+            return null;
+        }
+        int[] queueSizesInThisRack = new int[length];
+        int[] partitionIdsInThisRack = new int[length];
+        int lengthInThisRack = 0;
+        int maxSizePlus1InThisRack = -1;
+
+        for (int i = 0; i < length; i++) {
+            if (rack.equals(partitionLeaderRacks[i])) {
+                queueSizesInThisRack[lengthInThisRack] = queueSizes[i];
+                partitionIdsInThisRack[lengthInThisRack] = partitionIds[i];
+
+                if (queueSizes[i] > maxSizePlus1InThisRack)
+                    maxSizePlus1InThisRack = queueSizes[i];
+
+                lengthInThisRack += 1;
+            }
+        }
+        ++maxSizePlus1InThisRack;
+
+        invertAndFoldQueueSizeArray(queueSizesInThisRack, maxSizePlus1InThisRack, lengthInThisRack);
+        return new PartitionLoadStats(queueSizesInThisRack, partitionIdsInThisRack, lengthInThisRack);
+    }
+
+    private void invertAndFoldQueueSizeArray(int[] queueSizes, int maxSizePlus1, int length) {

Review Comment:
   The calling functions work the way that `queueSizes` is never empty, but I added the check to make it more robust to future changes.



-- 
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.

To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org

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