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'