KAFKA-17632: Fix RoundRobinPartitioner for even partition counts (apache#17620)

RoundRobinPartitioner does not handle the fact that on new batch creation, the partition method is called twice.

Reviewers: Viktor Somogyi-Vass <[email protected]>, Mickael Maison <[email protected]>

Author: urbandan

clients/src/main/java/org/apache/kafka/clients/producer/KafkaProducer.java

@@ -1074,6 +1074,8 @@ private Future<RecordMetadata> doSend(ProducerRecord<K, V> record, Callback call
 
             if (result.abortForNewBatch) {
                 int prevPartition = partition;
+                // IMPORTANT NOTE: the following onNewBatch and partition calls should not interrupted to allow
+                // the custom partitioner to correctly track its state
                 onNewBatch(record.topic(), cluster, prevPartition);
                 partition = partition(record, serializedKey, serializedValue, cluster);
                 if (log.isTraceEnabled()) {

clients/src/main/java/org/apache/kafka/clients/producer/Partitioner.java

@@ -49,6 +49,9 @@ public interface Partitioner extends Configurable, Closeable {
      * <p>
      * Notifies the partitioner a new batch is about to be created. When using the sticky partitioner,
      * this method can change the chosen sticky partition for the new batch.
+     * <p>
+     * After onNewBatch, the {@link #partition(String, Object, byte[], Object, byte[], Cluster)} method is called again
+     * which allows the implementation to "redirect" the message on new batch creation.
      * @param topic The topic name
      * @param cluster The current cluster metadata
      * @param prevPartition The partition previously selected for the record that triggered a new batch

clients/src/main/java/org/apache/kafka/clients/producer/RoundRobinPartitioner.java

@@ -18,6 +18,7 @@
 
 import org.apache.kafka.common.Cluster;
 import org.apache.kafka.common.PartitionInfo;
+import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.utils.Utils;
 
 import java.util.List;
@@ -36,6 +37,7 @@
  */
 public class RoundRobinPartitioner implements Partitioner {
     private final ConcurrentMap<String, AtomicInteger> topicCounterMap = new ConcurrentHashMap<>();
+    private final ThreadLocal<TopicPartition> previousPartition = new ThreadLocal<>();
 
     public void configure(Map<String, ?> configs) {}
 
@@ -51,6 +53,14 @@ public void configure(Map<String, ?> configs) {}
      */
     @Override
     public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
+        TopicPartition prevPartition = previousPartition.get();
+        if (prevPartition != null) {
+            previousPartition.remove();
+            if (topic.equals(prevPartition.topic())) {
+                return prevPartition.partition();
+            }
+        }
+
         int nextValue = nextValue(topic);
         List<PartitionInfo> availablePartitions = cluster.availablePartitionsForTopic(topic);
         if (!availablePartitions.isEmpty()) {
@@ -68,6 +78,11 @@ private int nextValue(String topic) {
         return counter.getAndIncrement();
     }
 
-    public void close() {}
+    @SuppressWarnings("deprecation")
+    @Override
+    public void onNewBatch(String topic, Cluster cluster, int prevPartition) {
+        previousPartition.set(new TopicPartition(topic, prevPartition));
+    }
 
+    public void close() {}
 }

clients/src/test/java/org/apache/kafka/clients/producer/RoundRobinPartitionerTest.java

@@ -97,6 +97,7 @@ public void testRoundRobinWithKeyBytes() {
         assertEquals(10, partitionCount.get(2).intValue());
     }
 
+    @SuppressWarnings("deprecation")
     @Test
     public void testRoundRobinWithNullKeyBytes() {
         final String topicA = "topicA";
@@ -113,6 +114,10 @@ public void testRoundRobinWithNullKeyBytes() {
         Partitioner partitioner = new RoundRobinPartitioner();
         for (int i = 0; i < 30; ++i) {
             int partition = partitioner.partition(topicA, null, null, null, null, testCluster);
+            // Simulate single-message batches
+            partitioner.onNewBatch(topicA, testCluster, partition);
+            int nextPartition = partitioner.partition(topicA, null, null, null, null, testCluster);
+            assertEquals(partition, nextPartition, "New batch creation should not affect the partition selection");
             Integer count = partitionCount.get(partition);
             if (null == count)
                 count = 0;
@@ -126,5 +131,42 @@ public void testRoundRobinWithNullKeyBytes() {
         assertEquals(10, partitionCount.get(0).intValue());
         assertEquals(10, partitionCount.get(1).intValue());
         assertEquals(10, partitionCount.get(2).intValue());
-    }    
+    }
+
+    @SuppressWarnings("deprecation")
+    @Test
+    public void testRoundRobinWithNullKeyBytesAndEvenPartitionCount() {
+        final String topicA = "topicA";
+        final String topicB = "topicB";
+
+        List<PartitionInfo> allPartitions = asList(new PartitionInfo(topicA, 0, NODES[0], NODES, NODES),
+                new PartitionInfo(topicA, 1, NODES[1], NODES, NODES), new PartitionInfo(topicA, 2, NODES[2], NODES, NODES),
+                new PartitionInfo(topicB, 0, NODES[0], NODES, NODES), new PartitionInfo(topicA, 3, NODES[0], NODES, NODES));
+        Cluster testCluster = new Cluster("clusterId", asList(NODES[0], NODES[1], NODES[2]), allPartitions,
+                Collections.emptySet(), Collections.emptySet());
+
+        final Map<Integer, Integer> partitionCount = new HashMap<>();
+
+        Partitioner partitioner = new RoundRobinPartitioner();
+        for (int i = 0; i < 40; ++i) {
+            int partition = partitioner.partition(topicA, null, null, null, null, testCluster);
+            // Simulate single-message batches
+            partitioner.onNewBatch(topicA, testCluster, partition);
+            int nextPartition = partitioner.partition(topicA, null, null, null, null, testCluster);
+            assertEquals(partition, nextPartition, "New batch creation should not affect the partition selection");
+            Integer count = partitionCount.get(partition);
+            if (null == count)
+                count = 0;
+            partitionCount.put(partition, count + 1);
+
+            if (i % 5 == 0) {
+                partitioner.partition(topicB, null, null, null, null, testCluster);
+            }
+        }
+
+        assertEquals(10, partitionCount.get(0).intValue());
+        assertEquals(10, partitionCount.get(1).intValue());
+        assertEquals(10, partitionCount.get(2).intValue());
+        assertEquals(10, partitionCount.get(3).intValue());
+    }
 }