Support FindCoordinatorRequest v2 in consumer and admin client (#2502)

Author: dpkp

kafka/admin/client.py

@@ -23,7 +23,8 @@
     ListGroupsRequest, DescribeGroupsRequest, DescribeAclsRequest, CreateAclsRequest, DeleteAclsRequest,
     DeleteGroupsRequest, DescribeLogDirsRequest
 )
-from kafka.protocol.commit import GroupCoordinatorRequest, OffsetFetchRequest
+from kafka.protocol.commit import OffsetFetchRequest
+from kafka.protocol.find_coordinator import FindCoordinatorRequest
 from kafka.protocol.metadata import MetadataRequest
 from kafka.protocol.types import Array
 from kafka.structs import TopicPartition, OffsetAndMetadata, MemberInformation, GroupInformation
@@ -285,17 +286,14 @@ def _find_coordinator_id_send_request(self, group_id):
         Returns:
             A message future
         """
-        # TODO add support for dynamically picking version of
-        # GroupCoordinatorRequest which was renamed to FindCoordinatorRequest.
-        # When I experimented with this, the coordinator value returned in
-        # GroupCoordinatorResponse_v1 didn't match the value returned by
-        # GroupCoordinatorResponse_v0 and I couldn't figure out why.
-        version = self._client.api_version(GroupCoordinatorRequest, max_version=0)
+        version = self._client.api_version(FindCoordinatorRequest, max_version=2)
         if version <= 0:
-            request = GroupCoordinatorRequest[version](group_id)
+            request = FindCoordinatorRequest[version](group_id)
+        elif version <= 2:
+            request = FindCoordinatorRequest[version](group_id, 0)
         else:
             raise NotImplementedError(
-                "Support for GroupCoordinatorRequest_v{} has not yet been added to KafkaAdminClient."
+                "Support for FindCoordinatorRequest_v{} has not yet been added to KafkaAdminClient."
                 .format(version))
         return self._send_request_to_node(self._client.least_loaded_node(), request)
 
@@ -308,18 +306,13 @@ def _find_coordinator_id_process_response(self, response):
         Returns:
             The node_id of the broker that is the coordinator.
         """
-        if response.API_VERSION <= 0:
-            error_type = Errors.for_code(response.error_code)
-            if error_type is not Errors.NoError:
-                # Note: When error_type.retriable, Java will retry... see
-                # KafkaAdminClient's handleFindCoordinatorError method
-                raise error_type(
-                    "FindCoordinatorRequest failed with response '{}'."
-                    .format(response))
-        else:
-            raise NotImplementedError(
-                "Support for FindCoordinatorRequest_v{} has not yet been added to KafkaAdminClient."
-                .format(response.API_VERSION))
+        error_type = Errors.for_code(response.error_code)
+        if error_type is not Errors.NoError:
+            # Note: When error_type.retriable, Java will retry... see
+            # KafkaAdminClient's handleFindCoordinatorError method
+            raise error_type(
+                "FindCoordinatorRequest failed with response '{}'."
+                .format(response))
         return response.coordinator_id
 
     def _find_coordinator_ids(self, group_ids):

kafka/cluster.py

@@ -21,7 +21,7 @@ class ClusterMetadata(object):
     A class to manage kafka cluster metadata.
 
     This class does not perform any IO. It simply updates internal state
-    given API responses (MetadataResponse, GroupCoordinatorResponse).
+    given API responses (MetadataResponse, FindCoordinatorResponse).
 
     Keyword Arguments:
         retry_backoff_ms (int): Milliseconds to backoff when retrying on
@@ -367,16 +367,16 @@ def add_group_coordinator(self, group, response):
         """Update with metadata for a group coordinator
 
         Arguments:
-            group (str): name of group from GroupCoordinatorRequest
-            response (GroupCoordinatorResponse): broker response
+            group (str): name of group from FindCoordinatorRequest
+            response (FindCoordinatorResponse): broker response
 
         Returns:
             string: coordinator node_id if metadata is updated, None on error
         """
         log.debug("Updating coordinator for %s: %s", group, response)
         error_type = Errors.for_code(response.error_code)
         if error_type is not Errors.NoError:
-            log.error("GroupCoordinatorResponse error: %s", error_type)
+            log.error("FindCoordinatorResponse error: %s", error_type)
             self._groups[group] = -1
             return
 

kafka/conn.py

