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