Add round robin partitioner

Author: Arseniy-Popov

aiokafka/partitioner.py

@@ -1,5 +1,8 @@
+import logging
 import random
 
+log = logging.getLogger(__name__)
+
 
 class DefaultPartitioner:
     """Default partitioner.
@@ -94,3 +97,37 @@ def murmur2(data):
     h &= 0xFFFFFFFF
 
     return h
+
+
+class RoundRobinPartitioner:
+    """
+    Round robin partitioner.
+
+    If partial ordering is not needed in your use case, round robin partitioning
+    achieves a more even distribution of messages across partitions and enables
+    a higher rate of consumption through a more uniform temporal spacing between
+    messages in a single partition.
+    """
+
+    def __init__(self):
+        self._index = 0
+
+    def __call__(self, key, all_partitions, available_partitions):
+        """
+        Get the next partition according to the round robin algorithm.
+        :param key: partitioning key, expects `None`
+        :param all_partitions: list of all partitions
+        :param available_partitions: list of available partitions
+        :return: one of the values from available_partitions or all_partitions
+        """
+        if key:
+            log.warning(
+                "Partitioning key is ignored by RoundRobinPartitioner - "
+                "use DefaultPartitioner instead."
+            )
+        partitions = available_partitions or all_partitions
+        if self._index >= len(partitions):
+            self._index = 0
+        partition = partitions[self._index]
+        self._index += 1
+        return partition

benchmark/simple_produce_bench.py

@@ -5,6 +5,8 @@
 from collections import Counter
 import random
 
+from aiokafka.partitioner import DefaultPartitioner, RoundRobinPartitioner
+
 
 class Benchmark:
 
@@ -26,9 +28,14 @@ def __init__(self, args):
         else:
             self._is_transactional = False
 
+        if args.partitioner == "default":
+            self._producer_kwargs["partitioner"] = DefaultPartitioner()
+        elif args.partitioner == "round-robin":
+            self._producer_kwargs["partitioner"] = RoundRobinPartitioner()
+
         self.transaction_size = args.transaction_size
 
-        self._partition = args.partition
+        self._partition = args.partition if args.partition != -1 else None
         self._stats_interval = 1
         self._stats = [Counter()]
 
@@ -117,7 +124,7 @@ def parse_args():
         help='Topic to produce messages to. Default {default}.')
     parser.add_argument(
         '--partition', type=int, default=0,
-        help='Partition to produce messages to. Default {default}.')
+        help='Partition to produce messages to. Default {default}. Set to -1 to omit.')
     parser.add_argument(
         '--uvloop', action='store_true',
         help='Use uvloop instead of asyncio default loop.')
@@ -130,6 +137,9 @@ def parse_args():
     parser.add_argument(
         '--transaction-size', type=int, default=100,
         help='Number of messages in transaction')
+    parser.add_argument(
+        '--partitioner', type=str, default="default", choices=["default", "round-robin"],
+        help='Partitioner, either `default` or `round-robin`')
     return parser.parse_args()
 
 

tests/test_partitioner.py

@@ -1,6 +1,6 @@
 import pytest
 
-from aiokafka.partitioner import DefaultPartitioner, murmur2
+from aiokafka.partitioner import DefaultPartitioner, RoundRobinPartitioner, murmur2
 
 
 def test_default_partitioner():
@@ -41,3 +41,21 @@ def test_murmur2_not_ascii():
     # Verify no regression of murmur2() bug encoding py2 bytes that don't ascii encode
     murmur2(b"\xa4")
     murmur2(b"\x81" * 1000)
+
+
+def test_round_robin_partitioner():
+    partitioner = RoundRobinPartitioner()
+
+    all_partitions = available_partitions = list(range(2))
+    assert partitioner(None, all_partitions, available_partitions) == 0
+    assert partitioner(None, all_partitions, available_partitions) == 1
+    assert partitioner(None, all_partitions, available_partitions) == 0
+    assert partitioner(None, all_partitions, available_partitions) == 1
+
+    all_partitions = available_partitions = list(range(4))
+    assert partitioner(None, all_partitions, available_partitions) == 2
+    assert partitioner(None, all_partitions, available_partitions) == 3
+
+    all_partitions = [50]
+    available_partitions = [70]
+    assert partitioner(None, all_partitions, available_partitions) == 70