@@ -1255,14 +1255,15 @@ def reset_override_configs():
         from kafka.protocol.admin import ListGroupsRequest
         from kafka.protocol.api_versions import ApiVersionsRequest
         from kafka.protocol.broker_api_versions import BROKER_API_VERSIONS
-        from kafka.protocol.commit import OffsetFetchRequest, GroupCoordinatorRequest
+        from kafka.protocol.commit import OffsetFetchRequest
+        from kafka.protocol.find_coordinator import FindCoordinatorRequest
 
         test_cases = [
             # All cases starting from 0.10 will be based on ApiVersionsResponse
             ((0, 11), ApiVersionsRequest[1]()),
             ((0, 10, 0), ApiVersionsRequest[0]()),
             ((0, 9), ListGroupsRequest[0]()),
-            ((0, 8, 2), GroupCoordinatorRequest[0]('kafka-python-default-group')),
+            ((0, 8, 2), FindCoordinatorRequest[0]('kafka-python-default-group')),
             ((0, 8, 1), OffsetFetchRequest[0]('kafka-python-default-group', [])),
             ((0, 8, 0), MetadataRequest[0](topics)),
         ]

kafka/coordinator/base.py

@@ -14,7 +14,8 @@
 from kafka.future import Future
 from kafka.metrics import AnonMeasurable
 from kafka.metrics.stats import Avg, Count, Max, Rate
-from kafka.protocol.commit import GroupCoordinatorRequest, OffsetCommitRequest
+from kafka.protocol.commit import OffsetCommitRequest
+from kafka.protocol.find_coordinator import FindCoordinatorRequest
 from kafka.protocol.group import (HeartbeatRequest, JoinGroupRequest,
                             LeaveGroupRequest, SyncGroupRequest)
 
@@ -660,14 +661,21 @@ def _send_group_coordinator_request(self):
 
         log.debug("Sending group coordinator request for group %s to broker %s",
                   self.group_id, node_id)
-        request = GroupCoordinatorRequest[0](self.group_id)
+        version = self._client.api_version(FindCoordinatorRequest, max_version=2)
+        if version == 0:
+            request = FindCoordinatorRequest[version](self.group_id)
+        else:
+            request = FindCoordinatorRequest[version](self.group_id, 0)
         future = Future()
         _f = self._client.send(node_id, request)
         _f.add_callback(self._handle_group_coordinator_response, future)
         _f.add_errback(self._failed_request, node_id, request, future)
         return future
 
     def _handle_group_coordinator_response(self, future, response):
+        if response.API_VERSION >= 1 and response.throttle_time_ms > 0:
+            log.warning("FindCoordinatorRequest throttled by broker (%d ms)", response.throttle_time_ms)
+
         log.debug("Received group coordinator response %s", response)
 
         error_type = Errors.for_code(response.error_code)

kafka/protocol/commit.py

@@ -207,49 +207,3 @@ class OffsetFetchRequest_v3(Request):
     OffsetFetchResponse_v0, OffsetFetchResponse_v1,
     OffsetFetchResponse_v2, OffsetFetchResponse_v3,
 ]
-
-
-class GroupCoordinatorResponse_v0(Response):
-    API_KEY = 10
-    API_VERSION = 0
-    SCHEMA = Schema(
-        ('error_code', Int16),
-        ('coordinator_id', Int32),
-        ('host', String('utf-8')),
-        ('port', Int32)
-    )
-
-
-class GroupCoordinatorResponse_v1(Response):
-    API_KEY = 10
-    API_VERSION = 1
-    SCHEMA = Schema(
-        ('error_code', Int16),
-        ('error_message', String('utf-8')),
-        ('coordinator_id', Int32),
-        ('host', String('utf-8')),
-        ('port', Int32)
-    )
-
-
-class GroupCoordinatorRequest_v0(Request):
-    API_KEY = 10
-    API_VERSION = 0
-    RESPONSE_TYPE = GroupCoordinatorResponse_v0
-    SCHEMA = Schema(
-        ('consumer_group', String('utf-8'))
-    )
-
-
-class GroupCoordinatorRequest_v1(Request):
-    API_KEY = 10
-    API_VERSION = 1
-    RESPONSE_TYPE = GroupCoordinatorResponse_v1
-    SCHEMA = Schema(
-        ('coordinator_key', String('utf-8')),
-        ('coordinator_type', Int8)
-    )
-
-
-GroupCoordinatorRequest = [GroupCoordinatorRequest_v0, GroupCoordinatorRequest_v1]
-GroupCoordinatorResponse = [GroupCoordinatorResponse_v0, GroupCoordinatorResponse_v1]

