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__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_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_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__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_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_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_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_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__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_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__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_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 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__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_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_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
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