Add rate balance module (#59)

The rate balance module is a variant on size balance, except that it scales the size used based on the retention of the topic. Topics that are retained longer
than the cluster default (now specified on the command line with --default-retention) will be weighted lower, such that if a topic has double the retention of
the cluster, its scaled size will be half its actual size on disk. This will more closely approximate data rate balancing, which was the original intention of
the size module.

kafka/tools/assigner/__main__.py

@@ -106,7 +106,7 @@ def main():
     tools_path = get_tools_path(args.tools_path)
     check_java_home()
 
-    cluster = Cluster.create_from_zookeeper(args.zookeeper)
+    cluster = Cluster.create_from_zookeeper(args.zookeeper, getattr(args, 'default_retention', 1))
     run_plugins_at_step(plugins, 'set_cluster', cluster)
 
     # If the module needs the partition sizes, call a size module to get the information

kafka/tools/assigner/actions/balance.py

@@ -41,6 +41,7 @@ def _add_args(cls, parser):
         balance_actions = get_modules(kafka.tools.assigner.actions.balancemodules, ActionBalanceModule)
         parser.add_argument('-t', '--types', help="Balance types to perform. Multiple may be specified and they will be run in order", required=True,
                             choices=[klass.name for klass in balance_actions], nargs='*')
+        parser.add_argument('--default-retention', help="Default cluster retention, in ms", required=False, type=int, default=345600000)
 
     def process_cluster(self):
         for bmodule in self.modules:

kafka/tools/assigner/actions/balancemodules/rate.py

@@ -0,0 +1,34 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from kafka.tools.assigner.actions.balancemodules.size import ActionBalanceSize
+from kafka.tools.assigner.actions import ActionBalanceModule
+
+
+class ActionBalanceRate(ActionBalanceModule):
+    name = "rate"
+    helpstr = "Move the busiest partitions in the cluster to even the total bytes rate per-broker for each replica position"
+
+    def __init__(self, args, cluster):
+        super(ActionBalanceRate, self).__init__(args, cluster)
+
+        # This module merely calls the balance "size" module with the scaled_size attr, so we're going to instantiate a copy of that
+        self._size = ActionBalanceSize(args, cluster, size_attr='scaled_size')
+
+    def process_cluster(self):
+        # Call the size module
+        self._size.process_cluster()

kafka/tools/assigner/actions/balancemodules/size.py

@@ -25,8 +25,13 @@ class ActionBalanceSize(ActionBalanceModule):
     name = "size"
     helpstr = "Move the largest partitions in the cluster to even the total size on disk per-broker for each replica position"
 
+    # We override this so we can select either the size or the scaled_size attributes to sort on
+    def __init__(self, args, cluster, size_attr='size'):
+        super(ActionBalanceSize, self).__init__(args, cluster)
+        self._size_attr = size_attr
+
     def process_cluster(self):
-        log.info("Starting partition balance by size")
+        log.info("Starting partition balance by {0}".format(self._size_attr))
 
         # Figure out the max RF for the cluster
         max_rf = self.cluster.max_replication_factor()
@@ -43,24 +48,24 @@ def process_cluster(self):
 
             # Create a sorted list of partitions to use at this position (descending size)
             # Throw out partitions that are 4K or less in size, as they are effectively empty
-            partitions[pos] = [p for p in self.cluster.partitions(self.args.exclude_topics) if (len(p.replicas) > pos) and (p.size > 4)]
-            partitions[pos].sort(key=attrgetter('size'), reverse=True)
+            partitions[pos] = [p for p in self.cluster.partitions(self.args.exclude_topics) if (len(p.replicas) > pos) and (getattr(p, self._size_attr) > 4)]
+            partitions[pos].sort(key=attrgetter(self._size_attr), reverse=True)
 
             # Calculate broker size at this position
             for broker in self.cluster.brokers:
                 if pos in self.cluster.brokers[broker].partitions:
-                    sizes[pos][broker] = sum([p.size for p in self.cluster.brokers[broker].partitions[pos]], 0)
+                    sizes[pos][broker] = sum([getattr(p, self._size_attr) for p in self.cluster.brokers[broker].partitions[pos]], 0)
                 else:
                     sizes[pos][broker] = 0
 
             # Calculate the median size of partitions (margin is median/2) and the average size per broker to target
             # Yes, I know the median calculation is slightly broken (it keeps integers). This is OK
-            targets[pos] = sum([p.size for p in partitions[pos]], 0) // len(self.cluster.brokers)
+            targets[pos] = sum([getattr(p, self._size_attr) for p in partitions[pos]], 0) // len(self.cluster.brokers)
             sizelen = len(partitions[pos])
             if not sizelen % 2:
-                margins[pos] = (partitions[pos][sizelen // 2].size + partitions[pos][sizelen // 2 - 1].size) // 4
+                margins[pos] = (getattr(partitions[pos][sizelen // 2], self._size_attr) + getattr(partitions[pos][sizelen // 2 - 1], self._size_attr)) // 4
             else:
-                margins[pos] = partitions[pos][sizelen // 2].size // 2
+                margins[pos] = getattr(partitions[pos][sizelen // 2], self._size_attr) // 2
 
         # Balance partitions for each replica position separately
         for pos in range(max_rf):
@@ -79,34 +84,36 @@ def process_cluster(self):
 
                 # Find partitions to move to this broker
                 for partition in partitions[pos]:
+                    partition_size = getattr(partition, self._size_attr)
+
                     # We can use this partition if all of the following are true: the partition has a replica at this position,
                     # it's size is less than or equal to the max move size, the broker at this replica position would not go out
                     # of range, and it doesn't already exist on this broker at this position
-                    if ((len(partition.replicas) <= pos) or (partition.size > max_move) or
-                       ((sizes[pos][partition.replicas[pos].id] - partition.size) < (targets[pos] - margins[pos])) or
+                    if ((len(partition.replicas) <= pos) or (partition_size > max_move) or
+                       ((sizes[pos][partition.replicas[pos].id] - partition_size) < (targets[pos] - margins[pos])) or
                        (partition.replicas[pos] == broker)):
                         continue
 
                     # We can only use a partition that this replica exists on if swapping positions wouldn't hurt balance of the other position or broker
                     source = partition.replicas[pos]
                     if broker in partition.replicas:
                         other_pos = partition.replicas.index(broker)
-                        if ((sizes[other_pos][broker_id] - partition.size < targets[other_pos] - margins[other_pos]) or
-                           (sizes[other_pos][source.id] + partition.size > targets[pos] + margins[pos]) or
-                           (sizes[pos][broker_id] + partition.size > targets[pos] + margins[pos]) or
-                           (sizes[pos][source.id] - partition.size < targets[pos] - margins[pos])):
+                        if ((sizes[other_pos][broker_id] - partition_size < targets[other_pos] - margins[other_pos]) or
+                           (sizes[other_pos][source.id] + partition_size > targets[pos] + margins[pos]) or
+                           (sizes[pos][broker_id] + partition_size > targets[pos] + margins[pos]) or
+                           (sizes[pos][source.id] - partition_size < targets[pos] - margins[pos])):
                             continue
 
                         partition.swap_replica_positions(source, broker)
-                        sizes[other_pos][broker_id] -= partition.size
-                        sizes[other_pos][source.id] += partition.size
+                        sizes[other_pos][broker_id] -= partition_size
+                        sizes[other_pos][source.id] += partition_size
                     else:
                         # Move the partition and adjust sizes
                         partition.swap_replicas(source, broker)
-                    sizes[pos][broker_id] += partition.size
-                    sizes[pos][source.id] -= partition.size
-                    min_move -= partition.size
-                    max_move -= partition.size
+                    sizes[pos][broker_id] += partition_size
+                    sizes[pos][source.id] -= partition_size
+                    min_move -= partition_size
+                    max_move -= partition_size
 
                     # If we have moved enough partitions, stop for this broker
                     if min_move <= 0:

kafka/tools/assigner/models/cluster.py

@@ -17,6 +17,7 @@
 
 import json
 from kazoo.client import KazooClient
+from kazoo.exceptions import KazooException
 
 from kafka.tools.assigner import log
 from kafka.tools.assigner.exceptions import ZookeeperException, ClusterConsistencyException
@@ -45,15 +46,26 @@ def add_topic_with_replicas(cluster, topic, topic_data):
     cluster.add_topic(newtopic)
 
 
+def set_topic_retention(topic, zk):
+    try:
+        zdata, zstat = zk.get("/config/topics/{0}".format(topic.name))
+        tdata = json.loads(zdata)
+        topic.retention = int(tdata['config']['retention.ms'])
+    except (KeyError, ValueError, KazooException):
+        # If we can't get the config override for any reason, just stick with whatever the default is
+        pass
+
+
 class Cluster(BaseModel):
     equality_attrs = ['brokers', 'topics']
 
-    def __init__(self):
+    def __init__(self, retention=1):
         self.brokers = {}
         self.topics = {}
+        self.retention = retention
 
     @classmethod
-    def create_from_zookeeper(cls, zkconnect):
+    def create_from_zookeeper(cls, zkconnect, default_retention=1):
         log.info("Connecting to zookeeper {0}".format(zkconnect))
         try:
             zk = KazooClient(zkconnect)
@@ -62,14 +74,15 @@ def create_from_zookeeper(cls, zkconnect):
             raise ZookeeperException("Cannot connect to Zookeeper: {0}".format(e))
 
         # Get broker list
-        cluster = cls()
+        cluster = cls(retention=default_retention)
         add_brokers_from_zk(cluster, zk)
 
         # Get current partition state
         log.info("Getting partition list from Zookeeper")
         for topic in zk.get_children("/brokers/topics"):
             zdata, zstat = zk.get("/brokers/topics/{0}".format(topic))
             add_topic_with_replicas(cluster, topic, json.loads(zdata))
+            set_topic_retention(cluster.topics[topic], zk)
 
         if cluster.num_topics() == 0:
             raise ZookeeperException("The cluster specified does not have any topics")
@@ -82,6 +95,7 @@ def create_from_zookeeper(cls, zkconnect):
 
     def clone(self):
         newcluster = Cluster()
+        newcluster.retention = self.retention
 
         # We're not going to clone in the subclasses because we need to map partitions between topics and brokers
         # So we do shallow copies and populate the partition information ourselves
@@ -108,6 +122,7 @@ def add_broker(self, broker):
 
     def add_topic(self, topic):
         topic.cluster = self
+        topic.retention = self.retention
         self.topics[topic.name] = topic
 
     # Iterate over all the partitions in this cluster

kafka/tools/assigner/models/partition.py

@@ -15,6 +15,8 @@
 # specific language governing permissions and limitations
 # under the License.
 
+from __future__ import division
+
 from kafka.tools.assigner.exceptions import ReplicaNotFoundException
 from kafka.tools.assigner.models import BaseModel
 
@@ -27,6 +29,7 @@ def __init__(self, topic, num):
         self.num = num
         self.replicas = []
         self.size = 0
+        self.scaled_size = 0
 
     # Shallow copy - do not copy replica list (zero length)
     def copy(self):
@@ -38,6 +41,7 @@ def copy(self):
     def set_size(self, size):
         if size > self.size:
             self.size = size
+            self.scaled_size = (self.topic.cluster.retention / self.topic.retention) * self.size
 
     def dict_for_reassignment(self):
         return {"topic": self.topic.name, "partition": self.num, "replicas": [broker.id for broker in self.replicas]}

kafka/tools/assigner/models/topic.py

@@ -26,6 +26,7 @@ def __init__(self, name, partitions):
         self.name = name
         self.partitions = []
         self.cluster = None
+        self.retention = 1
 
         for i in range(partitions):
             self.add_partition(Partition(self, i))

tests/tools/assigner/actions/balancemodules/test_rate.py

@@ -0,0 +1,44 @@
+import sys
+import unittest
+
+from argparse import Namespace
+from ..fixtures import set_up_cluster, set_up_subparser
+
+from kafka.tools.assigner.actions.balance import ActionBalance
+from kafka.tools.assigner.actions.balancemodules.rate import ActionBalanceRate
+
+
+class ActionBalanceRateTests(unittest.TestCase):
+    def setUp(self):
+        self.cluster = set_up_cluster()
+        self.cluster.topics['testTopic1'].partitions[0].set_size(1000)
+        self.cluster.topics['testTopic1'].partitions[1].set_size(1000)
+        self.cluster.topics['testTopic2'].partitions[0].set_size(2000)
+        self.cluster.topics['testTopic2'].partitions[1].set_size(2000)
+
+        (self.parser, self.subparsers) = set_up_subparser()
+        self.args = Namespace(exclude_topics=[])
+
+    def test_configure_args(self):
+        ActionBalance.configure_args(self.subparsers)
+        sys.argv = ['kafka-assigner', 'balance', '-t', 'rate']
+        parsed_args = self.parser.parse_args()
+        assert parsed_args.action == 'balance'
+
+    def test_create_class(self):
+        action = ActionBalanceRate(self.args, self.cluster)
+        assert isinstance(action, ActionBalanceRate)
+
+    # This is a duplicate of the size test, but we need at least one test to make sure it does work
+    def test_process_cluster_one_move(self):
+        b1 = self.cluster.brokers[1]
+        b2 = self.cluster.brokers[2]
+        self.cluster.topics['testTopic1'].partitions[0].swap_replica_positions(b1, b2)
+
+        action = ActionBalanceRate(self.args, self.cluster)
+        action.process_cluster()
+
+        assert sum([p.scaled_size for p in self.cluster.brokers[1].partitions[0]], 0) == 3000
+        assert sum([p.scaled_size for p in self.cluster.brokers[1].partitions[1]], 0) == 3000
+        assert sum([p.scaled_size for p in self.cluster.brokers[2].partitions[0]], 0) == 3000
+        assert sum([p.scaled_size for p in self.cluster.brokers[2].partitions[1]], 0) == 3000

tests/tools/assigner/actions/fixtures.py

@@ -7,6 +7,7 @@
 
 def set_up_cluster():
     cluster = Cluster()
+    cluster.retention = 100000
     cluster.add_broker(Broker(1, "brokerhost1.example.com"))
     cluster.add_broker(Broker(2, "brokerhost2.example.com"))
     cluster.brokers[1].rack = "a"

tests/tools/assigner/models/test_cluster_zookeeper.py

@@ -46,7 +46,9 @@ def test_cluster_create_missing_broker(self):
                                      (('{"jmx_port":7667,"timestamp":"1465289114807","endpoints":["PLAINTEXT://brokerhost2.example.com:9223",'
                                       '"SSL://brokerhost2.example.com:9224"],"host":"brokerhost2.example.com","version":1,"port":9223}'), None),
                                      ('{"version":1,"partitions":{"0":[1,2],"1":[3,1]}}', None),
-                                     ('{"version":1,"partitions":{"0":[3,1],"1":[1,2]}}', None)]
+                                     ('{"version":1,"config":{}}', None),
+                                     ('{"version":1,"partitions":{"0":[3,1],"1":[1,2]}}', None),
+                                     ('{"version":1,"config":{}}', None)]
         cluster = Cluster.create_from_zookeeper('zkconnect')
 
         assert len(cluster.brokers) == 3
@@ -63,7 +65,9 @@ def test_cluster_create_from_zookeeper(self):
                                      (('{"jmx_port":7667,"timestamp":"1465289114807","endpoints":["PLAINTEXT://brokerhost2.example.com:9223",'
                                       '"SSL://brokerhost2.example.com:9224"],"host":"brokerhost2.example.com","version":1,"port":9223}'), None),
                                      ('{"version":1,"partitions":{"0":[1,2],"1":[2,1]}}', None),
-                                     ('{"version":1,"partitions":{"0":[2,1],"1":[1,2]}}', None)]
+                                     ('{"version":1,"config":{}}', None),
+                                     ('{"version":1,"partitions":{"0":[2,1],"1":[1,2]}}', None),
+                                     ('{"version":1,"config":{}}', None)]
         cluster = Cluster.create_from_zookeeper('zkconnect')
 
         assert len(cluster.brokers) == 2
@@ -85,3 +89,18 @@ def test_cluster_create_from_zookeeper(self):
         assert len(t2.partitions) == 2
         assert t2.partitions[0].replicas == [b2, b1]
         assert t2.partitions[1].replicas == [b1, b2]
+
+    def test_cluster_create_with_retention(self):
+        self.mock_children.side_effect = [['1', '2'], ['testTopic1', 'testTopic2']]
+        self.mock_get.side_effect = [(('{"jmx_port":7667,"timestamp":"1465289114807","endpoints":["PLAINTEXT://brokerhost1.example.com:9223",'
+                                     '"SSL://brokerhost1.example.com:9224"],"host":"brokerhost1.example.com","version":1,"port":9223}'), None),
+                                     (('{"jmx_port":7667,"timestamp":"1465289114807","endpoints":["PLAINTEXT://brokerhost2.example.com:9223",'
+                                      '"SSL://brokerhost2.example.com:9224"],"host":"brokerhost2.example.com","version":1,"port":9223}'), None),
+                                     ('{"version":1,"partitions":{"0":[1,2],"1":[2,1]}}', None),
+                                     ('{"version":1,"config":{}}', None),
+                                     ('{"version":1,"partitions":{"0":[2,1],"1":[1,2]}}', None),
+                                     ('{"version":1,"config":{"retention.ms":"172800000"}}', None)]
+        cluster = Cluster.create_from_zookeeper('zkconnect')
+
+        assert cluster.topics['testTopic1'].retention == 1
+        assert cluster.topics['testTopic2'].retention == 172800000

tests/tools/assigner/models/test_topic_and_partition.py

@@ -1,18 +1,22 @@
 import unittest
 
+from kafka.tools.assigner.models.cluster import Cluster
 from kafka.tools.assigner.models.broker import Broker
 from kafka.tools.assigner.models.topic import Topic
 from kafka.tools.assigner.models.partition import Partition
 
 
 class TopicAndPartitionTests(unittest.TestCase):
     def setUp(self):
+        self.cluster = Cluster(retention=500000)
         self.topic = Topic('testTopic', 1)
+        self.topic.cluster = self.cluster
+        self.topic.cluster.topics['testTopic'] = self.topic
 
     def test_topic_create(self):
         assert self.topic.name == 'testTopic'
         assert len(self.topic.partitions) == 1
-        assert self.topic.cluster is None
+        assert self.topic.cluster == self.cluster
 
     def test_partition_create(self):
         assert len(self.topic.partitions) == 1
@@ -100,6 +104,16 @@ def test_set_partition_size_larger(self):
         self.topic.partitions[0].set_size(200)
         assert self.topic.partitions[0].size == 200
 
+    def test_set_partition_scaled_size_smaller(self):
+        self.topic.retention = 1000000
+        self.topic.partitions[0].set_size(100)
+        assert self.topic.partitions[0].scaled_size == 50
+
+    def test_set_partition_scaled_size_larger(self):
+        self.topic.retention = 100000
+        self.topic.partitions[0].set_size(100)
+        assert self.topic.partitions[0].scaled_size == 500
+
     def test_partition_dict_for_reassignment_without_replicas(self):
         expected = {"topic": 'testTopic', "partition": 0, "replicas": []}
         assert self.topic.partitions[0].dict_for_reassignment() == expected