Implement KIP-345 in aiokafka (rebase of #827)

aiokafka/consumer/assignors.py

@@ -0,0 +1,24 @@
+from aiokafka.coordinator.assignors.abstract import AbstractPartitionAssignor
+import abc
+
+
+class AbstractStaticPartitionAssignor(AbstractPartitionAssignor):
+    """
+    Abstract assignor implementation that also supports static assignments (KIP-345)
+    """
+
+    @abc.abstractmethod
+    def assign(self, cluster, members,  # lgtm[py/inheritance/signature-mismatch]
+               member_group_instance_ids):
+        """Perform group assignment given cluster metadata, member subscriptions
+           and group_instance_ids
+        Arguments:
+            cluster (ClusterMetadata): metadata for use in assignment
+            members (dict of {member_id: MemberMetadata}): decoded metadata for
+                each member in the group.
+            member_group_instance_ids (dict of {member_id: MemberMetadata}):
+                decoded metadata for each member in the group.
+        Returns:
+            dict: {member_id: MemberAssignment}
+        """
+        pass

aiokafka/consumer/consumer.py

@@ -65,6 +65,8 @@ class AIOKafkaConsumer:
             committing offsets. If None, auto-partition assignment (via
             group coordinator) and offset commits are disabled.
             Default: None
+        group_instance_id (str or None): name of the group instance ID used for
+            static membership (KIP-345)
         key_deserializer (Callable): Any callable that takes a
             raw message key and returns a deserialized key.
         value_deserializer (Callable, Optional): Any callable that takes a
@@ -229,6 +231,7 @@ def __init__(self, *topics, loop=None,
                  bootstrap_servers='localhost',
                  client_id='aiokafka-' + __version__,
                  group_id=None,
+                 group_instance_id=None,
                  key_deserializer=None, value_deserializer=None,
                  fetch_max_wait_ms=500,
                  fetch_max_bytes=52428800,
@@ -291,6 +294,7 @@ def __init__(self, *topics, loop=None,
             sasl_oauth_token_provider=sasl_oauth_token_provider)
 
         self._group_id = group_id
+        self._group_instance_id = group_instance_id
         self._heartbeat_interval_ms = heartbeat_interval_ms
         self._session_timeout_ms = session_timeout_ms
         self._retry_backoff_ms = retry_backoff_ms
@@ -382,6 +386,7 @@ async def start(self):
             self._coordinator = GroupCoordinator(
                 self._client, self._subscription,
                 group_id=self._group_id,
+                group_instance_id=self._group_instance_id,
                 heartbeat_interval_ms=self._heartbeat_interval_ms,
                 session_timeout_ms=self._session_timeout_ms,
                 retry_backoff_ms=self._retry_backoff_ms,

aiokafka/consumer/group_coordinator.py

@@ -15,13 +15,15 @@
     HeartbeatRequest, JoinGroupRequest, LeaveGroupRequest, SyncGroupRequest)
 from aiokafka.structs import OffsetAndMetadata, TopicPartition
 from aiokafka.util import create_future, create_task
+from aiokafka.consumer.assignors import AbstractStaticPartitionAssignor
+
 
 log = logging.getLogger(__name__)
 
 UNKNOWN_OFFSET = -1
 
 
-class BaseCoordinator:
+class BaseCoordinator(object):
 
     def __init__(self, client, subscription, *,
                  exclude_internal_topics=True):
@@ -209,6 +211,7 @@ class GroupCoordinator(BaseCoordinator):
 
     def __init__(self, client, subscription, *,
                  group_id='aiokafka-default-group',
+                 group_instance_id=None,
                  session_timeout_ms=10000, heartbeat_interval_ms=3000,
                  retry_backoff_ms=100,
                  enable_auto_commit=True, auto_commit_interval_ms=5000,
@@ -240,6 +243,7 @@ def __init__(self, client, subscription, *,
         self.generation = OffsetCommitRequest.DEFAULT_GENERATION_ID
         self.member_id = JoinGroupRequest[0].UNKNOWN_MEMBER_ID
         self.group_id = group_id
+        self._group_instance_id = group_instance_id
         self.coordinator_id = None
 
         # Coordination flags and futures
@@ -345,9 +349,11 @@ def maybe_leave_group(self):
         return task
 
     async def _maybe_leave_group(self):
-        if self.generation > 0:
+        if self.generation > 0 and self._group_instance_id is None:
             # this is a minimal effort attempt to leave the group. we do not
             # attempt any resending if the request fails or times out.
+            # Note: do not send this leave request if we are running in static
+            # partition assignment mode (when group_instance_id has been set).
             version = 0 if self._client.api_version < (0, 11, 0) else 1
             request = LeaveGroupRequest[version](self.group_id, self.member_id)
             try:
@@ -401,8 +407,18 @@ async def _perform_assignment(
             'Invalid assignment protocol: %s' % assignment_strategy
         member_metadata = {}
         all_subscribed_topics = set()
-        for member_id, metadata_bytes in members:
+        group_instance_id_mapping = {}
+        # for member_id, group_instance_id, metadata_bytes in members:
+        for member in members:
+            if len(member) == 2:
+                member_id, metadata_bytes = member
+                group_instance_id = None
+            elif len(member) == 3:
+                member_id, group_instance_id, metadata_bytes = member
+            else:
+                raise Exception("unknown protocol returned from assignment")
             metadata = ConsumerProtocol.METADATA.decode(metadata_bytes)
+            group_instance_id_mapping = group_instance_id
             member_metadata[member_id] = metadata
             all_subscribed_topics.update(metadata.subscription)
 
@@ -420,7 +436,12 @@ async def _perform_assignment(
                   " with subscriptions %s", self.group_id, assignor.name,
                   member_metadata)
 
-        assignments = assignor.assign(self._cluster, member_metadata)
+        if isinstance(assignor, AbstractStaticPartitionAssignor):
+            assignments = assignor.assign(
+                self._cluster, member_metadata, group_instance_id_mapping
+            )
+        else:
+            assignments = assignor.assign(self._cluster, member_metadata)
         log.debug("Finished assignment for group %s: %s",
                   self.group_id, assignments)
 
@@ -1202,46 +1223,67 @@ async def perform_group_join(self):
                 metadata = metadata.encode()
             group_protocol = (assignor.name, metadata)
             metadata_list.append(group_protocol)
+            # for KIP-394 we may have to send a second join request
+            try_join = True
+            while try_join:
+                try_join = False
+
+                if self._api_version < (0, 10, 1):
+                    request = JoinGroupRequest[0](
+                        self.group_id,
+                        self._session_timeout_ms,
+                        self._coordinator.member_id,
+                        ConsumerProtocol.PROTOCOL_TYPE,
+                        metadata_list,
+                    )
+                elif self._api_version < (0, 11, 0):
+                    request = JoinGroupRequest[1](
+                        self.group_id,
+                        self._session_timeout_ms,
+                        self._rebalance_timeout_ms,
+                        self._coordinator.member_id,
+                        ConsumerProtocol.PROTOCOL_TYPE,
+                        metadata_list,
+                    )
+                elif self._api_version < (2, 3, 0):
+                    request = JoinGroupRequest[2](
+                        self.group_id,
+                        self._session_timeout_ms,
+                        self._rebalance_timeout_ms,
+                        self._coordinator.member_id,
+                        ConsumerProtocol.PROTOCOL_TYPE,
+                        metadata_list,
+                    )
+                else:
+                    request = JoinGroupRequest[3](
+                        self.group_id,
+                        self._session_timeout_ms,
+                        self._rebalance_timeout_ms,
+                        self._coordinator.member_id,
+                        self._coordinator._group_instance_id,
+                        ConsumerProtocol.PROTOCOL_TYPE,
+                        metadata_list,
+                    )
 
-        if self._api_version < (0, 10, 1):
-            request = JoinGroupRequest[0](
-                self.group_id,
-                self._session_timeout_ms,
-                self._coordinator.member_id,
-                ConsumerProtocol.PROTOCOL_TYPE,
-                metadata_list)
-        elif self._api_version < (0, 11, 0):
-            request = JoinGroupRequest[1](
-                self.group_id,
-                self._session_timeout_ms,
-                self._rebalance_timeout_ms,
-                self._coordinator.member_id,
-                ConsumerProtocol.PROTOCOL_TYPE,
-                metadata_list)
-        else:
-            request = JoinGroupRequest[2](
-                self.group_id,
-                self._session_timeout_ms,
-                self._rebalance_timeout_ms,
-                self._coordinator.member_id,
-                ConsumerProtocol.PROTOCOL_TYPE,
-                metadata_list)
+                # create the request for the coordinator
+                log.debug("Sending JoinGroup (%s) to coordinator %s",
+                          request, self.coordinator_id)
+                try:
+                    response = await self._coordinator._send_req(request)
+                except Errors.KafkaError:
+                    # Return right away. It's a connection error, so backoff will be
+                    # handled by coordinator lookup
+                    return None
+                if not self._subscription.active:
+                    # Subscription changed. Ignore response and restart group join
+                    return None
 
-        # create the request for the coordinator
-        log.debug("Sending JoinGroup (%s) to coordinator %s",
-                  request, self.coordinator_id)
-        try:
-            response = await self._coordinator._send_req(request)
-        except Errors.KafkaError:
-            # Return right away. It's a connection error, so backoff will be
-            # handled by coordinator lookup
-            return None
+                error_type = Errors.for_code(response.error_code)
 
-        if not self._subscription.active:
-            # Subscription changed. Ignore response and restart group join
-            return None
+                if error_type is Errors.MemberIdRequired:
+                    self._coordinator.member_id = response.member_id
+                    try_join = True
 
-        error_type = Errors.for_code(response.error_code)
         if error_type is Errors.NoError:
             log.debug("Join group response %s", response)
             self._coordinator.member_id = response.member_id
@@ -1300,12 +1342,22 @@ async def perform_group_join(self):
 
     async def _on_join_follower(self):
         # send follower's sync group with an empty assignment
-        version = 0 if self._api_version < (0, 11, 0) else 1
-        request = SyncGroupRequest[version](
-            self.group_id,
-            self._coordinator.generation,
-            self._coordinator.member_id,
-            [])
+        if self._api_version < (2, 3, 0):
+            version = 0 if self._api_version < (0, 11, 0) else 1
+            request = SyncGroupRequest[version](
+                self.group_id,
+                self._coordinator.generation,
+                self._coordinator.member_id,
+                [],
+            )
+        else:
+            request = SyncGroupRequest[2](
+                self.group_id,
+                self._coordinator.generation,
+                self._coordinator.member_id,
+                self._coordinator._group_instance_id,
+                [],
+            )
         log.debug(
             "Sending follower SyncGroup for group %s to coordinator %s: %s",
             self.group_id, self.coordinator_id, request)
@@ -1337,12 +1389,22 @@ async def _on_join_leader(self, response):
                 assignment = assignment.encode()
             assignment_req.append((member_id, assignment))
 
-        version = 0 if self._api_version < (0, 11, 0) else 1
-        request = SyncGroupRequest[version](
-            self.group_id,
-            self._coordinator.generation,
-            self._coordinator.member_id,
-            assignment_req)
+        if self._api_version < (2, 3, 0):
+            version = 0 if self._api_version < (0, 11, 0) else 1
+            request = SyncGroupRequest[version](
+                self.group_id,
+                self._coordinator.generation,
+                self._coordinator.member_id,
+                assignment_req,
+            )
+        else:
+            request = SyncGroupRequest[2](
+                self.group_id,
+                self._coordinator.generation,
+                self._coordinator.member_id,
+                self._coordinator._group_instance_id,
+                assignment_req,
+            )
 
         log.debug(
             "Sending leader SyncGroup for group %s to coordinator %s: %s",

aiokafka/consumer/subscription_state.py

@@ -14,8 +14,6 @@
 
 log = logging.getLogger(__name__)
 
-(List, Future)
-
 
 class SubscriptionType(Enum):
 

aiokafka/errors.py

@@ -856,6 +856,15 @@ class ListenerNotFound(BrokerResponseError):
     )
 
 
+class MemberIdRequired(BrokerResponseError):
+    errno = 79
+    message = 'MEMBER_ID_REQUIRED'
+    description = (
+        'Consumer needs to have a valid member '
+        'id before actually entering group'
+    )
+
+
 def _iter_broker_errors():
     for name, obj in inspect.getmembers(sys.modules[__name__]):
         if (