kafka-python 3.0.0__py3-none-any.whl

kafka/coordinator/assignors/abstract.py

@@ -0,0 +1,90 @@
+from abc import ABC, abstractmethod, abstractproperty
+from enum import IntEnum
+
+from kafka.protocol.consumer.metadata import (
+    ConsumerProtocolSubscription, ConsumerProtocolAssignment,
+)
+
+
+class RebalanceProtocol(IntEnum):
+    """KIP-429: rebalance protocol mode for a partition assignor.
+
+    EAGER - pre-KIP-429 behaviour: every member revokes its full
+    assignment before JoinGroup, then receives a fresh assignment in
+    SyncGroup. Simple but causes a "stop the world" pause on every
+    rebalance.
+
+    COOPERATIVE - KIP-429 incremental rebalance: members keep their
+    existing assignment across JoinGroup; the leader's assignment
+    indicates the partitions that need to move; only revoked
+    partitions are released, and only newly-assigned partitions
+    invoke the listener. A second rebalance round picks up partitions
+    that were revoked in round 1.
+    """
+    EAGER = 0
+    COOPERATIVE = 1
+
+
+class AbstractPartitionAssignor(ABC):
+    """
+    Abstract assignor implementation which does some common grunt work (in particular collecting
+    partition counts which are always needed in assignors).
+    """
+
+    @abstractproperty
+    def name(self):
+        """.name should be a string identifying the assignor"""
+        pass
+
+    def supported_protocols(self):
+        """Return the list of :class:`RebalanceProtocol` modes this
+        assignor supports, in order of preference.
+
+        Default is ``[EAGER]`` - every legacy assignor (Range,
+        RoundRobin, the original Sticky from KIP-54) behaves this
+        way. Override in subclasses that participate in KIP-429
+        incremental cooperative rebalancing (e.g.
+        ``CooperativeStickyAssignor``).
+        """
+        return [RebalanceProtocol.EAGER]
+
+    @abstractmethod
+    def assign(self, cluster, members):
+        """Perform group assignment given cluster metadata and member subscriptions
+
+        Arguments:
+            cluster (ClusterMetadata): metadata for use in assignment
+            members ([JoinGroupResponseMember]): member_id and metadata
+                for each member in the group, including group_instance_id
+                when available (v5+). metadata is a decoded instance of
+                ConsumerProtocolSubscription.
+
+        Returns:
+            dict: {member_id: ConsumerProtocolAssignment}
+        """
+        pass
+
+    @abstractmethod
+    def metadata(self, topics):
+        """Generate ProtocolMetadata to be submitted via JoinGroupRequest.
+
+        Arguments:
+            topics (set): a member's subscribed topics
+
+        Returns:
+            ConsumerProtocolSubscription
+        """
+        pass
+
+    @abstractmethod
+    def on_assignment(self, assignment, generation):
+        """Callback that runs on each assignment.
+
+        This method can be used to update internal state, if any, of the
+        partition assignor.
+
+        Arguments:
+            assignment (ConsumerProtocolAssignment): the member's assignment
+            generation (int): generation id of assignment
+        """
+        pass

kafka/coordinator/assignors/cooperative_sticky.py

