def test_no_exceptions_when_only_subscribed_topic_is_deleted(mocker): cluster = create_cluster(mocker, topics={'t'}, topics_partitions={0, 1, 2}) subscriptions = { 'C': {'t'}, } member_metadata = make_member_metadata(subscriptions) sticky_assignment = StickyPartitionAssignor.assign(cluster, member_metadata) expected_assignment = { 'C': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t', [0, 1, 2])], b''), } assert_assignment(sticky_assignment, expected_assignment) subscriptions = { 'C': {}, } member_metadata = {} for member, topics in six.iteritems(subscriptions): member_metadata[member] = StickyPartitionAssignor._metadata( topics, sticky_assignment[member].partitions()) cluster = create_cluster(mocker, topics={}, topics_partitions={}) sticky_assignment = StickyPartitionAssignor.assign(cluster, member_metadata) expected_assignment = { 'C': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [], b''), } assert_assignment(sticky_assignment, expected_assignment)
def test_sticky_multiple_consumers_mixed_topic_subscriptions(mocker): partitions = {'t1': {0, 1, 2}, 't2': {0, 1}} cluster = create_cluster(mocker, topics={'t1', 't2'}, topic_partitions_lambda=lambda t: partitions[t]) subscriptions = { 'C1': {'t1'}, 'C2': {'t1', 't2'}, 'C3': {'t1'}, } member_metadata = make_member_metadata(subscriptions) sticky_assignment = StickyPartitionAssignor.assign(cluster, member_metadata) expected_assignment = { 'C1': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t1', [0, 2])], b''), 'C2': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t2', [0, 1])], b''), 'C3': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t1', [1])], b''), } assert_assignment(sticky_assignment, expected_assignment)
def test_join_complete(mocker, coordinator): coordinator._subscription.subscribe(topics=['foobar']) assignor = RoundRobinPartitionAssignor() coordinator.config['assignors'] = (assignor,) mocker.spy(assignor, 'on_assignment') assert assignor.on_assignment.call_count == 0 assignment = ConsumerProtocolMemberAssignment(0, [('foobar', [0, 1])], b'') coordinator._on_join_complete( 0, 'member-foo', 'roundrobin', assignment.encode()) assert assignor.on_assignment.call_count == 1 assignor.on_assignment.assert_called_with(assignment)
def test_subscription_listener_failure(mocker, coordinator): listener = mocker.MagicMock(spec=ConsumerRebalanceListener) coordinator._subscription.subscribe(topics=['foobar'], listener=listener) # exception raised in listener should not be re-raised by coordinator listener.on_partitions_revoked.side_effect = Exception('crash') coordinator._on_join_prepare(0, 'member-foo') assert listener.on_partitions_revoked.call_count == 1 assignment = ConsumerProtocolMemberAssignment(0, [('foobar', [0, 1])], b'') coordinator._on_join_complete(0, 'member-foo', 'roundrobin', assignment.encode()) assert listener.on_partitions_assigned.call_count == 1
def test_subscription_listener(mocker, coordinator): listener = mocker.MagicMock(spec=ConsumerRebalanceListener) coordinator._subscription.subscribe(topics=['foobar'], listener=listener) coordinator._on_join_prepare(0, 'member-foo') assert listener.on_partitions_revoked.call_count == 1 listener.on_partitions_revoked.assert_called_with(set([])) assignment = ConsumerProtocolMemberAssignment(0, [('foobar', [0, 1])], b'') coordinator._on_join_complete(0, 'member-foo', 'roundrobin', assignment.encode()) assert listener.on_partitions_assigned.call_count == 1 listener.on_partitions_assigned.assert_called_with( set([TopicPartition('foobar', 0), TopicPartition('foobar', 1)]))
def assign(cls, cluster, members): """Performs group assignment given cluster metadata and member subscriptions Arguments: cluster (ClusterMetadata): cluster metadata members (dict of {member_id: MemberMetadata}): decoded metadata for each member in the group. Returns: dict: {member_id: MemberAssignment} """ members_metadata = {} for consumer, member_metadata in six.iteritems(members): members_metadata[consumer] = cls.parse_member_metadata( member_metadata) executor = StickyAssignmentExecutor(cluster, members_metadata) executor.perform_initial_assignment() executor.balance() cls._latest_partition_movements = executor.partition_movements assignment = {} for member_id in members: assignment[member_id] = ConsumerProtocolMemberAssignment( cls.version, sorted(executor.get_final_assignment(member_id)), b"") return assignment
def assign(cls, cluster, member_metadata): consumers_per_topic = collections.defaultdict(list) for member, metadata in six.iteritems(member_metadata): for topic in metadata.subscription: consumers_per_topic[topic].append(member) # construct {member_id: {topic: [partition, ...]}} assignment = collections.defaultdict(dict) for topic, consumers_for_topic in six.iteritems(consumers_per_topic): partitions = cluster.partitions_for_topic(topic) if partitions is None: log.warning('No partition metadata for topic %s', topic) continue partitions = sorted(partitions) consumers_for_topic.sort() partitions_per_consumer = len(partitions) // len(consumers_for_topic) consumers_with_extra = len(partitions) % len(consumers_for_topic) for i, member in enumerate(consumers_for_topic): start = partitions_per_consumer * i start += min(i, consumers_with_extra) length = partitions_per_consumer if not i + 1 > consumers_with_extra: length += 1 assignment[member][topic] = partitions[start:start+length] protocol_assignment = {} for member_id in member_metadata: protocol_assignment[member_id] = ConsumerProtocolMemberAssignment( cls.version, sorted(assignment[member_id].items()), b'') return protocol_assignment
def test_assignment_updated_for_deleted_topic(mocker): def topic_partitions(topic): if topic == 't1': return {0} if topic == 't3': return set(range(100)) cluster = create_cluster(mocker, topics={'t1', 't3'}, topic_partitions_lambda=topic_partitions) subscriptions = { 'C': {'t1', 't2', 't3'}, } member_metadata = make_member_metadata(subscriptions) sticky_assignment = StickyPartitionAssignor.assign(cluster, member_metadata) expected_assignment = { 'C': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t1', [0]), ('t3', list(range(100)))], b''), } assert_assignment(sticky_assignment, expected_assignment)
def test_sticky_two_consumers_one_topic_two_partitions(mocker): cluster = create_cluster(mocker, topics={'t'}, topics_partitions={0, 1}) subscriptions = { 'C1': {'t'}, 'C2': {'t'}, } member_metadata = make_member_metadata(subscriptions) sticky_assignment = StickyPartitionAssignor.assign(cluster, member_metadata) expected_assignment = { 'C1': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t', [0])], b''), 'C2': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t', [1])], b''), } assert_assignment(sticky_assignment, expected_assignment)
def test_assignor_range(cluster): assignor = RangePartitionAssignor member_metadata = { 'C0': assignor.metadata(set(['t0', 't1'])), 'C1': assignor.metadata(set(['t0', 't1'])), } ret = assignor.assign(cluster, member_metadata) expected = { 'C0': ConsumerProtocolMemberAssignment(assignor.version, [('t0', [0, 1]), ('t1', [0, 1])], b''), 'C1': ConsumerProtocolMemberAssignment(assignor.version, [('t0', [2]), ('t1', [2])], b'') } assert ret == expected assert set(ret) == set(expected) for member in ret: assert ret[member].encode() == expected[member].encode()
def test_assignor_range(mocker): assignor = RangePartitionAssignor member_metadata = { 'C0': assignor.metadata({'t0', 't1'}), 'C1': assignor.metadata({'t0', 't1'}), } cluster = create_cluster(mocker, {'t0', 't1'}, topics_partitions={0, 1, 2}) ret = assignor.assign(cluster, member_metadata) expected = { 'C0': ConsumerProtocolMemberAssignment(assignor.version, [('t0', [0, 1]), ('t1', [0, 1])], b''), 'C1': ConsumerProtocolMemberAssignment(assignor.version, [('t0', [2]), ('t1', [2])], b'') } assert ret == expected assert set(ret) == set(expected) for member in ret: assert ret[member].encode() == expected[member].encode()
def test_perform_assignment(mocker, coordinator): member_metadata = { 'member-foo': ConsumerProtocolMemberMetadata(0, ['foo1'], b''), 'member-bar': ConsumerProtocolMemberMetadata(0, ['foo1'], b'') } assignments = { 'member-foo': ConsumerProtocolMemberAssignment(0, [('foo1', [0])], b''), 'member-bar': ConsumerProtocolMemberAssignment(0, [('foo1', [1])], b'') } mocker.patch.object(RoundRobinPartitionAssignor, 'assign') RoundRobinPartitionAssignor.assign.return_value = assignments ret = coordinator._perform_assignment( 'member-foo', 'roundrobin', [(member, metadata.encode()) for member, metadata in member_metadata.items()]) assert RoundRobinPartitionAssignor.assign.call_count == 1 RoundRobinPartitionAssignor.assign.assert_called_with( coordinator._client.cluster, member_metadata) assert ret == assignments
def test_sticky_one_consumer_nonexisting_topic(mocker): cluster = create_cluster(mocker, topics={}, topics_partitions={}) subscriptions = { 'C': {'t'}, } member_metadata = make_member_metadata(subscriptions) sticky_assignment = StickyPartitionAssignor.assign(cluster, member_metadata) expected_assignment = { 'C': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [], b''), } assert_assignment(sticky_assignment, expected_assignment)
def test_sticky_should_only_assign_partitions_from_subscribed_topics(mocker): cluster = create_cluster(mocker, topics={'t', 'other-t'}, topics_partitions={0, 1, 2}) subscriptions = { 'C': {'t'}, } member_metadata = make_member_metadata(subscriptions) sticky_assignment = StickyPartitionAssignor.assign(cluster, member_metadata) expected_assignment = { 'C': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t', [0, 1, 2])], b''), } assert_assignment(sticky_assignment, expected_assignment)
def test_sticky_one_consumer_multiple_topics(mocker): cluster = create_cluster(mocker, topics={'t1', 't2'}, topics_partitions={0, 1, 2}) subscriptions = { 'C': {'t1', 't2'}, } member_metadata = make_member_metadata(subscriptions) sticky_assignment = StickyPartitionAssignor.assign(cluster, member_metadata) expected_assignment = { 'C': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t1', [0, 1, 2]), ('t2', [0, 1, 2])], b''), } assert_assignment(sticky_assignment, expected_assignment)
def _protocol_assignments( self, assignments: ClientAssignmentMapping, cl_distribution: HostToPartitionMap, topic_groups: Mapping[str, int]) -> MemberAssignmentMapping: return { client: ConsumerProtocolMemberAssignment( self.version, sorted( assignment.kafka_protocol_assignment(self._table_manager)), self._compress( ClientMetadata( assignment=assignment, url=self._member_urls[client], changelog_distribution=cl_distribution, topic_groups=topic_groups, ).dumps(), ), ) for client, assignment in assignments.items() }
def assign(cls, cluster, member_metadata): # get all topics and check every memeber has the same all_topics = None for metadata in six.itervalues(member_metadata): if all_topics is None: all_topics = set(metadata.subscription) elif all_topics != set(metadata.subscription): diff = all_topics.symmetric_difference(metadata.subscription) raise UnmergeableTopcis( 'Topic(s) %s do not appear in all members', ', '.join(diff)) # get all partition numbers and check every topic has the same all_partitions = None for topic in all_topics: partitions = cluster.partitions_for_topic(topic) if partitions is None: raise UnmergeableTopcis('No partition metadata for topic %s', topic) if all_partitions is None: all_partitions = set(partitions) elif all_partitions != set(partitions): diff = all_partitions.symmetric_difference(partitions) raise UnmergeableTopcis( 'Partition(s) %s do not appear in all topics', ', '.join(str(p) for p in diff)) all_partitions = sorted(all_partitions) assignment = collections.defaultdict( lambda: collections.defaultdict(list)) # round robin assignation of the partition numbers member_iter = itertools.cycle(sorted(member_metadata.keys())) for partition in all_partitions: member_id = next(member_iter) for topic in all_topics: assignment[member_id][topic].append(partition) protocol_assignment = {} for member_id in member_metadata: protocol_assignment[member_id] = ConsumerProtocolMemberAssignment( cls.version, sorted(assignment[member_id].items()), b'') return protocol_assignment
def test_sticky_add_remove_consumer_one_topic(mocker): cluster = create_cluster(mocker, topics={'t'}, topics_partitions={0, 1, 2}) subscriptions = { 'C1': {'t'}, } member_metadata = make_member_metadata(subscriptions) assignment = StickyPartitionAssignor.assign(cluster, member_metadata) expected_assignment = { 'C1': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t', [0, 1, 2])], b''), } assert_assignment(assignment, expected_assignment) subscriptions = { 'C1': {'t'}, 'C2': {'t'}, } member_metadata = {} for member, topics in six.iteritems(subscriptions): member_metadata[member] = StickyPartitionAssignor._metadata( topics, assignment[member].partitions() if member in assignment else []) assignment = StickyPartitionAssignor.assign(cluster, member_metadata) verify_validity_and_balance(subscriptions, assignment) subscriptions = { 'C2': {'t'}, } member_metadata = {} for member, topics in six.iteritems(subscriptions): member_metadata[member] = StickyPartitionAssignor._metadata( topics, assignment[member].partitions()) assignment = StickyPartitionAssignor.assign(cluster, member_metadata) verify_validity_and_balance(subscriptions, assignment) assert len(assignment['C2'].assignment[0][1]) == 3
def assign(cls, cluster, member_metadata): all_topics = set() for metadata in six.itervalues(member_metadata): all_topics.update(metadata.subscription) all_topic_partitions = [] for topic in all_topics: partitions = cluster.partitions_for_topic(topic) if partitions is None: log.warning('No partition metadata for topic %s', topic) continue for partition in partitions: all_topic_partitions.append(TopicPartition(topic, partition)) all_topic_partitions.sort() # construct {member_id: {topic: [partition, ...]}} assignment = collections.defaultdict( lambda: collections.defaultdict(list)) member_iter = itertools.cycle(sorted(member_metadata.keys())) for partition in all_topic_partitions: member_id = next(member_iter) # Because we constructed all_topic_partitions from the set of # member subscribed topics, we should be safe assuming that # each topic in all_topic_partitions is in at least one member # subscription; otherwise this could yield an infinite loop while partition.topic not in member_metadata[ member_id].subscription: member_id = next(member_iter) assignment[member_id][partition.topic].append(partition.partition) protocol_assignment = {} for member_id in member_metadata: protocol_assignment[member_id] = ConsumerProtocolMemberAssignment( cls.version, sorted(assignment[member_id].items()), b'') return protocol_assignment
def test_sticky_assignor2(mocker): """ Given: there are three consumers C0, C1, C2, and three topics t0, t1, t2, with 1, 2, and 3 partitions respectively. Therefore, the partitions are t0p0, t1p0, t1p1, t2p0, t2p1, t2p2. C0 is subscribed to t0; C1 is subscribed to t0, t1; and C2 is subscribed to t0, t1, t2. Then: perform the assignment Expected: the assignment is - C0 [t0p0] - C1 [t1p0, t1p1] - C2 [t2p0, t2p1, t2p2] Then: remove C0 and perform the assignment Expected: the assignment is - C1 [t0p0, t1p0, t1p1] - C2 [t2p0, t2p1, t2p2] """ partitions = {'t0': {0}, 't1': {0, 1}, 't2': {0, 1, 2}} cluster = create_cluster(mocker, topics={'t0', 't1', 't2'}, topic_partitions_lambda=lambda t: partitions[t]) subscriptions = { 'C0': {'t0'}, 'C1': {'t0', 't1'}, 'C2': {'t0', 't1', 't2'}, } member_metadata = {} for member, topics in six.iteritems(subscriptions): member_metadata[member] = StickyPartitionAssignor._metadata(topics, []) sticky_assignment = StickyPartitionAssignor.assign(cluster, member_metadata) expected_assignment = { 'C0': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t0', [0])], b''), 'C1': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t1', [0, 1])], b''), 'C2': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t2', [0, 1, 2])], b''), } assert_assignment(sticky_assignment, expected_assignment) del subscriptions['C0'] member_metadata = {} for member, topics in six.iteritems(subscriptions): member_metadata[member] = StickyPartitionAssignor._metadata( topics, sticky_assignment[member].partitions()) sticky_assignment = StickyPartitionAssignor.assign(cluster, member_metadata) expected_assignment = { 'C1': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t0', [0]), ('t1', [0, 1])], b''), 'C2': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t2', [0, 1, 2])], b''), } assert_assignment(sticky_assignment, expected_assignment)
def test_sticky_add_remove_topic_two_consumers(mocker): cluster = create_cluster(mocker, topics={'t1', 't2'}, topics_partitions={0, 1, 2}) subscriptions = { 'C1': {'t1'}, 'C2': {'t1'}, } member_metadata = make_member_metadata(subscriptions) sticky_assignment = StickyPartitionAssignor.assign(cluster, member_metadata) expected_assignment = { 'C1': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t1', [0, 2])], b''), 'C2': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t1', [1])], b''), } assert_assignment(sticky_assignment, expected_assignment) subscriptions = { 'C1': {'t1', 't2'}, 'C2': {'t1', 't2'}, } member_metadata = {} for member, topics in six.iteritems(subscriptions): member_metadata[member] = StickyPartitionAssignor._metadata( topics, sticky_assignment[member].partitions()) sticky_assignment = StickyPartitionAssignor.assign(cluster, member_metadata) expected_assignment = { 'C1': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t1', [0, 2]), ('t2', [1])], b''), 'C2': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t1', [1]), ('t2', [0, 2])], b''), } assert_assignment(sticky_assignment, expected_assignment) subscriptions = { 'C1': {'t2'}, 'C2': {'t2'}, } member_metadata = {} for member, topics in six.iteritems(subscriptions): member_metadata[member] = StickyPartitionAssignor._metadata( topics, sticky_assignment[member].partitions()) sticky_assignment = StickyPartitionAssignor.assign(cluster, member_metadata) expected_assignment = { 'C1': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t2', [1])], b''), 'C2': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t2', [0, 2])], b''), } assert_assignment(sticky_assignment, expected_assignment)
def assign(self, cluster: ClusterMetadata, members: Dict[str, ConsumerProtocolMemberMetadata]) \ -> Dict[str, ConsumerProtocolMemberAssignment]: logger.info('Statefulset Partition Assignor') logger.debug(f'Cluster = {cluster}\nMembers = {members}') # Get all topic all_topics: Set = set() for key, metadata in members.items(): all_topics.update(metadata.subscription) # Get all partitions by topic name all_topic_partitions = [] for topic in all_topics: partitions = cluster.partitions_for_topic(topic) if partitions is None: logger.warning('No partition metadata for topic %s', topic) continue for partition in partitions: all_topic_partitions.append(TopicPartition(topic, partition)) # Sort partition all_topic_partitions.sort() # Create default dict with lambda assignment: DefaultDict[str, Any] = collections.defaultdict( lambda: collections.defaultdict(list)) advanced_assignor_dict = self.get_advanced_assignor_dict( all_topic_partitions) for topic, partitions in advanced_assignor_dict.items(): for member_id, member_data in members.items(): # Loads member assignors data user_data = json.loads(member_data.user_data) # Get number of partitions by topic name topic_number_partitions = len(partitions) # Logic assignors if nb_replica as same as topic_numbers_partitions (used by StoreBuilder for # assign each partitions to right instance, in this case nb_replica is same as topic_number_partitions) if user_data['nb_replica'] == topic_number_partitions: if user_data['assignor_policy'] == 'all': for partition in partitions: assignment[member_id][topic].append(partition) elif user_data['assignor_policy'] == 'only_own': if user_data['instance'] in partitions: assignment[member_id][topic].append( partitions[user_data['instance']]) else: raise BadAssignorPolicy else: # Todo Add repartition raise NotImplementedError logger.debug(f'Assignment = {assignment}') protocol_assignment = {} for member_id in members: protocol_assignment[member_id] = ConsumerProtocolMemberAssignment( self.version, sorted(assignment[member_id].items()), members[member_id].user_data) logger.debug(f'Protocol Assignment = {protocol_assignment}') return protocol_assignment
def test_sticky_assignor1(mocker): """ Given: there are three consumers C0, C1, C2, four topics t0, t1, t2, t3, and each topic has 2 partitions, resulting in partitions t0p0, t0p1, t1p0, t1p1, t2p0, t2p1, t3p0, t3p1. Each consumer is subscribed to all three topics. Then: perform fresh assignment Expected: the assignment is - C0: [t0p0, t1p1, t3p0] - C1: [t0p1, t2p0, t3p1] - C2: [t1p0, t2p1] Then: remove C1 consumer and perform the reassignment Expected: the new assignment is - C0 [t0p0, t1p1, t2p0, t3p0] - C2 [t0p1, t1p0, t2p1, t3p1] """ cluster = create_cluster(mocker, topics={'t0', 't1', 't2', 't3'}, topics_partitions={0, 1}) subscriptions = { 'C0': {'t0', 't1', 't2', 't3'}, 'C1': {'t0', 't1', 't2', 't3'}, 'C2': {'t0', 't1', 't2', 't3'}, } member_metadata = make_member_metadata(subscriptions) sticky_assignment = StickyPartitionAssignor.assign(cluster, member_metadata) expected_assignment = { 'C0': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t0', [0]), ('t1', [1]), ('t3', [0])], b''), 'C1': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t0', [1]), ('t2', [0]), ('t3', [1])], b''), 'C2': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t1', [0]), ('t2', [1])], b''), } assert_assignment(sticky_assignment, expected_assignment) del subscriptions['C1'] member_metadata = {} for member, topics in six.iteritems(subscriptions): member_metadata[member] = StickyPartitionAssignor._metadata( topics, sticky_assignment[member].partitions()) sticky_assignment = StickyPartitionAssignor.assign(cluster, member_metadata) expected_assignment = { 'C0': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t0', [0]), ('t1', [1]), ('t2', [0]), ('t3', [0])], b''), 'C2': ConsumerProtocolMemberAssignment(StickyPartitionAssignor.version, [('t0', [1]), ('t1', [0]), ('t2', [1]), ('t3', [1])], b''), } assert_assignment(sticky_assignment, expected_assignment)
def assign(self, cluster: ClusterMetadata, members: Dict[str, ConsumerProtocolMemberMetadata]) \ -> Dict[str, ConsumerProtocolMemberAssignment]: """Assign function was call by aiokafka for assign consumer on right topic partition. Args: cluster (ClusterMetadata): Kafka-python cluster metadata (more detail in kafka-python documentation) members (Dict[str, ConsumerProtocolMemberMetadata]): members dict which contains ConsumerProtocolMemberMetadata (more detail in kafka-python documentation) Returns: Dict[str, ConsumerProtocolMemberAssignment]: dict which contain members and assignment protocol (more detail in kafka-python documentation) """ self.logger.info('Statefulset Partition Assignor') self.logger.debug('Cluster = %s\nMembers = %s', cluster, members) # Get all topic all_topics: Set = set() for key, metadata in members.items(): self.logger.debug('Key = %s\nMetadata = %s', key, metadata) all_topics.update(metadata.subscription) # Get all partitions by topic name all_topic_partitions = [] for topic in all_topics: partitions = cluster.partitions_for_topic(topic) if partitions is None: self.logger.warning('No partition metadata for topic %s', topic) continue for partition in partitions: all_topic_partitions.append(TopicPartition(topic, partition)) # Sort partition all_topic_partitions.sort() # Create default dict with lambda assignment: DefaultDict[str, Any] = collections.defaultdict(lambda: collections.defaultdict(list)) advanced_assignor_dict = self.get_advanced_assignor_dict(all_topic_partitions) for topic, partitions in advanced_assignor_dict.items(): for member_id, member_data in members.items(): # Loads member assignors data user_data = json.loads(member_data.user_data) # Get number of partitions by topic name topic_number_partitions = len(partitions) # Logic assignors if nb_replica as same as topic_numbers_partitions (used by StoreBuilder for # assign each partitions to right instance, in this case nb_replica is same as topic_number_partitions) if user_data['nb_replica'] == topic_number_partitions: if user_data['assignor_policy'] == 'all': for partition in partitions: assignment[member_id][topic].append(partition) elif user_data['assignor_policy'] == 'only_own': if user_data['instance'] in partitions: assignment[member_id][topic].append(partitions[user_data['instance']]) else: raise BadAssignorPolicy else: raise NotImplementedError self.logger.debug('Assignment = %s', assignment) protocol_assignment = {} for member_id in members: protocol_assignment[member_id] = ConsumerProtocolMemberAssignment(self.version, sorted(assignment[member_id].items()), members[member_id].user_data) self.logger.debug('Protocol Assignment = %s', protocol_assignment) return protocol_assignment