Use Subscription wrapper to hold group_instance_id + member metadata

Author: dpkp

kafka/coordinator/assignors/abstract.py

@@ -23,8 +23,9 @@ def assign(self, cluster, members):
 
         Arguments:
             cluster (ClusterMetadata): metadata for use in assignment
-            members (dict of {member_id: MemberMetadata}): decoded metadata for
-                each member in the group.
+            members (dict of {member_id: Subscription}): decoded metadata
+                for each member in the group, including group_instance_id
+                when available.
 
         Returns:
             dict: {member_id: MemberAssignment}

kafka/coordinator/assignors/range.py

@@ -32,10 +32,10 @@ class RangePartitionAssignor(AbstractPartitionAssignor):
     version = 0
 
     @classmethod
-    def assign(cls, cluster, member_metadata):
+    def assign(cls, cluster, group_subscriptions):
         consumers_per_topic = collections.defaultdict(list)
-        for member, metadata in six.iteritems(member_metadata):
-            for topic in metadata.subscription:
+        for member, subscription in six.iteritems(group_subscriptions):
+            for topic in subscription.subscription:
                 consumers_per_topic[topic].append(member)
 
         # construct {member_id: {topic: [partition, ...]}}
@@ -61,7 +61,7 @@ def assign(cls, cluster, member_metadata):
                 assignment[member][topic] = partitions[start:start+length]
 
         protocol_assignment = {}