@@ -0,0 +1,167 @@
+"""KIP-429 cooperative sticky partition assignor.
+
+Wraps :class:`StickyPartitionAssignor` (KIP-54) with the two-phase
+"incremental cooperative" rebalancing protocol:
+
+  * Members keep their assignment across JoinGroup - no global revoke.
+  * The leader runs the sticky algorithm to compute the *ideal* final
+    assignment, then identifies any partition that is moving from one
+    owner to another and *removes it from the new owner's first-round
+    assignment*. The current owner sees its assignment shrink, revokes
+    the lost partition, and the broker is signaled (via
+    ``request_rejoin``) that another rebalance is needed.
+  * Round two: the freshly-revoked partition is owned by nobody; the
+    sticky algorithm now gives it to its intended new owner.
+
+This avoids the "stop the world" pause that EAGER mode imposes - each
+member only pauses while it's processing the specific partitions
+moving in or out of its own assignment.
+
+References:
+  * KIP-429: https://cwiki.apache.org/confluence/x/vAclBg
+  * Java: org.apache.kafka.clients.consumer.CooperativeStickyAssignor
+"""
+
+from collections import defaultdict
+
+from kafka.coordinator.assignors.abstract import RebalanceProtocol
+from kafka.coordinator.assignors.sticky.sticky_assignor import (
+    StickyAssignmentExecutor,
+    StickyAssignorMemberMetadataV1,
+    StickyPartitionAssignor,
+)
+from kafka.protocol.consumer.metadata import (
+    ConsumerProtocolAssignment, ConsumerProtocolSubscription,
+)
+from kafka.structs import TopicPartition
+
+
+# Wire version 1 of ConsumerProtocolSubscription is what KIP-429 added.
+# Members advertise their currently-owned partitions in the
+# ``owned_partitions`` field so the leader can compute the diff.
+_COOPERATIVE_SUBSCRIPTION_VERSION = 1
+
+
+class CooperativeStickyAssignor(StickyPartitionAssignor):
+    """KIP-429 cooperative variant of the sticky assignor.
+
+    Behaviorally identical to :class:`StickyPartitionAssignor` for
+    final partition placement (it inherits the same algorithm) - the
+    only difference is that movements are staged across two rebalance
+    rounds so no member ever sees a partition assigned to it while
+    another member still owns it.
+    """
+
+    name = "cooperative-sticky"
+    # Bump the wire metadata to v1 so OwnedPartitions is encoded on
+    # JoinGroup. The leader reads .owned_partitions to compute the
+    # set of partitions that are moving.
+    version = _COOPERATIVE_SUBSCRIPTION_VERSION
+
+    def supported_protocols(self):
+        # COOPERATIVE only - mixing this assignor with eager assignors
+        # in the same consumer is rejected at consumer init time
+        # (see KafkaConsumer.__init__ validation).
+        return [RebalanceProtocol.COOPERATIVE]
+
+    def metadata(self, topics):
+        # Encode OwnedPartitions (v1+) so the leader can compute the
+        # cooperative diff. The base class uses StickyAssignorUserData
+        # for the same purpose in v0 - under cooperative we surface
+        # the owned set via the dedicated schema field instead.
+        SubTP = ConsumerProtocolSubscription.TopicPartition
+        owned_partitions = []
+        if self.member_assignment is not None:
+            by_topic = defaultdict(list)
+            for tp in self.member_assignment:
+                by_topic[tp.topic].append(tp.partition)
+            owned_partitions = [
+                SubTP(topic=t, partitions=sorted(parts))
+                for t, parts in by_topic.items()
+            ]
+        return ConsumerProtocolSubscription(
+            version=self.version,
+            topics=sorted(topics),
+            user_data=b'',
+            owned_partitions=owned_partitions,
+        )
+
+    @classmethod
+    def parse_member_metadata(cls, metadata):
+        """Decode a member's ``ConsumerProtocolSubscription``.
+
+        Cooperative members carry owned partitions in the
+        ``owned_partitions`` schema field (v1+) rather than the
+        ``user_data`` blob the legacy sticky assignor uses. Returns
+        the same ``StickyAssignorMemberMetadataV1`` shape so the
+        underlying sticky algorithm can consume it unchanged.
+        """
+        member_partitions = []
+        # owned_partitions is a list of TopicPartition data containers
+        # (v1+); on v0 metadata the field is absent - treat as empty.
+        for tp in getattr(metadata, 'owned_partitions', None) or ():
+            for partition in tp.partitions:
+                member_partitions.append(TopicPartition(tp.topic, partition))
+
+        generation = metadata.generation_id
+        return StickyAssignorMemberMetadataV1(
+            partitions=member_partitions,
+            generation=metadata.generation_id, # requires schema v2, defaults to -1
+            subscription=list(metadata.topics),
+        )
+
+    def assign(self, cluster, members):
+        """Cooperative two-phase assignment.
+
+        1. Compute the ideal final sticky assignment.
+        2. Build a map of currently-owned partitions across all
+           members from their ``OwnedPartitions``.
+        3. For any partition whose final owner differs from its
+           current owner, remove it from the new owner's first-round
+           assignment. The current owner sees its assignment shrink,
+           revokes the partition, and re-joins; on round two the
+           partition is unowned and the algorithm assigns it.
+        """
+        members_metadata = {
+            member.member_id: self.parse_member_metadata(member.metadata)
+            for member in members
+        }
+        executor = StickyAssignmentExecutor(cluster, members_metadata)
+        executor.perform_initial_assignment()
+        executor.balance()
+        # Expose for diagnostic tests (matches parent behaviour).
+        self._latest_partition_movements = executor.partition_movements
+
+        # Map: partition -> current_owner_member_id (None if unowned).
+        currently_owned = {}
+        for member_id, parsed in members_metadata.items():
+            for tp in parsed.partitions:
+                currently_owned[tp] = member_id
+
+        # Build the round-1 assignment: drop any partition that's
+        # moving (final owner != current owner). The current owner
+        # will revoke it in _on_join_complete; the broker will see
+        # the consumer re-join and round 2 will land the partition
+        # on the intended new owner.
+        #
+        # ``executor.get_final_assignment`` returns the canonical wire
+        # shape: ``[(topic, [partition, ...]), ...]``. We rebuild the
+        # same shape after filtering so the encoder is happy.
+        cooperative = {}
+        for member in members:
+            member_id = member.member_id
+            kept_by_topic = defaultdict(list)
+            for topic, parts in executor.get_final_assignment(member_id):
+                for p in parts:
+                    tp = TopicPartition(topic, p)
+                    current_owner = currently_owned.get(tp)
+                    if current_owner is None or current_owner == member_id:
+                        # Either nobody owned it, or this member
+                        # already owns it. Safe to assign now.
+                        kept_by_topic[topic].append(p)
+                    # else: partition is moving; defer to round 2.
+            assigned_partitions = sorted(
+                (t, sorted(ps)) for t, ps in kept_by_topic.items())
+            cooperative[member_id] = ConsumerProtocolAssignment(
+                self.version, assigned_partitions, b'')
+        return cooperative

