def test_revoke_leadership_not_possible(
self,
create_balancer,
create_cluster_topology,
):
assignment = dict([
((u'T0', 0), ['2', '3']),
((u'T1', 0), ['3']),
((u'T1', 1), ['0', '1']),
((u'T2', 0), ['1', '0']),
])
ct = create_cluster_topology(assignment, broker_range(4))
cb = create_balancer(ct)
cb.revoke_leadership(['2', '3'])
new_leaders_per_broker = {
broker.id: broker.count_preferred_replica()
for broker in six.itervalues(ct.brokers)
}
# Broker '2' and '3' are to be revoked from leadership
# But (u'T0, 0) has replicas '2' and '3'
# and (u'T1', 0) has only single replica
# Assert no leadership changes
assert new_leaders_per_broker['0'] == 1
assert new_leaders_per_broker['1'] == 1
assert new_leaders_per_broker['2'] == 1
assert new_leaders_per_broker['3'] == 1
# Assert no partition movements
_, total_movements = \
calculate_partition_movement(assignment, ct.assignment)
assert total_movements == 0
def test__rebalance_groups_partition_cnt_case1(
self,
create_balancer,
create_cluster_topology,
):
# rg1 has 6 partitions
# rg2 has 2 partitions
# Both rg's are balanced(based on replica-count) initially
# Result: rg's will be balanced for partition-count
assignment = dict([
((u'T1', 1), ['0', '1', '2']),
((u'T1', 0), ['1']),
((u'T3', 0), ['1']),
((u'T2', 0), ['0', '1', '3']),
])
ct = create_cluster_topology(assignment, broker_range(4))
cb = create_balancer(ct)
# Re-balance replication-groups for partition-count
cb._rebalance_groups_partition_cnt()
# Verify both replication-groups have same partition-count
assert len(ct.rgs['rg1'].partitions) == len(ct.rgs['rg2'].partitions)
_, total_movements = \
calculate_partition_movement(assignment, ct.assignment)
# Verify minimum partition movements 2
assert total_movements == 2
net_imbal, _ = get_replication_group_imbalance_stats(
list(ct.rgs.values()),
list(ct.partitions.values()),
)
# Verify replica-count imbalance remains unaltered
assert net_imbal == 0
def test_revoke_leadership_single_broker(
self,
create_balancer,
create_cluster_topology,
):
assignment = dict([
((u'T0', 0), ['2', '0']),
((u'T1', 0), ['2', '1']),
((u'T1', 1), ['0', '2']),
((u'T2', 0), ['1', '0']),
])
ct = create_cluster_topology(assignment, broker_range(3))
cb = create_balancer(ct)
cb.revoke_leadership(['2'])
new_leaders_per_broker = {
broker.id: broker.count_preferred_replica()
for broker in six.itervalues(ct.brokers)
}
_, total_movements = \
calculate_partition_movement(assignment, ct.assignment)
# Get net imbalance statistics excluding brokers to be revoked
# leadership from
brokers = [
b for b in six.itervalues(ct.brokers) if b.id not in ['2', '3']
]
new_net_imbal = get_net_imbalance(get_broker_leader_counts(brokers))
# Verify that broker '2' is not leader of any partition
assert new_leaders_per_broker['2'] == 0
assert new_leaders_per_broker['1'] == 2
assert new_leaders_per_broker['0'] == 2
# Assert no partition movements
assert total_movements == 0
# Assert remaining brokers are balanced with leader count
assert new_net_imbal == 0
def test_revoke_leadership_not_possible(
self,
create_balancer,
create_cluster_topology,
):
assignment = dict(
[
((u'T0', 0), ['2', '3']),
((u'T1', 0), ['3']),
((u'T1', 1), ['0', '1']),
((u'T2', 0), ['1', '0']),
]
)
ct = create_cluster_topology(assignment, broker_range(4))
cb = create_balancer(ct)
cb.revoke_leadership(['2', '3'])
new_leaders_per_broker = {
broker.id: broker.count_preferred_replica()
for broker in six.itervalues(ct.brokers)
}
# Broker '2' and '3' are to be revoked from leadership
# But (u'T0, 0) has replicas '2' and '3'
# and (u'T1', 0) has only single replica
# Assert no leadership changes
assert new_leaders_per_broker['0'] == 1
assert new_leaders_per_broker['1'] == 1
assert new_leaders_per_broker['2'] == 1
assert new_leaders_per_broker['3'] == 1
# Assert no partition movements
_, total_movements = \
calculate_partition_movement(assignment, ct.assignment)
assert total_movements == 0
def test__rebalance_groups_partition_cnt_case1(self):
# rg1 has 6 partitions
# rg2 has 2 partitions
# Both rg's are balanced(based on replica-count) initially
# Result: rg's will be balanced for partition-count
assignment = dict(
[
((u'T1', 1), ['0', '1', '2']),
((u'T1', 0), ['1']),
((u'T3', 0), ['1']),
((u'T2', 0), ['0', '1', '3']),
]
)
ct = self.build_cluster_topology(assignment, self.srange(4))
# Re-balance replication-groups for partition-count
ct._rebalance_groups_partition_cnt()
# Verify both replication-groups have same partition-count
assert len(ct.rgs['rg1'].partitions) == len(ct.rgs['rg2'].partitions)
_, total_movements = \
calculate_partition_movement(assignment, ct.assignment)
# Verify minimum partition movements 2
assert total_movements == 2
net_imbal, _ = get_replication_group_imbalance_stats(
ct.rgs.values(),
ct.partitions.values(),
)
# Verify replica-count imbalance remains unaltered
assert net_imbal == 0
def test__rebalance_groups_partition_cnt_case6(
self,
create_balancer,
create_cluster_topology,
):
# rg1 has 5 partitions
# rg2 has 1 partitions
# rg3 has 1 partitions
# Result: rg's will be balanced for partition-count
# All rg's will be balanced with 2 partition-movements
# This test case covers the aspect that even if the partition
# count difference b/w the replication-groups is > 1,
# we still move onto next replication-group if either of the
# replication-groups reaches the optimal partition-count.
brokers = {
"0": {
"host": "host1"
},
"1": {
"host": "host2"
},
"2": {
"host": "host3"
},
"3": {
"host": "host4"
},
"5": {
"host": "host5"
},
}
assignment = dict([
((u'T0', 0), ['0', '2']),
((u'T1', 0), ['1', '0']),
((u'T2', 0), ['0', '5']),
((u'T3', 0), ['1']),
])
ct = create_cluster_topology(assignment, brokers)
cb = create_balancer(ct)
# Re-balance replication-groups for partition-count
cb._rebalance_groups_partition_cnt()
# Assert final partition counts in replication-groups
assert len(ct.rgs['rg1'].partitions) == 3
assert len(ct.rgs['rg2'].partitions) == 2
assert len(ct.rgs['rg3'].partitions) == 2
_, total_movements = \
calculate_partition_movement(assignment, ct.assignment)
# Verify minimum partition movements 2
assert total_movements == 2
def test__rebalance_groups_partition_cnt_case5(
self,
create_balancer,
create_cluster_topology,
):
# rg1 has 4 partitions
# rg2 has 2 partitions
# rg3 has 2 partitions
# Result: rg's will be balanced for partition-count
# All rg's will be balanced with just 1 partition-movement
brokers = {
"0": {
"host": "host1"
},
"1": {
"host": "host2"
},
"2": {
"host": "host3"
},
"3": {
"host": "host4"
},
"5": {
"host": "host5"
},
}
assignment = dict([
((u'T0', 0), ['0', '2']),
((u'T1', 0), ['1', '3']),
((u'T2', 0), ['0', '5']),
((u'T3', 0), ['1', '5']),
])
ct = create_cluster_topology(assignment, brokers)
cb = create_balancer(ct)
# Re-balance replication-groups for partition-count
cb._rebalance_groups_partition_cnt()
# Assert partition is moved from rg1 only
assert len(ct.rgs['rg1'].partitions) == 3
_, total_movements = \
calculate_partition_movement(assignment, ct.assignment)
# Verify minimum partition movements 1
assert total_movements == 1
net_imbal, _ = get_replication_group_imbalance_stats(
list(ct.rgs.values()),
list(ct.partitions.values()),
)
# Verify replica-count imbalance remains unaltered
assert net_imbal == 0
def test__rebalance_groups_partition_cnt_case6(
self,
create_balancer,
create_cluster_topology,
):
# rg1 has 5 partitions
# rg2 has 1 partitions
# rg3 has 1 partitions
# Result: rg's will be balanced for partition-count
# All rg's will be balanced with 2 partition-movements
# This test case covers the aspect that even if the partition
# count difference b/w the replication-groups is > 1,
# we still move onto next replication-group if either of the
# replication-groups reaches the optimal partition-count.
brokers = {
"0": {"host": "host1"},
"1": {"host": "host2"},
"2": {"host": "host3"},
"3": {"host": "host4"},
"5": {"host": "host5"},
}
assignment = dict(
[
((u'T0', 0), ['0', '2']),
((u'T1', 0), ['1', '0']),
((u'T2', 0), ['0', '5']),
((u'T3', 0), ['1']),
]
)
ct = create_cluster_topology(assignment, brokers)
cb = create_balancer(ct)
# Re-balance replication-groups for partition-count
cb._rebalance_groups_partition_cnt()
# Assert final partition counts in replication-groups
assert len(ct.rgs['rg1'].partitions) == 3
assert len(ct.rgs['rg2'].partitions) == 2
assert len(ct.rgs['rg3'].partitions) == 2
_, total_movements = \
calculate_partition_movement(assignment, ct.assignment)
# Verify minimum partition movements 2
assert total_movements == 2
def test__rebalance_groups_partition_cnt_case5(
self,
create_balancer,
create_cluster_topology,
):
# rg1 has 4 partitions
# rg2 has 2 partitions
# rg3 has 2 partitions
# Result: rg's will be balanced for partition-count
# All rg's will be balanced with just 1 partition-movement
brokers = {
"0": {"host": "host1"},
"1": {"host": "host2"},
"2": {"host": "host3"},
"3": {"host": "host4"},
"5": {"host": "host5"},
}
assignment = dict(
[
((u'T0', 0), ['0', '2']),
((u'T1', 0), ['1', '3']),
((u'T2', 0), ['0', '5']),
((u'T3', 0), ['1', '5']),
]
)
ct = create_cluster_topology(assignment, brokers)
cb = create_balancer(ct)
# Re-balance replication-groups for partition-count
cb._rebalance_groups_partition_cnt()
# Assert partition is moved from rg1 only
assert len(ct.rgs['rg1'].partitions) == 3
_, total_movements = \
calculate_partition_movement(assignment, ct.assignment)
# Verify minimum partition movements 1
assert total_movements == 1
net_imbal, _ = get_replication_group_imbalance_stats(
list(ct.rgs.values()),
list(ct.partitions.values()),
)
# Verify replica-count imbalance remains unaltered
assert net_imbal == 0
def test__rebalance_groups_partition_cnt_case3(
self,
create_balancer,
create_cluster_topology,
):
# 1 over-balanced, 1 under-balanced, 1 opt-balanced replication-group
# rg1 has 3 partitions
# rg2 has 2 partitions
# rg3 has 1 partition
# All rg's are balanced(based on replica-count) initially
# Result: rg's will be balanced for partition-count
assignment = dict(
[
((u'T1', 1), ['0', '2']),
((u'T3', 1), ['2']),
((u'T3', 0), ['0']),
((u'T2', 0), ['0', '5']),
]
)
brokers = {
'0': mock.MagicMock(),
'2': mock.MagicMock(),
'5': mock.MagicMock(),
}
ct = create_cluster_topology(assignment, brokers)
cb = create_balancer(ct)
# Re-balance brokers across replication-groups
cb._rebalance_groups_partition_cnt()
# Verify all replication-groups have same partition-count
assert len(ct.rgs['rg1'].partitions) == len(ct.rgs['rg2'].partitions)
assert len(ct.rgs['rg1'].partitions) == len(ct.rgs['rg3'].partitions)
_, total_movements = \
calculate_partition_movement(assignment, ct.assignment)
# Verify minimum partition movements
assert total_movements == 1
net_imbal, _ = get_replication_group_imbalance_stats(
list(ct.rgs.values()),
list(ct.partitions.values()),
)
# Verify replica-count imbalance remains 0
assert net_imbal == 0
def test__rebalance_groups_partition_cnt_case3(
self,
create_balancer,
create_cluster_topology,
):
# 1 over-balanced, 1 under-balanced, 1 opt-balanced replication-group
# rg1 has 3 partitions
# rg2 has 2 partitions
# rg3 has 1 partition
# All rg's are balanced(based on replica-count) initially
# Result: rg's will be balanced for partition-count
assignment = dict(
[
((u'T1', 1), ['0', '2']),
((u'T3', 1), ['2']),
((u'T3', 0), ['0']),
((u'T2', 0), ['0', '5']),
]
)
brokers = {
'0': mock.MagicMock(),
'2': mock.MagicMock(),
'5': mock.MagicMock(),
}
ct = create_cluster_topology(assignment, brokers)
cb = create_balancer(ct)
# Re-balance brokers across replication-groups
cb._rebalance_groups_partition_cnt()
# Verify all replication-groups have same partition-count
assert len(ct.rgs['rg1'].partitions) == len(ct.rgs['rg2'].partitions)
assert len(ct.rgs['rg1'].partitions) == len(ct.rgs['rg3'].partitions)
_, total_movements = \
calculate_partition_movement(assignment, ct.assignment)
# Verify minimum partition movements
assert total_movements == 1
net_imbal, _ = get_replication_group_imbalance_stats(
ct.rgs.values(),
ct.partitions.values(),
)
# Verify replica-count imbalance remains 0
assert net_imbal == 0
def test_revoke_leadership_multiple_brokers(
self,
create_balancer,
create_cluster_topology,
):
assignment = dict(
[
((u'T0', 0), ['2', '0']),
((u'T1', 0), ['2', '1']),
((u'T1', 1), ['0', '2']),
((u'T2', 0), ['0', '3']),
((u'T3', 0), ['3', '1']),
((u'T3', 1), ['3', '1']),
]
)
ct = create_cluster_topology(assignment, broker_range(4))
cb = create_balancer(ct)
cb.revoke_leadership(['2', '3'])
new_leaders_per_broker = {
broker.id: broker.count_preferred_replica()
for broker in six.itervalues(ct.brokers)
}
_, total_movements = \
calculate_partition_movement(assignment, ct.assignment)
# Get net imbalance statistics excluding brokers to be revoked
# leadership from
brokers = [
b for b in six.itervalues(ct.brokers)
if b.id not in ['2', '3']
]
new_net_imbal = get_net_imbalance(get_broker_leader_counts(brokers))
# Verify that broker '2' and '3' is not leader of any partition
assert new_leaders_per_broker['0'] == 3
assert new_leaders_per_broker['1'] == 3
assert new_leaders_per_broker['2'] == 0
assert new_leaders_per_broker['3'] == 0
# Assert no partition movements
assert total_movements == 0
# Assert remaining brokers are balanced with leader count
assert new_net_imbal == 0