-        for member_id in member_metadata:
+        for member_id in group_subscriptions:
             protocol_assignment[member_id] = ConsumerProtocolMemberAssignment_v0(
                 cls.version,
                 sorted(assignment[member_id].items()),

kafka/coordinator/assignors/roundrobin.py

@@ -49,10 +49,10 @@ class RoundRobinPartitionAssignor(AbstractPartitionAssignor):
     version = 0
 
     @classmethod
-    def assign(cls, cluster, member_metadata):
+    def assign(cls, cluster, group_subscriptions):
         all_topics = set()
-        for metadata in six.itervalues(member_metadata):
-            all_topics.update(metadata.subscription)
+        for subscription in six.itervalues(group_subscriptions):
+            all_topics.update(subscription.subscription)
 
         all_topic_partitions = []
         for topic in all_topics:
@@ -67,20 +67,20 @@ def assign(cls, cluster, member_metadata):
         # construct {member_id: {topic: [partition, ...]}}
         assignment = collections.defaultdict(lambda: collections.defaultdict(list))
 
-        member_iter = itertools.cycle(sorted(member_metadata.keys()))
+        member_iter = itertools.cycle(sorted(group_subscriptions.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:
+            while partition.topic not in group_subscriptions[member_id].subscription:
                 member_id = next(member_iter)
             assignment[member_id][partition.topic].append(partition.partition)
 
         protocol_assignment = {}
-        for member_id in member_metadata:
+        for member_id in group_subscriptions:
             protocol_assignment[member_id] = ConsumerProtocolMemberAssignment_v0(
                 cls.version,
                 sorted(assignment[member_id].items()),

kafka/coordinator/consumer.py

@@ -13,6 +13,7 @@
 from kafka.coordinator.assignors.roundrobin import RoundRobinPartitionAssignor
 from kafka.coordinator.assignors.sticky.sticky_assignor import StickyPartitionAssignor
 from kafka.coordinator.protocol import ConsumerProtocol
+from kafka.coordinator.subscription import Subscription
 import kafka.errors as Errors
 from kafka.future import Future
 from kafka.metrics import AnonMeasurable
@@ -333,17 +334,20 @@ def time_to_next_poll(self):
     def _perform_assignment(self, leader_id, assignment_strategy, members):
         assignor = self._lookup_assignor(assignment_strategy)
         assert assignor, 'Invalid assignment protocol: %s' % (assignment_strategy,)
-        member_metadata = {}
+        member_subscriptions = {}
         all_subscribed_topics = set()
         for member in members:
             if len(member) == 3:
                 member_id, group_instance_id, metadata_bytes = member
             else:
                 member_id, metadata_bytes = member
                 group_instance_id = None
-            metadata = ConsumerProtocol[0].METADATA.decode(metadata_bytes),
-            member_metadata[member_id] = metadata
-            all_subscribed_topics.update(metadata.subscription) # pylint: disable-msg=no-member
+            subscription = Subscription(
+                ConsumerProtocol[0].METADATA.decode(metadata_bytes),
+                group_instance_id
+            )
+            member_subscriptions[member_id] = subscription
+            all_subscribed_topics.update(subscription.subscription)
 
         # the leader will begin watching for changes to any of the topics
         # the group is interested in, which ensures that all metadata changes
@@ -361,9 +365,9 @@ def _perform_assignment(self, leader_id, assignment_strategy, members):
 
         log.debug("Performing assignment for group %s using strategy %s"
                   " with subscriptions %s", self.group_id, assignor.name,
-                  member_metadata)
+                  member_subscriptions)
 
-        assignments = assignor.assign(self._cluster, member_metadata)
+        assignments = assignor.assign(self._cluster, member_subscriptions)
 
         log.debug("Finished assignment for group %s: %s", self.group_id, assignments)
 

kafka/coordinator/subscription.py

@@ -0,0 +1,34 @@
+from __future__ import absolute_import
+
+
+class Subscription(object):
+    __slots__ = ('_metadata', '_group_instance_id')
+    def __init__(self, metadata, group_instance_id):
+        self._metadata = metadata
+        self._group_instance_id = group_instance_id
+
+    @property
+    def version(self):
+        return self._metadata.version
+
+    @property
+    def user_data(self):
+        return self._metadata.user_data
+
+    @property
+    def subscription(self):
+        return self._metadata.subscription
+
+    @property
+    def group_instance_id(self):
+        return self._group_instance_id
+
+    def encode(self):
+        return self._metadata.encode()
+
+    def __eq__(self, other):
+        return (
+            isinstance(other, Subscription) and
+            self._metadata == other._metadata and
+            self._group_instance_id == other._group_instance_id
+        )

test/test_assignors.py

@@ -11,6 +11,7 @@
 from kafka.coordinator.assignors.roundrobin import RoundRobinPartitionAssignor
 from kafka.coordinator.assignors.sticky.sticky_assignor import StickyPartitionAssignor
 from kafka.coordinator.protocol import ConsumerProtocolMemberAssignment_v0
+from kafka.coordinator.subscription import Subscription
 from kafka.vendor import six
 
 
@@ -34,13 +35,13 @@ def create_cluster(mocker, topics, topics_partitions=None, topic_partitions_lamb
 def test_assignor_roundrobin(mocker):
     assignor = RoundRobinPartitionAssignor
 
-    member_metadata = {
-        'C0': assignor.metadata({'t0', 't1'}),
-        'C1': assignor.metadata({'t0', 't1'}),
+    group_subscriptions = {
+        'C0': Subscription(assignor.metadata({'t0', 't1'}), None),
+        'C1': Subscription(assignor.metadata({'t0', 't1'}), None),
     }
 
     cluster = create_cluster(mocker, {'t0', 't1'}, topics_partitions={0, 1, 2})
-    ret = assignor.assign(cluster, member_metadata)
+    ret = assignor.assign(cluster, group_subscriptions)
     expected = {
         'C0': ConsumerProtocolMemberAssignment_v0(
             assignor.version, [('t0', [0, 2]), ('t1', [1])], b''),
@@ -56,13 +57,13 @@ def test_assignor_roundrobin(mocker):
 def test_assignor_range(mocker):
     assignor = RangePartitionAssignor
 
-    member_metadata = {
-        'C0': assignor.metadata({'t0', 't1'}),
-        'C1': assignor.metadata({'t0', 't1'}),
+    group_subscriptions = {
+        'C0': Subscription(assignor.metadata({'t0', 't1'}), None),
+        'C1': Subscription(assignor.metadata({'t0', 't1'}), None),
     }
 
     cluster = create_cluster(mocker, {'t0', 't1'}, topics_partitions={0, 1, 2})
-    ret = assignor.assign(cluster, member_metadata)
+    ret = assignor.assign(cluster, group_subscriptions)
     expected = {
         'C0': ConsumerProtocolMemberAssignment_v0(
             assignor.version, [('t0', [0, 1]), ('t1', [0, 1])], b''),

test/test_coordinator.py

@@ -14,6 +14,7 @@
 from kafka.coordinator.consumer import ConsumerCoordinator
 from kafka.coordinator.protocol import (
     ConsumerProtocolMemberMetadata_v0, ConsumerProtocolMemberAssignment_v0)
+from kafka.coordinator.subscription import Subscription
 import kafka.errors as Errors
 from kafka.future import Future
 from kafka.protocol.broker_api_versions import BROKER_API_VERSIONS
@@ -192,9 +193,9 @@ def test_subscription_listener_failure(mocker, coordinator):
 
 def test_perform_assignment(mocker, coordinator):
     coordinator._subscription.subscribe(topics=['foo1'])
-    member_metadata = {
-        'member-foo': ConsumerProtocolMemberMetadata_v0(0, ['foo1'], b''),
-        'member-bar': ConsumerProtocolMemberMetadata_v0(0, ['foo1'], b'')
+    group_subscriptions = {
+        'member-foo': Subscription(ConsumerProtocolMemberMetadata_v0(0, ['foo1'], b''), None),
+        'member-bar': Subscription(ConsumerProtocolMemberMetadata_v0(0, ['foo1'], b''), None),
     }
     assignments = {
         'member-foo': ConsumerProtocolMemberAssignment_v0(
@@ -208,12 +209,12 @@ def test_perform_assignment(mocker, coordinator):
 
     ret = coordinator._perform_assignment(
         'member-foo', 'roundrobin',
-        [(member, metadata.encode())
-         for member, metadata in member_metadata.items()])
+        [(member, subscription.encode())
+         for member, subscription in group_subscriptions.items()])
 
     assert RoundRobinPartitionAssignor.assign.call_count == 1
     RoundRobinPartitionAssignor.assign.assert_called_with(
-        coordinator._client.cluster, member_metadata)
+        coordinator._client.cluster, group_subscriptions)
     assert ret == assignments