kafka/coordinator/assignors/range.py

@@ -0,0 +1,81 @@
+import collections
+import itertools
+import logging
+
+from kafka.coordinator.assignors.abstract import (
+    AbstractPartitionAssignor,
+    ConsumerProtocolSubscription,
+    ConsumerProtocolAssignment,
+)
+
+log = logging.getLogger(__name__)
+
+
+class RangePartitionAssignor(AbstractPartitionAssignor):
+    """
+    The range assignor works on a per-topic basis. For each topic, we lay out
+    the available partitions in numeric order and the consumers in
+    lexicographic order. We then divide the number of partitions by the total
+    number of consumers to determine the number of partitions to assign to each
+    consumer. If it does not evenly divide, then the first few consumers will
+    have one extra partition.
+
+    For example, suppose there are two consumers C0 and C1, two topics t0 and
+    t1, and each topic has 3 partitions, resulting in partitions t0p0, t0p1,
+    t0p2, t1p0, t1p1, and t1p2.
+
+    The assignment will be:
+        C0: [t0p0, t0p1, t1p0, t1p1]
+        C1: [t0p2, t1p2]
+    """
+    name = 'range'
+    version = 0
+
+    @classmethod
+    def assign(cls, cluster, members):
+        consumers_per_topic = collections.defaultdict(list)
+        for member in members:
+            for topic in member.metadata.topics:
+                consumers_per_topic[topic].append((member.group_instance_id, member.member_id))
+
+        # construct {member_id: {topic: [partition, ...]}}
+        assignment = collections.defaultdict(dict)
+
+        for topic in consumers_per_topic:
+            # group by static members (True) v dynamic members (False)
+            grouped = {k: list(g) for k, g in itertools.groupby(consumers_per_topic[topic], key=lambda ids: ids[0] is not None)}
+            consumers_per_topic[topic] = sorted(grouped.get(True, [])) + sorted(grouped.get(False, [])) # sorted static members first, then sorted dynamic
+
+        for topic, consumers_for_topic in consumers_per_topic.items():
+            partitions = cluster.partitions_for_topic(topic)
+            if partitions is None:
+                log.warning('No partition metadata for topic %s', topic)
+                continue
+            partitions = sorted(partitions)
+
+            partitions_per_consumer = len(partitions) // len(consumers_for_topic)
+            consumers_with_extra = len(partitions) % len(consumers_for_topic)
+
+            for i, (_group_instance_id, member_id) 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_id][topic] = partitions[start:start+length]
+
+        protocol_assignment = {}
+        for member in members:
+            protocol_assignment[member.member_id] = ConsumerProtocolAssignment(
+                cls.version,
+                sorted(assignment[member.member_id].items()),
+                b'')
+        return protocol_assignment
+
+    @classmethod
+    def metadata(cls, topics):
+        return ConsumerProtocolSubscription(cls.version, list(topics), b'')
+
+    @classmethod
+    def on_assignment(cls, assignment, generation):
+        pass

