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)
Beispiel #4
0
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
Beispiel #5
0
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)]))
Beispiel #6
0
    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
Beispiel #7
0
    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)
Beispiel #10
0
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()
Beispiel #11
0
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()
Beispiel #12
0
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
Beispiel #13
0
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)
Beispiel #14
0
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)
Beispiel #15
0
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)
Beispiel #16
0
 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()
     }
Beispiel #17
0
    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
Beispiel #18
0
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
Beispiel #19
0
    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
Beispiel #20
0
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)
Beispiel #21
0
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)
Beispiel #22
0
    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
Beispiel #23
0
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)
Beispiel #24
0
    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