kafka/protocol/find_coordinator.py

@@ -0,0 +1,64 @@
+from __future__ import absolute_import
+
+from kafka.protocol.api import Request, Response
+from kafka.protocol.types import Array, Int8, Int16, Int32, Int64, Schema, String
+
+
+class FindCoordinatorResponse_v0(Response):
+    API_KEY = 10
+    API_VERSION = 0
+    SCHEMA = Schema(
+        ('error_code', Int16),
+        ('coordinator_id', Int32),
+        ('host', String('utf-8')),
+        ('port', Int32)
+    )
+
+
+class FindCoordinatorResponse_v1(Response):
+    API_KEY = 10
+    API_VERSION = 1
+    SCHEMA = Schema(
+        ('throttle_time_ms', Int32),
+        ('error_code', Int16),
+        ('error_message', String('utf-8')),
+        ('coordinator_id', Int32),
+        ('host', String('utf-8')),
+        ('port', Int32)
+    )
+
+
+class FindCoordinatorResponse_v2(Response):
+    API_KEY = 10
+    API_VERSION = 2
+    SCHEMA = FindCoordinatorResponse_v1.SCHEMA
+
+
+class FindCoordinatorRequest_v0(Request):
+    API_KEY = 10
+    API_VERSION = 0
+    RESPONSE_TYPE = FindCoordinatorResponse_v0
+    SCHEMA = Schema(
+        ('consumer_group', String('utf-8'))
+    )
+
+
+class FindCoordinatorRequest_v1(Request):
+    API_KEY = 10
+    API_VERSION = 1
+    RESPONSE_TYPE = FindCoordinatorResponse_v1
+    SCHEMA = Schema(
+        ('coordinator_key', String('utf-8')),
+        ('coordinator_type', Int8) # 0: consumer, 1: transaction
+    )
+
+
+class FindCoordinatorRequest_v2(Request):
+    API_KEY = 10
+    API_VERSION = 2
+    RESPONSE_TYPE = FindCoordinatorResponse_v2
+    SCHEMA = FindCoordinatorRequest_v1.SCHEMA
+
+
+FindCoordinatorRequest = [FindCoordinatorRequest_v0, FindCoordinatorRequest_v1, FindCoordinatorRequest_v2]
+FindCoordinatorResponse = [FindCoordinatorResponse_v0, FindCoordinatorResponse_v1, FindCoordinatorResponse_v2]

kafka/protocol/parser.py

@@ -4,7 +4,7 @@
 import logging
 
 import kafka.errors as Errors
-from kafka.protocol.commit import GroupCoordinatorResponse
+from kafka.protocol.find_coordinator import FindCoordinatorResponse
 from kafka.protocol.frame import KafkaBytes
 from kafka.protocol.types import Int32, TaggedFields
 from kafka.version import __version__
@@ -142,7 +142,7 @@ def _process_response(self, read_buffer):
         # 0.8.2 quirk
         if (recv_correlation_id == 0 and
             correlation_id != 0 and
-            request.RESPONSE_TYPE is GroupCoordinatorResponse[0] and
+            request.RESPONSE_TYPE is FindCoordinatorResponse[0] and
             (self._api_version == (0, 8, 2) or self._api_version is None)):
             log.warning('Kafka 0.8.2 quirk -- GroupCoordinatorResponse'
                         ' Correlation ID does not match request. This'

test/test_protocol.py

@@ -5,8 +5,8 @@
 import pytest
 
 from kafka.protocol.api import RequestHeader
-from kafka.protocol.commit import GroupCoordinatorRequest
 from kafka.protocol.fetch import FetchRequest, FetchResponse
+from kafka.protocol.find_coordinator import FindCoordinatorRequest
 from kafka.protocol.message import Message, MessageSet, PartialMessage
 from kafka.protocol.metadata import MetadataRequest
 from kafka.protocol.types import Int16, Int32, Int64, String, UnsignedVarInt32, CompactString, CompactArray, CompactBytes
@@ -168,7 +168,7 @@ def test_encode_message_header():
         b'client3',                        # ClientId
     ])
 
-    req = GroupCoordinatorRequest[0]('foo')
+    req = FindCoordinatorRequest[0]('foo')
     header = RequestHeader(req, correlation_id=4, client_id='client3')
     assert header.encode() == expect