kafka/coordinator/assignors/roundrobin.py

@@ -0,0 +1,101 @@
+import collections
+import itertools
+import logging
+
+from kafka.coordinator.assignors.abstract import (
+    AbstractPartitionAssignor,
+    ConsumerProtocolSubscription,
+    ConsumerProtocolAssignment,
+)
+from kafka.structs import TopicPartition
+
+log = logging.getLogger(__name__)
+
+
+class RoundRobinPartitionAssignor(AbstractPartitionAssignor):
+    """
+    The roundrobin assignor lays out all the available partitions and all the
+    available consumers. It then proceeds to do a roundrobin assignment from
+    partition to consumer. If the subscriptions of all consumer instances are
+    identical, then the partitions will be uniformly distributed. (i.e., the
+    partition ownership counts will be within a delta of exactly one across all
+    consumers.)
+
+    For example, suppose there are two consumers C0 and C1, two topics t0 and
+    t1, and each topic has 3 partitions, resulting in partitions t0p0, t0p1,
+    t0p2, t1p0, t1p1, and t1p2.
+
+    The assignment will be:
+        C0: [t0p0, t0p2, t1p1]
+        C1: [t0p1, t1p0, t1p2]
+
+    When subscriptions differ across consumer instances, the assignment process
+    still considers each consumer instance in round robin fashion but skips
+    over an instance if it is not subscribed to the topic. Unlike the case when
+    subscriptions are identical, this can result in imbalanced assignments.
+
+    For example, suppose we have three consumers C0, C1, C2, and three topics
+    t0, t1, t2, with unbalanced partitions t0p0, t1p0, t1p1, t2p0, t2p1, t2p2,
+    where C0 is subscribed to t0; C1 is subscribed to t0, t1; and C2 is
+    subscribed to t0, t1, t2.
+
+    The assignment will be:
+        C0: [t0p0]
+        C1: [t1p0]
+        C2: [t1p1, t2p0, t2p1, t2p2]
+    """
+    name = 'roundrobin'
+    version = 0
+
+    @classmethod
+    def assign(cls, cluster, members):
+        all_topics = set()
+        for member in members:
+            all_topics.update(member.metadata.topics)
+
+        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))
+
+        # Sort static and dynamic members separately to maintain stable static assignments
+        ungrouped = [(member.group_instance_id, member.member_id) for member in members]
+        grouped = {k: list(g) for k, g in itertools.groupby(ungrouped, key=lambda ids: ids[0] is not None)}
+        member_list = sorted(grouped.get(True, [])) + sorted(grouped.get(False, [])) # sorted static members first, then sorted dynamic
+        member_iter = itertools.cycle(member_list)
+        member_topics = {member.member_id: member.metadata.topics for member in members}
+
+        for partition in all_topic_partitions:
+            _group_instance_id, 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_topics[member_id]:
+                member_id = next(member_iter)
+            assignment[member_id][partition.topic].append(partition.partition)
+
+        protocol_assignment = {}
+        for member in members:
+            protocol_assignment[member.member_id] = ConsumerProtocolAssignment(
+                cls.version,
+                sorted(assignment[member.member_id].items()),
+                b'')
+        return protocol_assignment
+
+    @classmethod
+    def metadata(cls, topics):
+        return ConsumerProtocolSubscription(cls.version, list(topics), b'')
+
+    @classmethod
+    def on_assignment(cls, assignment, generation):
+        pass

