def setUp(self):
self.cluster = Cluster()
self.cluster.add_broker(Broker(1, "brokerhost1.example.com"))
self.cluster.add_broker(Broker(2, "brokerhost2.example.com"))
self.cluster.add_broker(Broker(3, "brokerhost3.example.com"))
self.cluster.add_topic(Topic("testTopic1", 2))
self.cluster.add_topic(Topic("testTopic2", 2))
def set_up_cluster_4broker():
cluster = Cluster()
cluster.add_broker(Broker(1, "brokerhost1.example.com"))
cluster.add_broker(Broker(2, "brokerhost2.example.com"))
cluster.add_broker(Broker(3, "brokerhost3.example.com"))
cluster.add_broker(Broker(4, "brokerhost4.example.com"))
cluster.brokers[1].rack = "a"
cluster.brokers[2].rack = "a"
cluster.brokers[3].rack = "b"
cluster.brokers[4].rack = "b"
cluster.add_topic(Topic("testTopic1", 4))
cluster.add_topic(Topic("testTopic2", 4))
cluster.add_topic(Topic("testTopic3", 4))
partition = cluster.topics['testTopic1'].partitions[0]
partition.add_replica(cluster.brokers[1], 0)
partition.add_replica(cluster.brokers[2], 1)
partition = cluster.topics['testTopic1'].partitions[1]
partition.add_replica(cluster.brokers[2], 0)
partition.add_replica(cluster.brokers[3], 1)
partition = cluster.topics['testTopic1'].partitions[2]
partition.add_replica(cluster.brokers[2], 0)
partition.add_replica(cluster.brokers[3], 1)
partition = cluster.topics['testTopic1'].partitions[3]
partition.add_replica(cluster.brokers[4], 0)
partition.add_replica(cluster.brokers[1], 1)
partition = cluster.topics['testTopic2'].partitions[0]
partition.add_replica(cluster.brokers[4], 0)
partition.add_replica(cluster.brokers[3], 1)
partition = cluster.topics['testTopic2'].partitions[1]
partition.add_replica(cluster.brokers[2], 0)
partition.add_replica(cluster.brokers[4], 1)
partition = cluster.topics['testTopic2'].partitions[2]
partition.add_replica(cluster.brokers[2], 0)
partition.add_replica(cluster.brokers[1], 1)
partition = cluster.topics['testTopic2'].partitions[3]
partition.add_replica(cluster.brokers[3], 0)
partition.add_replica(cluster.brokers[1], 1)
partition = cluster.topics['testTopic3'].partitions[0]
partition.add_replica(cluster.brokers[3], 0)
partition.add_replica(cluster.brokers[2], 1)
partition = cluster.topics['testTopic3'].partitions[1]
partition.add_replica(cluster.brokers[4], 0)
partition.add_replica(cluster.brokers[2], 1)
partition = cluster.topics['testTopic3'].partitions[2]
partition.add_replica(cluster.brokers[1], 0)
partition.add_replica(cluster.brokers[2], 1)
partition = cluster.topics['testTopic3'].partitions[3]
partition.add_replica(cluster.brokers[3], 0)
partition.add_replica(cluster.brokers[4], 1)
return cluster
def create_cluster_onehost(self):
cluster = Cluster()
cluster.add_broker(Broker(1, "brokerhost1.example.com"))
cluster.add_topic(Topic("testTopic1", 2))
cluster.add_topic(Topic("testTopic2", 2))
partition = cluster.topics['testTopic1'].partitions[0]
partition.add_replica(cluster.brokers[1], 0)
partition = cluster.topics['testTopic1'].partitions[1]
partition.add_replica(cluster.brokers[1], 0)
partition = cluster.topics['testTopic2'].partitions[0]
partition.add_replica(cluster.brokers[1], 0)
partition = cluster.topics['testTopic2'].partitions[1]
partition.add_replica(cluster.brokers[1], 0)
return cluster
def test_cluster_create_from_zookeeper(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,"partitions":{"0":[2,1],"1":[1,2]}}', None)
]
cluster = Cluster.create_from_zookeeper('zkconnect')
assert len(cluster.brokers) == 2
b1 = cluster.brokers[1]
b2 = cluster.brokers[2]
assert b1.hostname == 'brokerhost1.example.com'
assert b2.hostname == 'brokerhost2.example.com'
assert len(cluster.topics) == 2
assert 'testTopic1' in cluster.topics
assert 'testTopic2' in cluster.topics
t1 = cluster.topics['testTopic1']
assert len(t1.partitions) == 2
assert t1.partitions[0].replicas == [b1, b2]
assert t1.partitions[1].replicas == [b2, b1]
t2 = cluster.topics['testTopic2']
assert len(t2.partitions) == 2
assert t2.partitions[0].replicas == [b2, b1]
assert t2.partitions[1].replicas == [b1, b2]
def setUp(self):
self.cluster = Cluster()
self.cluster.add_broker(Broker(1, "brokerhost1.example.com"))
self.cluster.add_broker(Broker(2, "brokerhost2.example.com"))
self.cluster.add_broker(Broker(3, "brokerhost3.example.com"))
self.cluster.add_topic(Topic("testTopic1", 2))
self.cluster.add_topic(Topic("testTopic2", 2))
def test_cluster_create_from_zookeeper(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,"partitions":{"0":[2,1],"1":[1,2]}}', None)]
cluster = Cluster.create_from_zookeeper('zkconnect')
assert len(cluster.brokers) == 2
b1 = cluster.brokers[1]
b2 = cluster.brokers[2]
assert b1.hostname == 'brokerhost1.example.com'
assert b2.hostname == 'brokerhost2.example.com'
assert len(cluster.topics) == 2
assert 'testTopic1' in cluster.topics
assert 'testTopic2' in cluster.topics
t1 = cluster.topics['testTopic1']
assert len(t1.partitions) == 2
assert t1.partitions[0].replicas == [b1, b2]
assert t1.partitions[1].replicas == [b2, b1]
t2 = cluster.topics['testTopic2']
assert len(t2.partitions) == 2
assert t2.partitions[0].replicas == [b2, b1]
assert t2.partitions[1].replicas == [b1, b2]
def set_up_cluster():
cluster = Cluster()
cluster.add_broker(Broker(1, "brokerhost1.example.com"))
cluster.add_broker(Broker(2, "brokerhost2.example.com"))
cluster.add_topic(Topic("testTopic1", 2))
cluster.add_topic(Topic("testTopic2", 2))
partition = cluster.topics['testTopic1'].partitions[0]
partition.add_replica(cluster.brokers[1], 0)
partition.add_replica(cluster.brokers[2], 1)
partition = cluster.topics['testTopic1'].partitions[1]
partition.add_replica(cluster.brokers[2], 0)
partition.add_replica(cluster.brokers[1], 1)
partition = cluster.topics['testTopic2'].partitions[0]
partition.add_replica(cluster.brokers[2], 0)
partition.add_replica(cluster.brokers[1], 1)
partition = cluster.topics['testTopic2'].partitions[1]
partition.add_replica(cluster.brokers[1], 0)
partition.add_replica(cluster.brokers[2], 1)
return cluster
def setUp(self):
self.aparser = argparse.ArgumentParser(
prog='kafka-assigner',
description='Rejigger Kafka cluster partitions')
self.subparsers = self.aparser.add_subparsers(
help='Select manipulation module to use')
self.args = argparse.Namespace()
self.cluster = Cluster()
self.plugin = PluginModule()
def main():
# Start by loading all the modules
action_map = get_module_map(kafka.tools.assigner.actions, kafka.tools.assigner.actions.ActionModule)
sizer_map = get_module_map(kafka.tools.assigner.sizers, kafka.tools.assigner.sizers.SizerModule)
plugins = get_all_plugins()
# Set up and parse all CLI arguments
args = set_up_arguments(action_map, sizer_map, plugins)
run_plugins_at_step(plugins, 'set_arguments', args)
tools_path = get_tools_path(args.tools_path)
check_java_home()
cluster = Cluster.create_from_zookeeper(args.zookeeper)
run_plugins_at_step(plugins, 'set_cluster', cluster)
# If the module needs the partition sizes, call a size module to get the information
check_and_get_sizes(action_map[args.action], args, cluster, sizer_map)
run_plugins_at_step(plugins, 'after_sizes')
print_leadership("before", cluster, args.leadership)
# Clone the cluster, and run the action to generate a new cluster state
newcluster = cluster.clone()
action_to_run = action_map[args.action](args, newcluster)
action_to_run.process_cluster()
run_plugins_at_step(plugins, 'set_new_cluster', action_to_run.cluster)
print_leadership("after", newcluster, args.leadership)
move_partitions = cluster.changed_partitions(action_to_run.cluster)
batches = split_partitions_into_batches(move_partitions, batch_size=args.moves, use_class=Reassignment)
run_plugins_at_step(plugins, 'set_batches', batches)
log.info("Partition moves required: {0}".format(len(move_partitions)))
log.info("Number of batches: {0}".format(len(batches)))
dry_run = is_dry_run(args)
for i, batch in enumerate(batches):
log.info("Executing partition reassignment {0}/{1}: {2}".format(i + 1, len(batches), repr(batch)))
batch.execute(i + 1, len(batches), args.zookeeper, tools_path, plugins, dry_run)
run_plugins_at_step(plugins, 'before_ple')
if not args.skip_ple:
all_cluster_partitions = [p for p in action_to_run.cluster.partitions()]
batches = split_partitions_into_batches(all_cluster_partitions, batch_size=args.ple_size, use_class=ReplicaElection)
log.info("Number of replica elections: {0}".format(len(batches)))
run_preferred_replica_elections(batches, args, tools_path, plugins, dry_run)
run_plugins_at_step(plugins, 'finished')
return os.EX_OK
def test_cluster_create_missing_broker(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":[3,1]}}', None),
('{"version":1,"partitions":{"0":[3,1],"1":[1,2]}}', None)]
cluster = Cluster.create_from_zookeeper('zkconnect')
assert len(cluster.brokers) == 3
b1 = cluster.brokers[1]
b3 = cluster.brokers[3]
assert cluster.brokers[3].hostname is None
assert cluster.topics['testTopic1'].partitions[1].replicas == [b3, b1]
assert cluster.topics['testTopic2'].partitions[0].replicas == [b3, b1]
def create_cluster_twohost(self):
self.cluster = Cluster()
self.cluster.add_broker(Broker(1, "brokerhost1.example.com"))
self.cluster.add_broker(Broker(2, "brokerhost2.example.com"))
self.cluster.add_topic(Topic("testTopic1", 2))
self.cluster.add_topic(Topic("testTopic2", 2))
partition = self.cluster.topics['testTopic1'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
partition.add_replica(self.cluster.brokers[2], 1)
partition = self.cluster.topics['testTopic1'].partitions[1]
partition.add_replica(self.cluster.brokers[2], 0)
partition.add_replica(self.cluster.brokers[1], 1)
partition = self.cluster.topics['testTopic2'].partitions[0]
partition.add_replica(self.cluster.brokers[2], 0)
partition.add_replica(self.cluster.brokers[1], 1)
partition = self.cluster.topics['testTopic2'].partitions[1]
partition.add_replica(self.cluster.brokers[1], 0)
partition.add_replica(self.cluster.brokers[2], 1)
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
def create_cluster_onehost(self):
cluster = Cluster()
cluster.add_broker(Broker(1, "brokerhost1.example.com"))
cluster.add_topic(Topic("testTopic1", 2))
cluster.add_topic(Topic("testTopic2", 2))
partition = cluster.topics['testTopic1'].partitions[0]
partition.add_replica(cluster.brokers[1], 0)
partition = cluster.topics['testTopic1'].partitions[1]
partition.add_replica(cluster.brokers[1], 0)
partition = cluster.topics['testTopic2'].partitions[0]
partition.add_replica(cluster.brokers[1], 0)
partition = cluster.topics['testTopic2'].partitions[1]
partition.add_replica(cluster.brokers[1], 0)
return cluster
def test_cluster_create_missing_broker(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":[3,1]}}', None),
('{"version":1,"partitions":{"0":[3,1],"1":[1,2]}}', None)
]
cluster = Cluster.create_from_zookeeper('zkconnect')
assert len(cluster.brokers) == 3
b1 = cluster.brokers[1]
b3 = cluster.brokers[3]
assert cluster.brokers[3].hostname is None
assert cluster.topics['testTopic1'].partitions[1].replicas == [b3, b1]
assert cluster.topics['testTopic2'].partitions[0].replicas == [b3, b1]
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"
cluster.brokers[2].rack = "b"
cluster.add_topic(Topic("testTopic1", 2))
cluster.add_topic(Topic("testTopic2", 2))
partition = cluster.topics['testTopic1'].partitions[0]
partition.add_replica(cluster.brokers[1], 0)
partition.add_replica(cluster.brokers[2], 1)
partition = cluster.topics['testTopic1'].partitions[1]
partition.add_replica(cluster.brokers[2], 0)
partition.add_replica(cluster.brokers[1], 1)
partition = cluster.topics['testTopic2'].partitions[0]
partition.add_replica(cluster.brokers[2], 0)
partition.add_replica(cluster.brokers[1], 1)
partition = cluster.topics['testTopic2'].partitions[1]
partition.add_replica(cluster.brokers[1], 0)
partition.add_replica(cluster.brokers[2], 1)
return cluster
def set_up_cluster():
cluster = Cluster()
cluster.add_broker(Broker(1, "brokerhost1.example.com"))
cluster.add_broker(Broker(2, "brokerhost2.example.com"))
cluster.add_topic(Topic("testTopic1", 2))
cluster.add_topic(Topic("testTopic2", 2))
partition = cluster.topics['testTopic1'].partitions[0]
partition.add_replica(cluster.brokers[2], 0)
partition.add_replica(cluster.brokers[1], 1)
partition.size = 1001
partition = cluster.topics['testTopic1'].partitions[1]
partition.add_replica(cluster.brokers[1], 0)
partition.add_replica(cluster.brokers[2], 1)
partition.size = 1002
partition = cluster.topics['testTopic2'].partitions[0]
partition.add_replica(cluster.brokers[1], 0)
partition.add_replica(cluster.brokers[2], 1)
partition.size = 2001
partition = cluster.topics['testTopic2'].partitions[1]
partition.add_replica(cluster.brokers[1], 0)
partition.add_replica(cluster.brokers[2], 1)
partition.size = 2002
return cluster
class SizerSSHTests(unittest.TestCase):
def setUp(self):
self.args = Namespace()
self.args.datadir = "/path/to/data"
self.patcher_load_keys = patch('kafka.tools.assigner.sizers.ssh.paramiko.SSHClient.load_system_host_keys', autospec=True)
self.patcher_connect = patch('kafka.tools.assigner.sizers.ssh.paramiko.SSHClient.connect', autospec=True)
self.patcher_exec_command = patch('kafka.tools.assigner.sizers.ssh.paramiko.SSHClient.exec_command', autospec=True)
self.mock_load_keys = self.patcher_load_keys.start()
self.mock_connect = self.patcher_connect.start()
self.mock_exec_command = self.patcher_exec_command.start()
def tearDown(self):
self.patcher_load_keys.stop()
self.patcher_connect.stop()
self.patcher_exec_command.stop()
def create_cluster_onehost(self):
self.cluster = Cluster()
self.cluster.add_broker(Broker(1, "brokerhost1.example.com"))
self.cluster.add_topic(Topic("testTopic1", 2))
self.cluster.add_topic(Topic("testTopic2", 2))
partition = self.cluster.topics['testTopic1'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
partition = self.cluster.topics['testTopic1'].partitions[1]
partition.add_replica(self.cluster.brokers[1], 0)
partition = self.cluster.topics['testTopic2'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
partition = self.cluster.topics['testTopic2'].partitions[1]
partition.add_replica(self.cluster.brokers[1], 0)
def create_cluster_twohost(self):
self.cluster = Cluster()
self.cluster.add_broker(Broker(1, "brokerhost1.example.com"))
self.cluster.add_broker(Broker(2, "brokerhost2.example.com"))
self.cluster.add_topic(Topic("testTopic1", 2))
self.cluster.add_topic(Topic("testTopic2", 2))
partition = self.cluster.topics['testTopic1'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
partition.add_replica(self.cluster.brokers[2], 1)
partition = self.cluster.topics['testTopic1'].partitions[1]
partition.add_replica(self.cluster.brokers[2], 0)
partition.add_replica(self.cluster.brokers[1], 1)
partition = self.cluster.topics['testTopic2'].partitions[0]
partition.add_replica(self.cluster.brokers[2], 0)
partition.add_replica(self.cluster.brokers[1], 1)
partition = self.cluster.topics['testTopic2'].partitions[1]
partition.add_replica(self.cluster.brokers[1], 0)
partition.add_replica(self.cluster.brokers[2], 1)
def test_sizer_create(self):
self.create_cluster_onehost()
sizer = SizerSSH(self.args, self.cluster)
self.mock_load_keys.assert_called_with(ANY)
assert isinstance(sizer, SizerSSH)
def test_sizer_run_missing_host(self):
self.create_cluster_onehost()
self.cluster.brokers[1].hostname = None
sizer = SizerSSH(self.args, self.cluster)
self.assertRaises(UnknownBrokerException, sizer.get_partition_sizes)
def test_sizer_run_badinput(self):
self.create_cluster_onehost()
m_stdout = StringIO(u'foo\nbar\nbaz\n')
m_stderr = StringIO()
m_stdin = StringIO()
self.mock_exec_command.return_value = (m_stdin, m_stdout, m_stderr)
sizer = SizerSSH(self.args, self.cluster)
sizer.get_partition_sizes()
self.mock_connect.assert_called_once_with(ANY, "brokerhost1.example.com", allow_agent=True)
self.mock_exec_command.assert_called_once_with(ANY, "du -sk /path/to/data/*")
assert self.cluster.topics['testTopic1'].partitions[0].size == 0
assert self.cluster.topics['testTopic1'].partitions[1].size == 0
assert self.cluster.topics['testTopic2'].partitions[0].size == 0
assert self.cluster.topics['testTopic2'].partitions[1].size == 0
def test_sizer_run_setsizes_singlehost(self):
self.create_cluster_onehost()
m_stdout = StringIO(u'1001\t/path/to/data/testTopic1-0\n'
"1002\t/path/to/data/testTopic1-1\n"
"2001\t/path/to/data/testTopic2-0\n"
"2002\t/path/to/data/testTopic2-1\n")
m_stderr = StringIO()
m_stdin = StringIO()
self.mock_exec_command.return_value = (m_stdin, m_stdout, m_stderr)
sizer = SizerSSH(self.args, self.cluster)
sizer.get_partition_sizes()
self.mock_connect.assert_called_once_with(ANY, "brokerhost1.example.com", allow_agent=True)
self.mock_exec_command.assert_called_once_with(ANY, "du -sk /path/to/data/*")
assert self.cluster.topics['testTopic1'].partitions[0].size == 1001
assert self.cluster.topics['testTopic1'].partitions[1].size == 1002
assert self.cluster.topics['testTopic2'].partitions[0].size == 2001
assert self.cluster.topics['testTopic2'].partitions[1].size == 2002
def test_sizer_run_setsizes_twohost(self):
self.create_cluster_twohost()
m_stdout_1 = StringIO(u'1001\t/path/to/data/testTopic1-0\n'
u'1002\t/path/to/data/testTopic1-1\n'
u'2001\t/path/to/data/testTopic2-0\n'
u'2002\t/path/to/data/testTopic2-1\n')
m_stdout_2 = StringIO(u'1001\t/path/to/data/testTopic1-0\n'
u'4002\t/path/to/data/testTopic1-1\n'
u'1011\t/path/to/data/testTopic2-0\n'
u'2002\t/path/to/data/testTopic2-1\n')
m_stderr = StringIO()
m_stdin = StringIO()
self.mock_exec_command.side_effect = [(m_stdin, m_stdout_1, m_stderr), (m_stdin, m_stdout_2, m_stderr)]
sizer = SizerSSH(self.args, self.cluster)
sizer.get_partition_sizes()
self.mock_connect.assert_has_calls([call(ANY, "brokerhost1.example.com", allow_agent=True), call(ANY, "brokerhost2.example.com", allow_agent=True)])
self.mock_exec_command.assert_has_calls([call(ANY, "du -sk /path/to/data/*"), call(ANY, "du -sk /path/to/data/*")])
assert self.cluster.topics['testTopic1'].partitions[0].size == 1001
assert self.cluster.topics['testTopic1'].partitions[1].size == 4002
assert self.cluster.topics['testTopic2'].partitions[0].size == 2001
assert self.cluster.topics['testTopic2'].partitions[1].size == 2002
def test_sizer_run_setsizes_badtopic(self):
self.create_cluster_onehost()
m_stdout = StringIO(u'1001\t/path/to/data/testTopic1-0\n'
u'1002\t/path/to/data/testTopic1-1\n'
u'5002\t/path/to/data/badTopic1-1\n'
u'2001\t/path/to/data/testTopic2-0\n'
u'2002\t/path/to/data/testTopic2-1\n')
m_stderr = StringIO()
m_stdin = StringIO()
self.mock_exec_command.return_value = (m_stdin, m_stdout, m_stderr)
sizer = SizerSSH(self.args, self.cluster)
sizer.get_partition_sizes()
self.mock_connect.assert_called_once_with(ANY, "brokerhost1.example.com", allow_agent=True)
self.mock_exec_command.assert_called_once_with(ANY, "du -sk /path/to/data/*")
assert self.cluster.topics['testTopic1'].partitions[0].size == 1001
assert self.cluster.topics['testTopic1'].partitions[1].size == 1002
assert self.cluster.topics['testTopic2'].partitions[0].size == 2001
assert self.cluster.topics['testTopic2'].partitions[1].size == 2002
def test_sizer_run_setsizes_badpartition(self):
self.create_cluster_onehost()
m_stdout = StringIO(u'1001\t/path/to/data/testTopic1-0\n'
u'1002\t/path/to/data/testTopic1-1\n'
u'5002\t/path/to/data/testTopic1-2\n'
u'2001\t/path/to/data/testTopic2-0\n'
u'2002\t/path/to/data/testTopic2-1\n')
m_stderr = StringIO()
m_stdin = StringIO()
self.mock_exec_command.return_value = (m_stdin, m_stdout, m_stderr)
sizer = SizerSSH(self.args, self.cluster)
sizer.get_partition_sizes()
self.mock_connect.assert_called_once_with(ANY, "brokerhost1.example.com", allow_agent=True)
self.mock_exec_command.assert_called_once_with(ANY, "du -sk /path/to/data/*")
assert self.cluster.topics['testTopic1'].partitions[0].size == 1001
assert self.cluster.topics['testTopic1'].partitions[1].size == 1002
assert self.cluster.topics['testTopic2'].partitions[0].size == 2001
assert self.cluster.topics['testTopic2'].partitions[1].size == 2002
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
class PartitionOperationTests(unittest.TestCase):
def setUp(self):
self.cluster = Cluster()
self.cluster.add_broker(Broker(1, "brokerhost1.example.com"))
self.cluster.add_broker(Broker(2, "brokerhost2.example.com"))
self.cluster.add_broker(Broker(3, "brokerhost3.example.com"))
self.cluster.add_topic(Topic("testTopic1", 2))
self.cluster.add_topic(Topic("testTopic2", 2))
def add_topics(self):
partition = self.cluster.topics['testTopic1'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
partition.add_replica(self.cluster.brokers[2], 1)
partition = self.cluster.topics['testTopic1'].partitions[1]
partition.add_replica(self.cluster.brokers[2], 0)
partition.add_replica(self.cluster.brokers[1], 1)
partition = self.cluster.topics['testTopic2'].partitions[0]
partition.add_replica(self.cluster.brokers[2], 0)
partition.add_replica(self.cluster.brokers[1], 1)
partition = self.cluster.topics['testTopic2'].partitions[1]
partition.add_replica(self.cluster.brokers[1], 0)
partition.add_replica(self.cluster.brokers[2], 1)
def test_partition_add_broker_partition(self):
partition = self.cluster.topics['testTopic1'].partitions[0]
partition._add_broker_partition(0, self.cluster.brokers[1])
assert self.cluster.brokers[1].partitions[0] == [partition]
def test_partition_add_broker_partition_two(self):
partition = self.cluster.topics['testTopic1'].partitions[0]
partition._add_broker_partition(0, self.cluster.brokers[1])
partition2 = self.cluster.topics['testTopic2'].partitions[1]
partition2._add_broker_partition(0, self.cluster.brokers[1])
assert self.cluster.brokers[1].partitions[0] == [partition, partition2]
def test_partition_add_replica(self):
partition = self.cluster.topics['testTopic1'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
assert self.cluster.brokers[1].partitions[0] == [partition]
assert self.cluster.topics['testTopic1'].partitions[0].replicas == [self.cluster.brokers[1]]
def test_partition_add_replica_two(self):
partition = self.cluster.topics['testTopic1'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
partition.add_replica(self.cluster.brokers[2], 1)
assert self.cluster.brokers[1].partitions[0] == [partition]
assert self.cluster.brokers[2].partitions[1] == [partition]
assert self.cluster.topics['testTopic1'].partitions[0].replicas == [self.cluster.brokers[1], self.cluster.brokers[2]]
def test_partition_remove_broker_partition(self):
partition = self.cluster.topics['testTopic1'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
partition._remove_broker_partition(self.cluster.brokers[1])
assert self.cluster.brokers[1].partitions[0] == []
def test_partition_remove_broker_partition_two(self):
partition = self.cluster.topics['testTopic1'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
partition2 = self.cluster.topics['testTopic2'].partitions[1]
partition2.add_replica(self.cluster.brokers[1], 0)
partition._remove_broker_partition(self.cluster.brokers[1])
assert self.cluster.brokers[1].partitions[0] == [partition2]
def test_partition_remove_replica(self):
partition = self.cluster.topics['testTopic1'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
partition.remove_replica(self.cluster.brokers[1])
assert self.cluster.brokers[1].partitions[0] == []
assert self.cluster.topics['testTopic1'].partitions[0].replicas == []
def test_partition_remove_replica_single(self):
partition = self.cluster.topics['testTopic1'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
partition.add_replica(self.cluster.brokers[2], 1)
partition.remove_replica(self.cluster.brokers[1])
assert self.cluster.brokers[1].partitions[0] == []
assert self.cluster.brokers[2].partitions[1] == [partition]
assert self.cluster.topics['testTopic1'].partitions[0].replicas == [self.cluster.brokers[2]]
def test_partition_remove_replica_nonexistent(self):
self.assertRaises(ReplicaNotFoundException, self.cluster.topics['testTopic1'].partitions[0].remove_replica, self.cluster.brokers[1])
def test_partition_swap_replicas(self):
partition = self.cluster.topics['testTopic1'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
partition.swap_replicas(self.cluster.brokers[1], self.cluster.brokers[2])
assert self.cluster.brokers[2].partitions[0] == [partition]
assert self.cluster.topics['testTopic1'].partitions[0].replicas == [self.cluster.brokers[2]]
def test_partition_swap_replicas_nonexistent(self):
self.assertRaises(ReplicaNotFoundException,
self.cluster.topics['testTopic1'].partitions[0].swap_replicas,
self.cluster.brokers[1],
self.cluster.brokers[2])
def test_partition_swap_replica_positions(self):
partition = self.cluster.topics['testTopic1'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
partition.add_replica(self.cluster.brokers[2], 1)
partition.swap_replica_positions(self.cluster.brokers[1], self.cluster.brokers[2])
assert self.cluster.brokers[2].partitions[0] == [partition]
assert self.cluster.brokers[1].partitions[1] == [partition]
assert self.cluster.topics['testTopic1'].partitions[0].replicas == [self.cluster.brokers[2], self.cluster.brokers[1]]
def test_partition_swap_replica_positions_nonexistent(self):
partition = self.cluster.topics['testTopic1'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
partition.add_replica(self.cluster.brokers[2], 1)
self.assertRaises(ReplicaNotFoundException, partition.swap_replica_positions, self.cluster.brokers[1], self.cluster.brokers[3])
def test_cluster_clone(self):
# Should have a consistent cluster state
self.add_topics()
newcluster = self.cluster.clone()
assert self.cluster is not newcluster
for bid in newcluster.brokers:
assert newcluster.brokers[bid] == self.cluster.brokers[bid]
for tname in newcluster.topics:
assert newcluster.topics[tname] == self.cluster.topics[tname]
for partition in newcluster.partitions():
assert partition == self.cluster.topics[partition.topic.name].partitions[partition.num]
assert partition.replicas == self.cluster.topics[partition.topic.name].partitions[partition.num].replicas
def test_cluster_changed_partitions(self):
self.add_topics()
newcluster = self.cluster.clone()
newcluster.topics['testTopic1'].partitions[0].replicas.reverse()
difference = self.cluster.changed_partitions(newcluster)
assert difference == [newcluster.topics['testTopic1'].partitions[0]]
def test_cluster_changed_partitions_inconsistent(self):
self.add_topics()
badcluster = Cluster()
self.assertRaises(ClusterConsistencyException, badcluster.changed_partitions, self.cluster)
def setUp(self):
self.cluster = Cluster()
self.args = Namespace()
self.args.property = []
def setUp(self):
self.cluster = Cluster()
self.args = Namespace()
class PartitionOperationTests(unittest.TestCase):
def setUp(self):
self.cluster = Cluster()
self.cluster.add_broker(Broker(1, "brokerhost1.example.com"))
self.cluster.add_broker(Broker(2, "brokerhost2.example.com"))
self.cluster.add_broker(Broker(3, "brokerhost3.example.com"))
self.cluster.add_topic(Topic("testTopic1", 2))
self.cluster.add_topic(Topic("testTopic2", 2))
def add_topics(self):
partition = self.cluster.topics['testTopic1'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
partition.add_replica(self.cluster.brokers[2], 1)
partition = self.cluster.topics['testTopic1'].partitions[1]
partition.add_replica(self.cluster.brokers[2], 0)
partition.add_replica(self.cluster.brokers[1], 1)
partition = self.cluster.topics['testTopic2'].partitions[0]
partition.add_replica(self.cluster.brokers[2], 0)
partition.add_replica(self.cluster.brokers[1], 1)
partition = self.cluster.topics['testTopic2'].partitions[1]
partition.add_replica(self.cluster.brokers[1], 0)
partition.add_replica(self.cluster.brokers[2], 1)
def test_partition_add_broker_partition(self):
partition = self.cluster.topics['testTopic1'].partitions[0]
partition._add_broker_partition(0, self.cluster.brokers[1])
assert self.cluster.brokers[1].partitions[0] == [partition]
def test_partition_add_broker_partition_two(self):
partition = self.cluster.topics['testTopic1'].partitions[0]
partition._add_broker_partition(0, self.cluster.brokers[1])
partition2 = self.cluster.topics['testTopic2'].partitions[1]
partition2._add_broker_partition(0, self.cluster.brokers[1])
assert self.cluster.brokers[1].partitions[0] == [partition, partition2]
def test_partition_add_replica(self):
partition = self.cluster.topics['testTopic1'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
assert self.cluster.brokers[1].partitions[0] == [partition]
assert self.cluster.topics['testTopic1'].partitions[0].replicas == [self.cluster.brokers[1]]
def test_partition_add_replica_two(self):
partition = self.cluster.topics['testTopic1'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
partition.add_replica(self.cluster.brokers[2], 1)
assert self.cluster.brokers[1].partitions[0] == [partition]
assert self.cluster.brokers[2].partitions[1] == [partition]
assert self.cluster.topics['testTopic1'].partitions[0].replicas == [self.cluster.brokers[1], self.cluster.brokers[2]]
def test_partition_remove_broker_partition(self):
partition = self.cluster.topics['testTopic1'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
partition._remove_broker_partition(self.cluster.brokers[1])
assert self.cluster.brokers[1].partitions[0] == []
def test_partition_remove_broker_partition_two(self):
partition = self.cluster.topics['testTopic1'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
partition2 = self.cluster.topics['testTopic2'].partitions[1]
partition2.add_replica(self.cluster.brokers[1], 0)
partition._remove_broker_partition(self.cluster.brokers[1])
assert self.cluster.brokers[1].partitions[0] == [partition2]
def test_partition_remove_replica(self):
partition = self.cluster.topics['testTopic1'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
partition.remove_replica(self.cluster.brokers[1])
assert self.cluster.brokers[1].partitions[0] == []
assert self.cluster.topics['testTopic1'].partitions[0].replicas == []
def test_partition_remove_replica_single(self):
partition = self.cluster.topics['testTopic1'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
partition.add_replica(self.cluster.brokers[2], 1)
partition.remove_replica(self.cluster.brokers[1])
assert self.cluster.brokers[1].partitions[0] == []
assert self.cluster.brokers[2].partitions[1] == [partition]
assert self.cluster.topics['testTopic1'].partitions[0].replicas == [self.cluster.brokers[2]]
def test_partition_remove_replica_nonexistent(self):
self.assertRaises(ReplicaNotFoundException, self.cluster.topics['testTopic1'].partitions[0].remove_replica, self.cluster.brokers[1])
def test_partition_swap_replicas(self):
partition = self.cluster.topics['testTopic1'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
partition.swap_replicas(self.cluster.brokers[1], self.cluster.brokers[2])
assert self.cluster.brokers[2].partitions[0] == [partition]
assert self.cluster.topics['testTopic1'].partitions[0].replicas == [self.cluster.brokers[2]]
def test_partition_swap_replicas_nonexistent(self):
self.assertRaises(ReplicaNotFoundException,
self.cluster.topics['testTopic1'].partitions[0].swap_replicas,
self.cluster.brokers[1],
self.cluster.brokers[2])
def test_partition_swap_replica_positions(self):
partition = self.cluster.topics['testTopic1'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
partition.add_replica(self.cluster.brokers[2], 1)
partition.swap_replica_positions(self.cluster.brokers[1], self.cluster.brokers[2])
assert self.cluster.brokers[2].partitions[0] == [partition]
assert self.cluster.brokers[1].partitions[1] == [partition]
assert self.cluster.topics['testTopic1'].partitions[0].replicas == [self.cluster.brokers[2], self.cluster.brokers[1]]
def test_partition_swap_replica_positions_nonexistent(self):
partition = self.cluster.topics['testTopic1'].partitions[0]
partition.add_replica(self.cluster.brokers[1], 0)
partition.add_replica(self.cluster.brokers[2], 1)
self.assertRaises(ReplicaNotFoundException, partition.swap_replica_positions, self.cluster.brokers[1], self.cluster.brokers[3])
def test_cluster_clone(self):
# Should have a consistent cluster state
self.add_topics()
newcluster = self.cluster.clone()
assert self.cluster is not newcluster
for bid in newcluster.brokers:
assert newcluster.brokers[bid] == self.cluster.brokers[bid]
for tname in newcluster.topics:
assert newcluster.topics[tname] == self.cluster.topics[tname]
for partition in newcluster.partitions([]):
assert partition == self.cluster.topics[partition.topic.name].partitions[partition.num]
assert partition.replicas == self.cluster.topics[partition.topic.name].partitions[partition.num].replicas
def test_cluster_changed_partitions(self):
self.add_topics()
newcluster = self.cluster.clone()
newcluster.topics['testTopic1'].partitions[0].replicas.reverse()
difference = self.cluster.changed_partitions(newcluster)
assert difference == [newcluster.topics['testTopic1'].partitions[0]]
def test_cluster_changed_partitions_inconsistent(self):
self.add_topics()
badcluster = Cluster()
self.assertRaises(ClusterConsistencyException, badcluster.changed_partitions, self.cluster)
def test_cluster_changed_partitions_inconsistent(self):
self.add_topics()
badcluster = Cluster()
self.assertRaises(ClusterConsistencyException, badcluster.changed_partitions, self.cluster)
class SimpleClusterTests(unittest.TestCase):
def setUp(self):
self.cluster = Cluster()
def add_brokers(self, num):
for i in range(1, num + 1):
broker = Broker(i, "brokerhost{0}.example.com".format(i))
self.cluster.add_broker(broker)
def add_topics(self, num, partition_count=2):
for i in range(1, num + 1):
topic = Topic("testTopic{0}".format(i), partition_count)
self.cluster.add_topic(topic)
def add_partitions_to_broker(self, broker_id, pos, num):
topic = Topic('testTopic', num)
self.cluster.brokers[broker_id].partitions[pos] = []
for i in range(num):
self.cluster.brokers[broker_id].partitions[pos].append(
topic.partitions[i])
def test_cluster_create(self):
assert self.cluster.brokers == {}
assert self.cluster.topics == {}
def test_cluster_add_broker(self):
self.add_brokers(1)
assert len(self.cluster.brokers) == 1
for bid in self.cluster.brokers.keys():
assert self.cluster.brokers[bid].cluster is self.cluster
def test_cluster_num_brokers(self):
self.add_brokers(2)
assert self.cluster.num_brokers() == 2
def test_cluster_add_topic(self):
self.add_topics(1)
assert len(self.cluster.topics) == 1
for tname in self.cluster.topics.keys():
assert self.cluster.topics[tname].cluster is self.cluster
def test_cluster_num_topics(self):
self.add_topics(2)
assert self.cluster.num_topics() == 2
def test_cluster_partition_iterator(self):
self.add_topics(2)
seen_partitions = {}
for partition in self.cluster.partitions():
seen_partitions["{0}:{1}".format(partition.topic.name,
partition.num)] = 1
assert seen_partitions == {
'testTopic1:0': 1,
'testTopic1:1': 1,
'testTopic2:0': 1,
'testTopic2:1': 1
}
def test_cluster_max_replication_factor(self):
self.add_brokers(2)
self.add_partitions_to_broker(1, 0, 1)
self.add_partitions_to_broker(1, 1, 1)
self.add_partitions_to_broker(2, 2, 2)
assert self.cluster.max_replication_factor() == 3
def test_cluster_log_info(self):
self.add_brokers(2)
with LogCapture() as l:
self.cluster.log_broker_summary()
l.check(('kafka-assigner', 'INFO',
'Broker 1: partitions=0/0 (0.00%), size=0'),
('kafka-assigner', 'INFO',
'Broker 2: partitions=0/0 (0.00%), size=0'))
def setUp(self):
self.cluster = Cluster()
class SimpleClusterTests(unittest.TestCase):
def setUp(self):
self.cluster = Cluster()
def add_brokers(self, num):
for i in range(1, num + 1):
broker = Broker(i, "brokerhost{0}.example.com".format(i))
self.cluster.add_broker(broker)
def add_topics(self, num, partition_count=2):
for i in range(1, num + 1):
topic = Topic("testTopic{0}".format(i), partition_count)
self.cluster.add_topic(topic)
def add_partitions_to_broker(self, broker_id, pos, num):
topic = Topic('testTopic', num)
self.cluster.brokers[broker_id].partitions[pos] = []
for i in range(num):
self.cluster.brokers[broker_id].partitions[pos].append(topic.partitions[i])
def test_cluster_create(self):
assert self.cluster.brokers == {}
assert self.cluster.topics == {}
def test_cluster_add_broker(self):
self.add_brokers(1)
assert len(self.cluster.brokers) == 1
for bid in self.cluster.brokers.keys():
assert self.cluster.brokers[bid].cluster is self.cluster
def test_cluster_num_brokers(self):
self.add_brokers(2)
assert self.cluster.num_brokers() == 2
def test_cluster_add_topic(self):
self.add_topics(1)
assert len(self.cluster.topics) == 1
for tname in self.cluster.topics.keys():
assert self.cluster.topics[tname].cluster is self.cluster
def test_cluster_num_topics(self):
self.add_topics(2)
assert self.cluster.num_topics() == 2
def test_cluster_partition_iterator(self):
self.add_topics(2)
seen_partitions = {}
for partition in self.cluster.partitions([]):
seen_partitions["{0}:{1}".format(partition.topic.name, partition.num)] = 1
assert seen_partitions == {'testTopic1:0': 1, 'testTopic1:1': 1, 'testTopic2:0': 1, 'testTopic2:1': 1}
def test_cluster_partition_iterator_with_exclude(self):
self.add_topics(2)
seen_partitions = {}
for partition in self.cluster.partitions(['testTopic1']):
seen_partitions["{0}:{1}".format(partition.topic.name, partition.num)] = 1
assert seen_partitions == {'testTopic2:0': 1, 'testTopic2:1': 1}
def test_cluster_max_replication_factor(self):
self.add_brokers(2)
self.add_partitions_to_broker(1, 0, 1)
self.add_partitions_to_broker(1, 1, 1)
self.add_partitions_to_broker(2, 2, 2)
assert self.cluster.max_replication_factor() == 3
def test_cluster_log_info(self):
self.add_brokers(2)
with LogCapture() as l:
self.cluster.log_broker_summary()
l.check(('kafka-assigner', 'INFO', 'Broker 1: partitions=0/0 (0.00%), size=0'),
('kafka-assigner', 'INFO', 'Broker 2: partitions=0/0 (0.00%), size=0'))
def setUp(self):
self.cluster = Cluster()
def main():
# Start by loading all the modules
action_map = get_action_map()
sizer_map = get_sizer_map()
plugins_list = get_plugins_list()
# Instantiate all plugins
plugins = [plugin() for plugin in plugins_list]
# Set up and parse all CLI arguments
args = set_up_arguments(action_map, sizer_map, plugins)
for plugin in plugins:
plugin.set_arguments(args)
tools_path = get_tools_path(args.tools_path)
check_java_home()
cluster = Cluster.create_from_zookeeper(args.zookeeper)
for plugin in plugins:
plugin.set_cluster(cluster)
# If the module needs the partition sizes, call a size module to get the information
check_and_get_sizes(action_map[args.action], args, cluster, sizer_map)
for plugin in plugins:
plugin.after_sizes()
if args.leadership:
log.info("Cluster Leadership Balance (before):")
cluster.log_broker_summary()
# Clone the cluster, and run the action to generate a new cluster state
newcluster = cluster.clone()
action_to_run = action_map[args.action](args, newcluster)
action_to_run.process_cluster()
for plugin in plugins:
plugin.set_new_cluster(action_to_run.cluster)
if args.leadership:
log.info("Cluster Leadership Balance (after):")
newcluster.log_broker_summary()
move_partitions = cluster.changed_partitions(action_to_run.cluster)
batches = split_partitions_into_batches(move_partitions, batch_size=args.moves, use_class=Reassignment)
for plugin in plugins:
plugin.set_batches(batches)
log.info("Partition moves required: {0}".format(len(move_partitions)))
log.info("Number of batches: {0}".format(len(batches)))
dry_run = args.generate or not args.execute
if dry_run:
log.info("--execute flag NOT specified. DRY RUN ONLY")
for i, batch in enumerate(batches):
log.info("Executing partition reassignment {0}/{1}: {2}".format(i + 1, len(batches), repr(batch)))
batch.execute(i + 1, len(batches), args.zookeeper, tools_path, plugins, dry_run)
for plugin in plugins:
plugin.before_ple()
if not args.skip_ple:
batches = split_partitions_into_batches(move_partitions, batch_size=args.moves, use_class=ReplicaElection)
log.info("Number of replica elections: {0}".format(len(batches)))
run_preferred_replica_elections(batches, args, tools_path, plugins, dry_run)
for plugin in plugins:
plugin.finished()
return 0
