def test__select_broker_pair(self, create_partition):
p10 = create_partition('t1', 0)
p11 = create_partition('t1', 1)
p20 = create_partition('t2', 0)
b0 = create_broker('b0', [p10, p20, p11])
b1 = create_broker('b1', [p10, p20])
b2 = create_broker('b2', [p20, p11])
b3 = create_broker('b3', [p10, p20, p11])
rg_source = ReplicationGroup('rg1', set([b0, b1, b2, b3]))
b4 = create_broker('b4', [p20, p11])
b5 = create_broker('b5', [p20])
b6 = create_broker('b6', [p10])
b7 = create_broker('b7', [p11])
rg_dest = ReplicationGroup('rg2', set([b4, b5, b6, b7]))
# Select best-suitable brokers for moving partition 'p10'
broker_source, broker_dest = rg_source._select_broker_pair(
rg_dest, p10)
# source-broker can't be b2 since it doesn't have p10
# source-broker shouldn't be b1 since it has lesser partitions
# than b3 and b0
assert broker_source in [b0, b3]
# dest-broker shouldn't be b4 since it has more partitions than b4
# dest-broker can't be b6 since it has already partition p10
assert broker_dest in [b5, b7]
def test_update_sibling_count(self):
t1 = Topic('topic1', 2)
t2 = Topic('topic2', 2)
t3 = Topic('topic3', 1)
p10 = create_and_attach_partition(t1, 0)
p11 = create_and_attach_partition(t1, 1)
p12 = create_and_attach_partition(t1, 2)
p20 = create_and_attach_partition(t2, 0)
p21 = create_and_attach_partition(t2, 1)
p22 = create_and_attach_partition(t2, 2)
p30 = create_and_attach_partition(t3, 0)
p31 = create_and_attach_partition(t3, 1)
b1 = create_broker('b1', [p10, p11, p20, p21, p30, p31])
b2 = create_broker('b2', [p12, p21, p22])
b3 = create_broker('b3', [p10, p11, p22])
rg = ReplicationGroup('rg', set([b1, b2, b3]))
sibling_distance = {
b2: {
b1: {t1: -1, t2: 0, t3: -2},
b3: {t1: 1, t2: -1, t3: -2}
},
}
# Move a p10 from b1 to b2
b1.move_partition(p10, b2)
# NOTE: b2: b1: t1: -1 -> 1 and b2: b3: t1: 1 -> 0
expected = {
b2: {
b1: {t1: 1, t2: 0, t3: -2},
b3: {t1: 0, t2: -1, t3: -2},
},
}
actual = rg.update_sibling_distance(sibling_distance, b2, t1)
assert dict(actual) == expected
def test_get_target_brokers_3(self, create_partition):
p10 = create_partition('topic1', 0)
p11 = create_partition('topic1', 1)
p12 = create_partition('topic1', 2)
p20 = create_partition('topic2', 0)
p21 = create_partition('topic2', 1)
p30 = create_partition('topic3', 0)
p31 = create_partition('topic3', 1)
b1 = create_broker('b1', [p10, p11, p20])
b2 = create_broker('b2', [p12, p21, p30, p31])
b3 = create_broker('b3', [])
rg = ReplicationGroup('rg', set([b1, b2, b3]))
over_loaded = [b1, b2]
under_loaded = [b3]
target = rg._get_target_brokers(
over_loaded,
under_loaded,
rg.generate_sibling_distance(),
)
# b3 is the best broker fit because of number of partition
# Both p10 and p11 are a good fit. p20 can't be considered because it is
# already in b3.
assert target == (b2, b3, p30) or target == (b2, b3, p31)
def test_move_partition(self, create_partition):
p10 = create_partition('t1', 0)
p11 = create_partition('t1', 1)
p20 = create_partition('t2', 0)
b1 = create_broker('b1', [p10, p20, p11])
b2 = create_broker('b2', [p10, p11])
# 2 p1 replicas are in rg_source
rg_source = ReplicationGroup('rg1', set([b1, b2]))
b3 = create_broker('b3', [p10, p11])
b4 = create_broker('b4', [p20])
# 1 p1 replica is in rg_dest
rg_dest = ReplicationGroup('rg2', set([b3, b4]))
# Move partition p1 from rg1 to rg2
rg_source.move_partition(rg_dest, p10)
# partition-count of p1 for rg1 should be 1
assert rg_source.count_replica(p10) == 1
# partition-count of p1 for rg2 should be 2
assert rg_dest.count_replica(p10) == 2
# Rest of the partitions are untouched
assert rg_source.count_replica(p20) == 1
assert rg_dest.count_replica(p20) == 1
assert rg_source.count_replica(p11) == 2
assert rg_dest.count_replica(p11) == 1
def test_generate_sibling_count(self):
t1 = Topic('topic1', 2)
t2 = Topic('topic2', 2)
t3 = Topic('topic3', 1)
p10 = create_and_attach_partition(t1, 0)
p11 = create_and_attach_partition(t1, 1)
p12 = create_and_attach_partition(t1, 2)
p20 = create_and_attach_partition(t2, 0)
p21 = create_and_attach_partition(t2, 1)
p22 = create_and_attach_partition(t2, 2)
p30 = create_and_attach_partition(t3, 0)
p31 = create_and_attach_partition(t3, 1)
b1 = create_broker('b1', [p10, p11, p20, p21, p30, p31])
b2 = create_broker('b2', [p12, p21, p22])
b3 = create_broker('b3', [p10, p11, p22])
rg = ReplicationGroup('rg', set([b1, b2, b3]))
expected = {
b1: {b2: {t1: 1, t2: 0, t3: 2}, b3: {t1: 0, t2: 1, t3: 2}},
b2: {b1: {t1: -1, t2: 0, t3: -2}, b3: {t1: -1, t2: 1, t3: 0}},
b3: {b1: {t1: 0, t2: -1, t3: -2}, b2: {t1: 1, t2: -1, t3: 0}},
}
actual = rg.generate_sibling_distance()
assert dict(actual) == expected
def test_rebalance_brokers_for_topic_partition_imbalance(self, create_partition):
# Broker Topics:Partition
# 1 All partitions: 4 from topic1, 2 from topic2 and 1 each from topic3 and topic4
# 2 Newly added broker:(empty)
# total partitions: 8
p10 = create_partition('topic1', 0)
p11 = create_partition('topic1', 1)
p12 = create_partition('topic1', 2)
p13 = create_partition('topic1', 3)
p20 = create_partition('topic2', 0)
p21 = create_partition('topic2', 1)
p30 = create_partition('topic3', 0)
p40 = create_partition('topic4', 0)
b1 = create_broker('b1', [p10, p11, p12, p13, p20, p21, p30, p40])
b2 = create_broker('b2', [])
rg = ReplicationGroup('test_rg', set([b1, b2]))
rg.rebalance_brokers()
possible_topics1 = set([p10.topic, p11.topic, p20.topic, p30.topic])
possible_topics2 = set([p10.topic, p11.topic, p20.topic, p40.topic])
assert (
b1.topics == possible_topics1 and
b2.topics == possible_topics2
) or (
b1.topics == possible_topics2 and
b2.topics == possible_topics1
)
def test_get_target_brokers_3(self, create_partition):
p10 = create_partition('topic1', 0)
p11 = create_partition('topic1', 1)
p12 = create_partition('topic1', 2)
p20 = create_partition('topic2', 0)
p21 = create_partition('topic2', 1)
p30 = create_partition('topic3', 0)
p31 = create_partition('topic3', 1)
b1 = create_broker('b1', [p10, p11, p20])
b2 = create_broker('b2', [p12, p21, p30, p31])
b3 = create_broker('b3', [])
rg = ReplicationGroup('rg', set([b1, b2, b3]))
over_loaded = [b1, b2]
under_loaded = [b3]
target = rg._get_target_brokers(
over_loaded,
under_loaded,
rg.generate_sibling_distance(),
)
# b3 is the best broker fit because of number of partition
# Both p10 and p11 are a good fit. p20 can't be considered because it is
# already in b3.
assert target == (b2, b3, p30) or target == (b2, b3, p31)
def test__select_broker_pair(self, create_partition):
p10 = create_partition('t1', 0)
p11 = create_partition('t1', 1)
p20 = create_partition('t2', 0)
b0 = create_broker('b0', [p10, p20, p11])
b1 = create_broker('b1', [p10, p20])
b2 = create_broker('b2', [p20, p11])
b3 = create_broker('b3', [p10, p20, p11])
rg_source = ReplicationGroup('rg1', set([b0, b1, b2, b3]))
b4 = create_broker('b4', [p20, p11])
b5 = create_broker('b5', [p20])
b6 = create_broker('b6', [p10])
b7 = create_broker('b7', [p11])
rg_dest = ReplicationGroup('rg2', set([b4, b5, b6, b7]))
# Select best-suitable brokers for moving partition 'p10'
broker_source, broker_dest = rg_source._select_broker_pair(rg_dest, p10)
# source-broker can't be b2 since it doesn't have p10
# source-broker shouldn't be b1 since it has lesser partitions
# than b3 and b0
assert broker_source in [b0, b3]
# dest-broker shouldn't be b4 since it has more partitions than b4
# dest-broker can't be b6 since it has already partition p10
assert broker_dest in [b5, b7]
def test_add_broker(self):
rg = ReplicationGroup(
'test_rg',
set([sentinel.broker1, sentinel.broker2]),
)
rg.add_broker(sentinel.broker)
assert sentinel.broker in rg.brokers
def test_add_broker(self):
rg = ReplicationGroup(
'test_rg',
set([sentinel.broker1, sentinel.broker2]),
)
rg.add_broker(sentinel.broker)
assert sentinel.broker in rg.brokers
def test_generate_sibling_count(self):
t1 = Topic('topic1', 2)
t2 = Topic('topic2', 2)
t3 = Topic('topic3', 1)
p10 = create_and_attach_partition(t1, 0)
p11 = create_and_attach_partition(t1, 1)
p12 = create_and_attach_partition(t1, 2)
p20 = create_and_attach_partition(t2, 0)
p21 = create_and_attach_partition(t2, 1)
p22 = create_and_attach_partition(t2, 2)
p30 = create_and_attach_partition(t3, 0)
p31 = create_and_attach_partition(t3, 1)
b1 = create_broker('b1', [p10, p11, p20, p21, p30, p31])
b2 = create_broker('b2', [p12, p21, p22])
b3 = create_broker('b3', [p10, p11, p22])
rg = ReplicationGroup('rg', set([b1, b2, b3]))
expected = {
b1: {
b2: {
t1: 1,
t2: 0,
t3: 2
},
b3: {
t1: 0,
t2: 1,
t3: 2
}
},
b2: {
b1: {
t1: -1,
t2: 0,
t3: -2
},
b3: {
t1: -1,
t2: 1,
t3: 0
}
},
b3: {
b1: {
t1: 0,
t2: -1,
t3: -2
},
b2: {
t1: 1,
t2: -1,
t3: 0
}
},
}
actual = rg.generate_sibling_distance()
assert dict(actual) == expected
def test_remove_replica_invalid_broker(self, create_partition):
p10 = create_partition('topic1', 0)
p20 = create_partition('topic2', 0)
b1 = create_broker('b1', [p10])
b2 = create_broker('b2', [p20])
b3 = create_broker('b3', [p10])
rg = ReplicationGroup('test_rg', set([b1, b2]))
with pytest.raises(AssertionError):
rg.remove_replica(p10, [b2, b3])
def test_remove_replica_invalid_broker(self, create_partition):
p10 = create_partition('topic1', 0)
p20 = create_partition('topic2', 0)
b1 = create_broker('b1', [p10])
b2 = create_broker('b2', [p20])
b3 = create_broker('b3', [p10])
rg = ReplicationGroup('test_rg', set([b1, b2]))
with pytest.raises(AssertionError):
rg.remove_replica(p10, [b2, b3])
def test__elect_dest_broker_partition_conflict(self, create_partition):
p1 = create_partition('t1', 0)
p2 = create_partition('t2', 0)
p3 = create_partition('t1', 1)
b1 = create_broker('b1', [p1, p2, p3]) # b1 -> t1: {0,1}, t2: 0
rg = ReplicationGroup('test_rg', set([b1]))
# p1 already exists in b1
# This should never happen and we expect the application to fail badly
victim_partition = p1
assert None is rg._elect_dest_broker(victim_partition)
def test__elect_dest_broker_partition_conflict(self, create_partition):
p1 = create_partition('t1', 0)
p2 = create_partition('t2', 0)
p3 = create_partition('t1', 1)
b1 = create_broker('b1', [p1, p2, p3]) # b1 -> t1: {0,1}, t2: 0
rg = ReplicationGroup('test_rg', set([b1]))
# p1 already exists in b1
# This should never happen and we expect the application to fail badly
victim_partition = p1
assert None is rg._elect_dest_broker(victim_partition)
def test_acquire_partition_error(self, create_partition):
p10 = create_partition('t1', 0)
p11 = create_partition('t1', 1)
p20 = create_partition('t2', 0)
b1 = create_broker('b1', [p10, p11, p20])
rg = ReplicationGroup('test_rg', set([b1]))
b3 = create_broker('b3', [p11])
# Not eligible because broker b3 has already a replica of partition p11
with pytest.raises(NotEligibleGroupError):
rg.acquire_partition(p11, b3)
def test_acquire_partition_error(self, create_partition):
p10 = create_partition('t1', 0)
p11 = create_partition('t1', 1)
p20 = create_partition('t2', 0)
b1 = create_broker('b1', [p10, p11, p20])
rg = ReplicationGroup('test_rg', set([b1]))
b3 = create_broker('b3', [p11])
# Not eligible because broker b3 has already a replica of partition p11
with pytest.raises(NotEligibleGroupError):
rg.acquire_partition(p11, b3)
def test_add_replica(self, create_partition):
p10 = create_partition('topic1', 0)
p20 = create_partition('topic2', 0)
b1 = create_broker('b1', [p10, p20])
b2 = create_broker('b2', [p10, p20])
b3 = create_broker('b3', [p20])
rg = ReplicationGroup('test_rg', set([b1, b2, b3]))
rg.add_replica(p10)
assert b1.partitions == set([p10, p20])
assert b2.partitions == set([p10, p20])
assert b3.partitions == set([p10, p20])
def test_decommission_not_enough_replicas(self, create_partition):
p10 = create_partition('topic1', 0)
p11 = create_partition('topic1', 1)
p20 = create_partition('topic2', 0, replication_factor=2)
b1 = create_broker('b1', [p10, p11, p20])
b2 = create_broker('b1', [p20])
rg = ReplicationGroup('rg', set([b1, b2]))
b2.mark_decommissioned()
rg.rebalance_brokers()
assert not b2.empty()
def test_decommission_not_enough_replicas(self, create_partition):
p10 = create_partition('topic1', 0)
p11 = create_partition('topic1', 1)
p20 = create_partition('topic2', 0, replication_factor=2)
b1 = create_broker('b1', [p10, p11, p20])
b2 = create_broker('b1', [p20])
rg = ReplicationGroup('rg', set([b1, b2]))
b2.mark_decommissioned()
rg.rebalance_brokers()
assert not b2.empty()
def test_add_replica(self, create_partition):
p10 = create_partition('topic1', 0)
p20 = create_partition('topic2', 0)
b1 = create_broker('b1', [p10, p20])
b2 = create_broker('b2', [p10, p20])
b3 = create_broker('b3', [p20])
rg = ReplicationGroup('test_rg', set([b1, b2, b3]))
rg.add_replica(p10)
assert b1.partitions == set([p10, p20])
assert b2.partitions == set([p10, p20])
assert b3.partitions == set([p10, p20])
def test_decommission_no_remaining_brokers(self, create_partition):
p10 = create_partition('topic1', 0)
p11 = create_partition('topic1', 1)
p20 = create_partition('topic2', 0)
b1 = create_broker('b1', [p10, p11, p20])
b2 = create_broker('b2', [])
b2.mark_inactive()
rg = ReplicationGroup('rg', set([b1, b2]))
b1.mark_decommissioned()
# Two brokers b1 decommissioned b2 inactive
with pytest.raises(EmptyReplicationGroupError):
rg.rebalance_brokers()
def test_decommission_no_remaining_brokers(self, create_partition):
p10 = create_partition('topic1', 0)
p11 = create_partition('topic1', 1)
p20 = create_partition('topic2', 0)
b1 = create_broker('b1', [p10, p11, p20])
b2 = create_broker('b2', [])
b2.mark_inactive()
rg = ReplicationGroup('rg', set([b1, b2]))
b1.mark_decommissioned()
# Two brokers b1 decommissioned b2 inactive
with pytest.raises(EmptyReplicationGroupError):
rg.rebalance_brokers()
def test_update_sibling_count(self):
t1 = Topic('topic1', 2)
t2 = Topic('topic2', 2)
t3 = Topic('topic3', 1)
p10 = create_and_attach_partition(t1, 0)
p11 = create_and_attach_partition(t1, 1)
p12 = create_and_attach_partition(t1, 2)
p20 = create_and_attach_partition(t2, 0)
p21 = create_and_attach_partition(t2, 1)
p22 = create_and_attach_partition(t2, 2)
p30 = create_and_attach_partition(t3, 0)
p31 = create_and_attach_partition(t3, 1)
b1 = create_broker('b1', [p10, p11, p20, p21, p30, p31])
b2 = create_broker('b2', [p12, p21, p22])
b3 = create_broker('b3', [p10, p11, p22])
rg = ReplicationGroup('rg', set([b1, b2, b3]))
sibling_distance = {
b2: {
b1: {
t1: -1,
t2: 0,
t3: -2
},
b3: {
t1: 1,
t2: -1,
t3: -2
}
},
}
# Move a p10 from b1 to b2
b1.move_partition(p10, b2)
# NOTE: b2: b1: t1: -1 -> 1 and b2: b3: t1: 1 -> 0
expected = {
b2: {
b1: {
t1: 1,
t2: 0,
t3: -2
},
b3: {
t1: 0,
t2: -1,
t3: -2
},
},
}
actual = rg.update_sibling_distance(sibling_distance, b2, t1)
assert dict(actual) == expected
def test_count_replica(self, create_partition):
p10 = create_partition('t1', 0)
p11 = create_partition('t1', 1)
p12 = create_partition('t1', 2)
p13 = create_partition('t1', 3)
b1 = create_broker('b1', [p10, p11, p12])
b2 = create_broker('b2', [p10, p13])
rg = ReplicationGroup('test_rg', set([b1, b2]))
assert rg.count_replica(p10) == 2
assert rg.count_replica(p11) == 1
assert rg.count_replica(p12) == 1
assert rg.count_replica(p13) == 1
assert rg.count_replica(create_partition('t1', 4)) == 0
def test__elect_dest_broker(self, create_partition):
p10 = create_partition('t1', 0)
p11 = create_partition('t1', 1)
p20 = create_partition('t2', 0)
p30 = create_partition('t3', 0)
b1 = create_broker('b1', [p10, p20, p11]) # b1 -> t1: {0,1}, t2: 0
b2 = create_broker('b2', [p10, p30]) # b2 -> t1: 0, t3: 0
rg = ReplicationGroup('test_rg', set([b1, b2]))
# Since p30 is already in b2 so the preferred destination will be b1
# although b2 has less partitions.
victim_partition = p30
actual = rg._elect_dest_broker(victim_partition)
assert actual == b1
def test__elect_dest_broker(self, create_partition):
p10 = create_partition('t1', 0)
p11 = create_partition('t1', 1)
p20 = create_partition('t2', 0)
p30 = create_partition('t3', 0)
b1 = create_broker('b1', [p10, p20, p11]) # b1 -> t1: {0,1}, t2: 0
b2 = create_broker('b2', [p10, p30]) # b2 -> t1: 0, t3: 0
rg = ReplicationGroup('test_rg', set([b1, b2]))
# Since p30 is already in b2 so the preferred destination will be b1
# although b2 has less partitions.
victim_partition = p30
actual = rg._elect_dest_broker(victim_partition)
assert actual == b1
def test_count_replica(self, create_partition):
p10 = create_partition('t1', 0)
p11 = create_partition('t1', 1)
p12 = create_partition('t1', 2)
p13 = create_partition('t1', 3)
b1 = create_broker('b1', [p10, p11, p12])
b2 = create_broker('b2', [p10, p13])
rg = ReplicationGroup('test_rg', set([b1, b2]))
assert rg.count_replica(p10) == 2
assert rg.count_replica(p11) == 1
assert rg.count_replica(p12) == 1
assert rg.count_replica(p13) == 1
assert rg.count_replica(create_partition('t1', 4)) == 0
def test_acquire_partition(self, create_partition):
p10 = create_partition('t1', 0)
p11 = create_partition('t1', 1)
p12 = create_partition('t1', 2)
p13 = create_partition('t1', 3)
p20 = create_partition('t2', 0)
p21 = create_partition('t2', 1)
b1 = create_broker('b1', [p10, p11, p20])
b2 = create_broker('b2', [p12, p21])
rg = ReplicationGroup('test_rg', set([b1, b2]))
b3 = create_broker('b3', [p13])
rg.acquire_partition(p13, b3)
assert b3.empty()
assert p13 in b2.partitions
def test__elect_source_broker(self, create_partition):
p10 = create_partition('t1', 0)
p20 = create_partition('t2', 0)
p11 = create_partition('t1', 1)
b1 = create_broker('b1', [p10, p11]) # b1 -> t1: {0,1}, t2: 0
p30 = create_partition('t3', 0)
b2 = create_broker('b2', [p10, p20, p30]) # b2 -> t1: 0, t2: 1, t3: 0
rg = ReplicationGroup('test_rg', set([b1, b2]))
# b1 has 2 partitions (p1 and p3) for same topic t1
# b2 has only 1 partition (p1) for topic t1
# source broker should be b1 to reduce the number of partitions of the
# same topic
victim_partition = p10
actual = rg._elect_source_broker(victim_partition)
assert actual == b1
def test__elect_source_broker(self, create_partition):
p10 = create_partition('t1', 0)
p20 = create_partition('t2', 0)
p11 = create_partition('t1', 1)
b1 = create_broker('b1', [p10, p11]) # b1 -> t1: {0,1}, t2: 0
p30 = create_partition('t3', 0)
b2 = create_broker('b2', [p10, p20, p30]) # b2 -> t1: 0, t2: 1, t3: 0
rg = ReplicationGroup('test_rg', set([b1, b2]))
# b1 has 2 partitions (p1 and p3) for same topic t1
# b2 has only 1 partition (p1) for topic t1
# source broker should be b1 to reduce the number of partitions of the
# same topic
victim_partition = p10
actual = rg._elect_source_broker(victim_partition)
assert actual == b1
def test_acquire_partition(self, create_partition):
p10 = create_partition('t1', 0)
p11 = create_partition('t1', 1)
p12 = create_partition('t1', 2)
p13 = create_partition('t1', 3)
p20 = create_partition('t2', 0)
p21 = create_partition('t2', 1)
b1 = create_broker('b1', [p10, p11, p20])
b2 = create_broker('b2', [p12, p21])
rg = ReplicationGroup('test_rg', set([b1, b2]))
b3 = create_broker('b3', [p13])
rg.acquire_partition(p13, b3)
assert b3.empty()
assert p13 in b2.partitions
def test__elect_dest_broker_prefer_less_siblings(self, create_partition):
p10 = create_partition('t1', 0)
p11 = create_partition('t1', 1)
p20 = create_partition('t2', 0)
b1 = create_broker('b1', [p10, p11, p20]) # b1 -> t1: {0,1}, t2: 0
p30 = create_partition('t3', 0)
p31 = create_partition('t3', 1)
p32 = create_partition('t3', 2)
b2 = create_broker('b2', [p10, p30, p31, p32]) # b2 -> t1: 0, t3: 0
rg = ReplicationGroup('test_rg', set([b1, b2]))
# Since t1 has two partitions in b1 but only one in b2,
# the preferred destination should be b2
victim_partition_p12 = create_partition('t1', 2)
actual = rg._elect_dest_broker(victim_partition_p12)
assert actual == b2
def test__elect_dest_broker_prefer_less_siblings(self, create_partition):
p10 = create_partition('t1', 0)
p11 = create_partition('t1', 1)
p20 = create_partition('t2', 0)
b1 = create_broker('b1', [p10, p11, p20]) # b1 -> t1: {0,1}, t2: 0
p30 = create_partition('t3', 0)
p31 = create_partition('t3', 1)
p32 = create_partition('t3', 2)
b2 = create_broker('b2', [p10, p30, p31, p32]) # b2 -> t1: 0, t3: 0
rg = ReplicationGroup('test_rg', set([b1, b2]))
# Since t1 has two partitions in b1 but only one in b2,
# the preferred destination should be b2
victim_partition_p12 = create_partition('t1', 2)
actual = rg._elect_dest_broker(victim_partition_p12)
assert actual == b2
def test__get_target_brokers_1(self, create_partition):
p10 = create_partition('topic1', 0, replication_factor=2)
p11 = create_partition('topic1', 1, replication_factor=2)
p20 = create_partition('topic2', 0, replication_factor=2)
p30 = create_partition('topic3', 0)
b1 = create_broker('b1', [p10, p11, p20, p30])
b2 = create_broker('b2', [p10, p20])
rg = ReplicationGroup('rg', set([b1, b2]))
over_loaded = [b1]
under_loaded = [b2]
target = rg._get_target_brokers(
over_loaded,
under_loaded,
rg.generate_sibling_distance(),
)
# p30 is the best fit because topic3 doesn't have any partition in
# broker 2
assert target == (b1, b2, p30) or target == (b1, b2, p11)
def test__get_target_brokers_1(self, create_partition):
p10 = create_partition('topic1', 0, replication_factor=2)
p11 = create_partition('topic1', 1, replication_factor=2)
p20 = create_partition('topic2', 0, replication_factor=2)
p30 = create_partition('topic3', 0)
b1 = create_broker('b1', [p10, p11, p20, p30])
b2 = create_broker('b2', [p10, p20])
rg = ReplicationGroup('rg', set([b1, b2]))
over_loaded = [b1]
under_loaded = [b2]
target = rg._get_target_brokers(
over_loaded,
under_loaded,
rg.generate_sibling_distance(),
)
# p30 is the best fit because topic3 doesn't have any partition in
# broker 2
assert target == (b1, b2, p30) or target == (b1, b2, p11)
def test_decommission_multi_replicas_edge_case(self, create_partition):
# After decommissioning there will be 2 brokers with 5 partitions and 1
# with 4. Because of replication constraints (2 replicas of the same
# partition can't live in the same broker) the algorithm has to be smart
# enough to find a solution.
p10 = create_partition('topic1', 0)
p11 = create_partition('topic1', 1)
p20 = create_partition('topic2', 0, replication_factor=3)
p21 = create_partition('topic2', 1, replication_factor=3)
p22 = create_partition('topic2', 2, replication_factor=3)
p23 = create_partition('topic2', 3, replication_factor=3)
b1 = create_broker('b1', [p10, p20, p23])
b2 = create_broker('b2', [p11, p21, p22, p23])
b3 = create_broker('b3', [p20, p21, p22, p23])
b4 = create_broker('b4', [p20, p21, p22])
rg = ReplicationGroup('rg', set([b1, b2, b3, b4]))
b1.mark_decommissioned()
rg.rebalance_brokers()
assert b1.empty()
def test_decommission_multi_replicas_edge_case(self, create_partition):
# After decommissioning there will be 2 brokers with 5 partitions and 1
# with 4. Because of replication constraints (2 replicas of the same
# partition can't live in the same broker) the algorithm has to be smart
# enough to find a solution.
p10 = create_partition('topic1', 0)
p11 = create_partition('topic1', 1)
p20 = create_partition('topic2', 0, replication_factor=3)
p21 = create_partition('topic2', 1, replication_factor=3)
p22 = create_partition('topic2', 2, replication_factor=3)
p23 = create_partition('topic2', 3, replication_factor=3)
b1 = create_broker('b1', [p10, p20, p23])
b2 = create_broker('b2', [p11, p21, p22, p23])
b3 = create_broker('b3', [p20, p21, p22, p23])
b4 = create_broker('b4', [p20, p21, p22])
rg = ReplicationGroup('rg', set([b1, b2, b3, b4]))
b1.mark_decommissioned()
rg.rebalance_brokers()
assert b1.empty()
def test_partitions(self):
mock_brokers = set([
Mock(spec=Broker, partitions=set([sentinel.p1, sentinel.p2])),
Mock(spec=Broker, partitions=set([sentinel.p3, sentinel.p1])),
])
rg = ReplicationGroup('test_rg', mock_brokers)
expected = [
sentinel.p1,
sentinel.p2,
sentinel.p3,
sentinel.p1,
]
assert sorted(expected) == sorted(rg.partitions)
def test_rebalance_brokers_for_topic_partition_imbalance(
self, create_partition):
# Broker Topics:Partition
# 1 All partitions: 4 from topic1, 2 from topic2 and 1 each from topic3 and topic4
# 2 Newly added broker:(empty)
# total partitions: 8
p10 = create_partition('topic1', 0)
p11 = create_partition('topic1', 1)
p12 = create_partition('topic1', 2)
p13 = create_partition('topic1', 3)
p20 = create_partition('topic2', 0)
p21 = create_partition('topic2', 1)
p30 = create_partition('topic3', 0)
p40 = create_partition('topic4', 0)
b1 = create_broker('b1', [p10, p11, p12, p13, p20, p21, p30, p40])
b2 = create_broker('b2', [])
rg = ReplicationGroup('test_rg', set([b1, b2]))
rg.rebalance_brokers()
possible_topics1 = set([p10.topic, p11.topic, p20.topic, p30.topic])
possible_topics2 = set([p10.topic, p11.topic, p20.topic, p40.topic])
assert (b1.topics == possible_topics1 and b2.topics
== possible_topics2) or (b1.topics == possible_topics2
and b2.topics == possible_topics1)
def rg_balanced(create_partition):
# Broker Topics:Partition
# 1 topic1:0, topic1:1, topic2:0, topic4:0
# 2 topic1:2, topic1:3, topic2:1
# 3 topic2:1, topic3:0, topic4:0
# total partitions: 9
p10 = create_partition('topic1', 0)
p11 = create_partition('topic1', 1)
p12 = create_partition('topic1', 2)
p13 = create_partition('topic1', 3)
p20 = create_partition('topic2', 0)
p21 = create_partition('topic2', 1)
p30 = create_partition('topic3', 0)
p40 = create_partition('topic4', 0, replication_factor=2)
b1 = create_broker('b1', [p10, p11, p20, p40])
b2 = create_broker('b2', [p12, p13, p21])
b3 = create_broker('b3', [p21, p30, p40])
return ReplicationGroup('test_rg', set([b1, b2, b3]))
def test_add_broker_empty(self):
rg = ReplicationGroup('test_rg', None)
rg.add_broker(sentinel.broker)
assert set([sentinel.broker]) == rg.brokers
def test_rebalance_empty_replication_group(self):
rg = ReplicationGroup('empty_rg')
with pytest.raises(EmptyReplicationGroupError):
rg.rebalance_brokers()
def test_invalid_brokers_type_list(self):
with pytest.raises(TypeError):
ReplicationGroup('test_rg', [sentinel.broker1, sentinel.broker2])
def test_id(self):
rg = ReplicationGroup('test_rg', None)
assert 'test_rg' == rg.id
def test_move_partition(self, create_partition):
p10 = create_partition('t1', 0)
p11 = create_partition('t1', 1)
p20 = create_partition('t2', 0)
b1 = create_broker('b1', [p10, p20, p11])
b2 = create_broker('b2', [p10, p11])
# 2 p1 replicas are in rg_source
rg_source = ReplicationGroup('rg1', set([b1, b2]))
b3 = create_broker('b3', [p10, p11])
b4 = create_broker('b4', [p20])
# 1 p1 replica is in rg_dest
rg_dest = ReplicationGroup('rg2', set([b3, b4]))
# Move partition p1 from rg1 to rg2
rg_source.move_partition(rg_dest, p10)
# partition-count of p1 for rg1 should be 1
assert rg_source.count_replica(p10) == 1
# partition-count of p1 for rg2 should be 2
assert rg_dest.count_replica(p10) == 2
# Rest of the partitions are untouched
assert rg_source.count_replica(p20) == 1
assert rg_dest.count_replica(p20) == 1
assert rg_source.count_replica(p11) == 2
assert rg_dest.count_replica(p11) == 1
def test_rebalance_empty_replication_group(self):
rg = ReplicationGroup('empty_rg')
with pytest.raises(EmptyReplicationGroupError):
rg.rebalance_brokers()
def test_add_broker_empty(self):
rg = ReplicationGroup('test_rg', None)
rg.add_broker(sentinel.broker)
assert set([sentinel.broker]) == rg.brokers