kafka/coordinator/assignors/sticky/StickyAssignorUserData.json

@@ -0,0 +1,37 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+{
+  "type": "data",
+  "name": "StickyAssignorUserData",
+  // StickyAssignor currently always encodes with version 1.
+  // To decode, versions are attempted in reverse order until one succeeds.
+  // If no decoding is possible, the assignor ignores the previous user data.
+
+  // Version 1 added the "generation" field
+  "validVersions": "0-1",
+  "flexibleVersions": "none",
+  "fields": [
+    { "name": "PreviousAssignment", "type": "[]TopicPartition", "versions": "0+", "fields": [
+        { "name": "Topic", "type": "string", "mapKey": true, "versions": "0+", "entityType": "topicName",
+          "about": "The topic name."},
+        { "name": "Partitions", "type": "[]int32", "versions": "0+",
+          "about": "The partition ids."}
+      ]
+    },
+    { "name": "Generation", "type": "int32", "versions": "1+", "default": "-1", "ignorable": true,
+      "about": "The generation id of the previous assignment."}
+  ]
+}

kafka/coordinator/assignors/sticky/partition_movements.py

@@ -0,0 +1,149 @@
+import logging
+from collections import defaultdict, namedtuple
+from copy import deepcopy
+
+log = logging.getLogger(__name__)
+
+
+ConsumerPair = namedtuple("ConsumerPair", ["src_member_id", "dst_member_id"])
+"""
+Represents a pair of Kafka consumer ids involved in a partition reassignment.
+Each ConsumerPair corresponds to a particular partition or topic, indicates that the particular partition or some
+partition of the particular topic was moved from the source consumer to the destination consumer
+during the rebalance. This class helps in determining whether a partition reassignment results in cycles among
+the generated graph of consumer pairs.
+"""
+
+
+def is_sublist(source, target):
+    """Checks if one list is a sublist of another.
+
+    Arguments:
+      source: the list in which to search for the occurrence of target.
+      target: the list to search for as a sublist of source
+
+    Returns:
+      true if target is in source; false otherwise
+    """
+    for index in (i for i, e in enumerate(source) if e == target[0]):
+        if tuple(source[index: index + len(target)]) == target:
+            return True
+    return False
+
+
+class PartitionMovements:
+    """
+    This class maintains some data structures to simplify lookup of partition movements among consumers.
+    At each point of time during a partition rebalance it keeps track of partition movements
+    corresponding to each topic, and also possible movement (in form a ConsumerPair object) for each partition.
+    """
+
+    def __init__(self):
+        self.partition_movements_by_topic = defaultdict(
+            lambda: defaultdict(set)
+        )
+        self.partition_movements = {}
+
+    def move_partition(self, partition, old_consumer, new_consumer):
+        pair = ConsumerPair(src_member_id=old_consumer, dst_member_id=new_consumer)
+        if partition in self.partition_movements:
+            # this partition has previously moved
+            existing_pair = self._remove_movement_record_of_partition(partition)
+            if existing_pair.dst_member_id != old_consumer:
+                raise ValueError()
+            if existing_pair.src_member_id != new_consumer:
+                # the partition is not moving back to its previous consumer
+                self._add_partition_movement_record(
+                    partition, ConsumerPair(src_member_id=existing_pair.src_member_id, dst_member_id=new_consumer)
+                )
+        else:
+            self._add_partition_movement_record(partition, pair)
+
+    def get_partition_to_be_moved(self, partition, old_consumer, new_consumer):
+        if partition.topic not in self.partition_movements_by_topic:
+            return partition
+        if partition in self.partition_movements:
+            # this partition has previously moved
+            if old_consumer != self.partition_movements[partition].dst_member_id:
+                raise ValueError()
+            old_consumer = self.partition_movements[partition].src_member_id
+        reverse_pair = ConsumerPair(src_member_id=new_consumer, dst_member_id=old_consumer)
+        if reverse_pair not in self.partition_movements_by_topic[partition.topic]:
+            return partition
+
+        return next(iter(self.partition_movements_by_topic[partition.topic][reverse_pair]))
+
+    def are_sticky(self):
+        for topic, movements in self.partition_movements_by_topic.items():
+            movement_pairs = set(movements.keys())
+            if self._has_cycles(movement_pairs):
+                log.error(
+                    "Stickiness is violated for topic {}\n"
+                    "Partition movements for this topic occurred among the following consumer pairs:\n"
+                    "{}".format(topic, movement_pairs)
+                )
+                return False
+        return True
+
+    def _remove_movement_record_of_partition(self, partition):
+        pair = self.partition_movements[partition]
+        del self.partition_movements[partition]
+
+        self.partition_movements_by_topic[partition.topic][pair].remove(partition)
+        if not self.partition_movements_by_topic[partition.topic][pair]:
+            del self.partition_movements_by_topic[partition.topic][pair]
+        if not self.partition_movements_by_topic[partition.topic]:
+            del self.partition_movements_by_topic[partition.topic]
+
+        return pair
+
+    def _add_partition_movement_record(self, partition, pair):
+        self.partition_movements[partition] = pair
+        self.partition_movements_by_topic[partition.topic][pair].add(partition)
+
+    def _has_cycles(self, consumer_pairs):
+        cycles = set()
+        for pair in consumer_pairs:
+            reduced_pairs = deepcopy(consumer_pairs)
+            reduced_pairs.remove(pair)
+            path = [pair.src_member_id]
+            if self._is_linked(pair.dst_member_id, pair.src_member_id, reduced_pairs, path) and not self._is_subcycle(
+                path, cycles
+            ):
+                cycles.add(tuple(path))
+                log.error("A cycle of length {} was found: {}".format(len(path) - 1, path))
+
+        # for now we want to make sure there is no partition movements of the same topic between a pair of consumers.
+        # the odds of finding a cycle among more than two consumers seem to be very low (according to various randomized
+        # tests with the given sticky algorithm) that it should not worth the added complexity of handling those cases.
+        for cycle in cycles:
+            if len(cycle) == 3:  # indicates a cycle of length 2
+                return True
+        return False
+
+    @staticmethod
+    def _is_subcycle(cycle, cycles):
+        super_cycle = deepcopy(cycle)
+        super_cycle = super_cycle[:-1]
+        super_cycle.extend(cycle)
+        for found_cycle in cycles:
+            if len(found_cycle) == len(cycle) and is_sublist(super_cycle, found_cycle):
+                return True
+        return False
+
+    def _is_linked(self, src, dst, pairs, current_path):
+        if src == dst:
+            return False
+        if not pairs:
+            return False
+        if ConsumerPair(src, dst) in pairs:
+            current_path.append(src)
+            current_path.append(dst)
+            return True
+        for pair in pairs:
+            if pair.src_member_id == src:
+                reduced_set = deepcopy(pairs)
+                reduced_set.remove(pair)
+                current_path.append(pair.src_member_id)
+                return self._is_linked(pair.dst_member_id, dst, reduced_set, current_path)
+        return False