def test_rebalance_empty_brokers_and_exclude_topics(self):
distribution = {
('t0', '0'): ['1', '2'],
('t0', '1'): ['2', '3'],
('t1', '0'): ['2', '3'],
('t1', '1'): ['3', '4'],
('t1', '2'): ['4', '5'],
('t2', '0'): ['3', '4'],
('t2', '1'): ['4', '5'],
('t2', '2'): ['5', '6'],
}
brokers, zk = self._create_zk_for_topics(distribution)
o = OptimizedRebalanceChange(zk, brokers, ['2', '3'], ['t1'])
while o.run([]):
pass
assert distribution[('t1', '0')] == ['2', '3']
assert distribution[('t1', '1')] == ['3', '4']
assert distribution[('t1', '2')] == ['4', '5']
distribution.pop(('t1', '0'))
distribution.pop(('t1', '1'))
distribution.pop(('t1', '2'))
brokers = [item for sublist in distribution.values() for item in sublist]
assert '2' not in brokers
assert '3' not in brokers
def test_rebalance_empty_brokers_and_exclude_topics(self):
distribution = {
('t0', '0'): ['1', '2'],
('t0', '1'): ['2', '3'],
('t1', '0'): ['2', '3'],
('t1', '1'): ['3', '4'],
('t1', '2'): ['4', '5'],
('t2', '0'): ['3', '4'],
('t2', '1'): ['4', '5'],
('t2', '2'): ['5', '6'],
}
brokers, zk = self._create_zk_for_topics(distribution)
o = OptimizedRebalanceChange(zk, brokers, ['2', '3'], ['t1'])
while o.run([]):
pass
assert distribution[('t1', '0')] == ['2', '3']
assert distribution[('t1', '1')] == ['3', '4']
assert distribution[('t1', '2')] == ['4', '5']
distribution.pop(('t1', '0'))
distribution.pop(('t1', '1'))
distribution.pop(('t1', '2'))
brokers = [
item for sublist in distribution.values() for item in sublist
]
assert '2' not in brokers
assert '3' not in brokers
def test_rebalance_can_run(self):
o = OptimizedRebalanceChange(object(), [])
blocked_actions = ['restart', 'start', 'stop', 'rebalance']
# Check that can run in exact cases
for a in blocked_actions:
assert not o.can_run([a])
assert o.can_run(['xxx'])
assert o.can_run([])
def test_rebalance_recovered_with_additional_copy2(self):
distribution = {
('t0', '0'): ['2', '1'],
('t0', '1'): ['1', '2'],
('t0', '2'): ['3', '4']
}
_, zk = _create_zk_for_topics(distribution, ['1', '2', '4'])
o = OptimizedRebalanceChange(zk, ['1', '2', '4'])
while o.run([]):
pass
_verify_balanced(['1', '2', '4'], distribution)
def test_leader_partition_limit(self):
distribution = {
('t0', '0'): ['1', '2'],
('t0', '1'): ['1', '2'],
('t0', '2'): ['1', '2'],
('t1', '2'): ['1', '2'],
}
_, zk = _create_zk_for_topics(distribution, ['2', '3'])
o = OptimizedRebalanceChange(zk, ['2', '3'])
while o.run([]):
pass
_verify_balanced(['2', '3'], distribution)
def test_rebalance_empty_one_broker(self):
distribution = {
('t0', '0'): ['1', '2'],
('t0', '1'): ['2', '3'],
('t1', '0'): ['2', '3'],
('t1', '1'): ['3', '4'],
}
brokers, zk = _create_zk_for_topics(distribution)
o = OptimizedRebalanceChange(zk, brokers, ['2'], [])
while o.run([]):
pass
_verify_empty_brokers(('2'), distribution)
def test_rebalance_on_filled1(self):
distribution = {
('t0', '0'): ['2'],
('t0', '1'): ['1'],
('t0', '2'): ['1'],
('t0', '3'): ['1'],
}
brokers, zk = _create_zk_for_topics(distribution)
o = OptimizedRebalanceChange(zk, brokers)
# broker to partitions
while o.run([]):
pass
_verify_balanced(('1', '2'), distribution)
def test_rebalance_with_many_topics(self):
distribution = {}
topic_count = 1000
partition_count = 8
for i in range(0, topic_count):
topic = 't{}'.format(i)
distribution.update({(topic, str(partition)): ['1', '2', '3'] for partition in range(0, partition_count)})
_, zk = _create_zk_for_topics(distribution, broker_ids=['1', '2', '3', '4', '5'])
o = OptimizedRebalanceChange(zk, ['1', '2', '3', '4', '5'])
steps = 0
while o.run([]):
steps += 1
_verify_balanced(['1', '2', '3', '4', '5'], distribution)
def test_rebalance_fail_with_not_enough_replicas(self):
distribution = {
('t0', '0'): ['2', '1', '3'],
('t0', '1'): ['1', '2'],
}
_, zk = _create_zk_for_topics(distribution, broker_ids=['1', '3'])
o = OptimizedRebalanceChange(zk, ['1', '3'])
try:
while o.run([]):
pass
assert False, "Balancing can not work with low replication factor"
except Exception:
pass
def test_rebalance_with_dead_brokers(self):
distribution = {
('t0', '0'): ['2', '1'],
('t0', '1'): ['1', '2'],
('t0', '2'): ['1', '2'],
('t0', '3'): ['1', '2'],
('t0', '4'): ['1', '2'],
('t0', '5'): ['1', '2'],
('t0', '6'): ['1', '2'],
}
_, zk = _create_zk_for_topics(distribution, broker_ids=['1', '3'])
o = OptimizedRebalanceChange(zk, ['1', '3'])
while o.run([]):
pass
_verify_balanced(['1', '3'], distribution)
def check(self) -> Change:
with self.zk.lock():
data = self.zk.take_action(
self.broker_manager.id_manager.detect_broker_id())
if not data:
return None
if 'name' not in data:
_LOG.error(
'Action name can not be restored from {}, skipping'.format(
data))
return None
try:
if data['name'] == 'restart':
return RestartBrokerChange(self.zk, self.broker_manager,
lambda: False)
elif data['name'] == 'rebalance':
return OptimizedRebalanceChange(self.zk,
self.zk.get_broker_ids())
elif data['name'] == 'migrate':
return MigrationChange(self.zk, data['from'], data['to'],
data['shrink'])
elif data['name'] == 'fatboyslim':
return SwapPartitionsChange(
self.zk, lambda x: load_swap_data(
x, self.api_port, int(data['threshold_kb'])))
else:
_LOG.error('Action {} not supported'.format(data))
except Exception as e:
_LOG.error('Failed to create action from {}'.format(data),
exc_info=e)
return None
def createChange(self,
zk,
broker_ids,
empty_brokers,
exclude_topics,
parallelism=1):
return OptimizedRebalanceChange(zk, broker_ids, empty_brokers,
exclude_topics, parallelism)
def check(self):
if self.executed:
return None
if not self.broker.is_running_and_registered():
return None
_LOG.info("Rebalance on start, triggering rebalance")
self.executed = True
return OptimizedRebalanceChange(self.zk, self.zk.get_broker_ids())
def check(self):
if not self.broker.is_running_and_registered():
return None
new_list = self.zk.get_broker_ids()
if not new_list == self.old_broker_list:
_LOG.info(
'Broker list changed from {} to {}, triggering rebalance'.
format(self.old_broker_list, new_list))
self.old_broker_list = new_list
return OptimizedRebalanceChange(self.zk, new_list, [], [])
return None
def test_rebalance_on_empty1(self):
brokers, zk = _create_zk_for_topics({})
o = OptimizedRebalanceChange(zk, brokers)
while o.run([]):
pass
def test_rebalance_get_name(self):
o = OptimizedRebalanceChange(object(), [])
assert o.get_name() == 'rebalance'
