Support ListOffsets v3 in consumer (#2501)

Author: dpkp

kafka/conn.py

@@ -26,7 +26,7 @@
 from kafka.oauth.abstract import AbstractTokenProvider
 from kafka.protocol.admin import SaslHandShakeRequest, DescribeAclsRequest, DescribeClientQuotasRequest
 from kafka.protocol.commit import OffsetFetchRequest
-from kafka.protocol.offset import OffsetRequest
+from kafka.protocol.list_offsets import ListOffsetsRequest
 from kafka.protocol.produce import ProduceRequest
 from kafka.protocol.metadata import MetadataRequest
 from kafka.protocol.fetch import FetchRequest
@@ -1202,7 +1202,7 @@ def _infer_broker_version_from_api_versions(self, api_versions):
             ((2, 5), DescribeAclsRequest[2]),
             ((2, 4), ProduceRequest[8]),
             ((2, 3), FetchRequest[11]),
-            ((2, 2), OffsetRequest[5]),
+            ((2, 2), ListOffsetsRequest[5]),
             ((2, 1), FetchRequest[10]),
             ((2, 0), FetchRequest[8]),
             ((1, 1), FetchRequest[7]),

kafka/consumer/fetcher.py

@@ -13,8 +13,8 @@
 from kafka.future import Future
 from kafka.metrics.stats import Avg, Count, Max, Rate
 from kafka.protocol.fetch import FetchRequest
-from kafka.protocol.offset import (
-    OffsetRequest, OffsetResetStrategy, UNKNOWN_OFFSET
+from kafka.protocol.list_offsets import (
+    ListOffsetsRequest, OffsetResetStrategy, UNKNOWN_OFFSET
 )
 from kafka.record import MemoryRecords
 from kafka.serializer import Deserializer
@@ -272,7 +272,7 @@ def _retrieve_offsets(self, timestamps, timeout_ms=float("inf")):
             if not timestamps:
                 return {}
 
-            future = self._send_offset_requests(timestamps)
+            future = self._send_list_offsets_requests(timestamps)
             self._client.poll(future=future, timeout_ms=remaining_ms)
 
             if future.succeeded():
@@ -519,7 +519,7 @@ def _deserialize(self, f, topic, bytes_):
             return f.deserialize(topic, bytes_)
         return f(bytes_)
 
-    def _send_offset_requests(self, timestamps):
+    def _send_list_offsets_requests(self, timestamps):
         """Fetch offsets for each partition in timestamps dict. This may send
         request to multiple nodes, based on who is Leader for partition.
 
@@ -564,13 +564,13 @@ def on_fail(err):
                 list_offsets_future.failure(err)
 
         for node_id, timestamps in six.iteritems(timestamps_by_node):
-            _f = self._send_offset_request(node_id, timestamps)
+            _f = self._send_list_offsets_request(node_id, timestamps)
             _f.add_callback(on_success)
             _f.add_errback(on_fail)
         return list_offsets_future
 
-    def _send_offset_request(self, node_id, timestamps):
-        version = self._client.api_version(OffsetRequest, max_version=1)
+    def _send_list_offsets_request(self, node_id, timestamps):
+        version = self._client.api_version(ListOffsetsRequest, max_version=3)
         by_topic = collections.defaultdict(list)
         for tp, timestamp in six.iteritems(timestamps):
             if version >= 1:
@@ -579,28 +579,39 @@ def _send_offset_request(self, node_id, timestamps):
                 data = (tp.partition, timestamp, 1)
             by_topic[tp.topic].append(data)
 
-        request = OffsetRequest[version](-1, list(six.iteritems(by_topic)))
+        if version <= 1:
+            request = ListOffsetsRequest[version](
+                    -1,
+                    list(six.iteritems(by_topic)))
+        else:
+            request = ListOffsetsRequest[version](
+                    -1,
+                    self._isolation_level,
+                    list(six.iteritems(by_topic)))
+
 
         # Client returns a future that only fails on network issues
         # so create a separate future and attach a callback to update it
         # based on response error codes
         future = Future()
 
         _f = self._client.send(node_id, request)
-        _f.add_callback(self._handle_offset_response, future)
+        _f.add_callback(self._handle_list_offsets_response, future)
         _f.add_errback(lambda e: future.failure(e))
         return future
 
-    def _handle_offset_response(self, future, response):
-        """Callback for the response of the list offset call above.
+    def _handle_list_offsets_response(self, future, response):
+        """Callback for the response of the ListOffsets api call
 
         Arguments:
             future (Future): the future to update based on response
-            response (OffsetResponse): response from the server
+            response (ListOffsetsResponse): response from the server
 
         Raises:
             AssertionError: if response does not match partition
         """
+        if response.API_VERSION >= 2 and response.throttle_time_ms > 0:
+            log.warning("ListOffsetsRequest throttled by broker (%d ms)", response.throttle_time_ms)
         timestamp_offset_map = {}
         for topic, part_data in response.topics:
             for partition_info in part_data:
@@ -610,18 +621,18 @@ def _handle_offset_response(self, future, response):
                 if error_type is Errors.NoError:
                     if response.API_VERSION == 0:
                         offsets = partition_info[2]
-                        assert len(offsets) <= 1, 'Expected OffsetResponse with one offset'
+                        assert len(offsets) <= 1, 'Expected ListOffsetsResponse with one offset'
                         if not offsets:
                             offset = UNKNOWN_OFFSET
                         else:
                             offset = offsets[0]
-                        log.debug("Handling v0 ListOffsetResponse response for %s. "
+                        log.debug("Handling v0 ListOffsetsResponse response for %s. "
                                   "Fetched offset %s", partition, offset)
                         if offset != UNKNOWN_OFFSET:
                             timestamp_offset_map[partition] = (offset, None)
                     else:
                         timestamp, offset = partition_info[2:]
-                        log.debug("Handling ListOffsetResponse response for %s. "
+                        log.debug("Handling ListOffsetsResponse response for %s. "
                                   "Fetched offset %s, timestamp %s",
                                   partition, offset, timestamp)
                         if offset != UNKNOWN_OFFSET:
@@ -638,7 +649,7 @@ def _handle_offset_response(self, future, response):
                     future.failure(error_type(partition))
                     return
                 elif error_type is Errors.UnknownTopicOrPartitionError:
-                    log.warning("Received unknown topic or partition error in ListOffset "
+                    log.warning("Received unknown topic or partition error in ListOffsets "
                              "request for partition %s. The topic/partition " +
                              "may not exist or the user may not have Describe access "
                              "to it.", partition)

kafka/consumer/group.py

@@ -16,7 +16,7 @@
 from kafka.coordinator.assignors.range import RangePartitionAssignor
 from kafka.coordinator.assignors.roundrobin import RoundRobinPartitionAssignor
 from kafka.metrics import MetricConfig, Metrics
-from kafka.protocol.offset import OffsetResetStrategy
+from kafka.protocol.list_offsets import OffsetResetStrategy
 from kafka.structs import TopicPartition
 from kafka.version import __version__
 

kafka/consumer/subscription_state.py

@@ -7,7 +7,7 @@
 from kafka.vendor import six
 
 from kafka.errors import IllegalStateError
-from kafka.protocol.offset import OffsetResetStrategy
+from kafka.protocol.list_offsets import OffsetResetStrategy
 from kafka.structs import OffsetAndMetadata
 
 log = logging.getLogger(__name__)

kafka/protocol/offset.py renamed to kafka/protocol/list_offsets.py

@@ -12,7 +12,7 @@ class OffsetResetStrategy(object):
     NONE = 0
 
 
-class OffsetResponse_v0(Response):
+class ListOffsetsResponse_v0(Response):
     API_KEY = 2
     API_VERSION = 0
     SCHEMA = Schema(
@@ -24,7 +24,7 @@ class OffsetResponse_v0(Response):
                 ('offsets', Array(Int64))))))
     )
 
-class OffsetResponse_v1(Response):
+class ListOffsetsResponse_v1(Response):
     API_KEY = 2
     API_VERSION = 1
     SCHEMA = Schema(
@@ -38,7 +38,7 @@ class OffsetResponse_v1(Response):
     )
 
 
-class OffsetResponse_v2(Response):
+class ListOffsetsResponse_v2(Response):
     API_KEY = 2
     API_VERSION = 2
     SCHEMA = Schema(
@@ -53,16 +53,16 @@ class OffsetResponse_v2(Response):
     )
 
 
-class OffsetResponse_v3(Response):
+class ListOffsetsResponse_v3(Response):
     """
     on quota violation, brokers send out responses before throttling
     """
     API_KEY = 2
     API_VERSION = 3
-    SCHEMA = OffsetResponse_v2.SCHEMA
+    SCHEMA = ListOffsetsResponse_v2.SCHEMA
 
 
-class OffsetResponse_v4(Response):
+class ListOffsetsResponse_v4(Response):
     """
     Add leader_epoch to response
     """
@@ -81,19 +81,19 @@ class OffsetResponse_v4(Response):
     )
 
 
-class OffsetResponse_v5(Response):
+class ListOffsetsResponse_v5(Response):
     """
     adds a new error code, OFFSET_NOT_AVAILABLE
     """
     API_KEY = 2
     API_VERSION = 5
-    SCHEMA = OffsetResponse_v4.SCHEMA
+    SCHEMA = ListOffsetsResponse_v4.SCHEMA
 
 
-class OffsetRequest_v0(Request):
+class ListOffsetsRequest_v0(Request):
     API_KEY = 2
     API_VERSION = 0
-    RESPONSE_TYPE = OffsetResponse_v0
+    RESPONSE_TYPE = ListOffsetsResponse_v0
     SCHEMA = Schema(
         ('replica_id', Int32),
         ('topics', Array(
@@ -107,10 +107,10 @@ class OffsetRequest_v0(Request):
         'replica_id': -1
     }
 
-class OffsetRequest_v1(Request):
+class ListOffsetsRequest_v1(Request):
     API_KEY = 2
     API_VERSION = 1
-    RESPONSE_TYPE = OffsetResponse_v1
+    RESPONSE_TYPE = ListOffsetsResponse_v1
     SCHEMA = Schema(
         ('replica_id', Int32),
         ('topics', Array(
@@ -124,10 +124,10 @@ class OffsetRequest_v1(Request):
     }
 
 
-class OffsetRequest_v2(Request):
+class ListOffsetsRequest_v2(Request):
     API_KEY = 2
     API_VERSION = 2
-    RESPONSE_TYPE = OffsetResponse_v2
+    RESPONSE_TYPE = ListOffsetsResponse_v2
     SCHEMA = Schema(
         ('replica_id', Int32),
         ('isolation_level', Int8),  # <- added isolation_level
@@ -142,23 +142,23 @@ class OffsetRequest_v2(Request):
     }
 
 
-class OffsetRequest_v3(Request):
+class ListOffsetsRequest_v3(Request):
     API_KEY = 2
     API_VERSION = 3
-    RESPONSE_TYPE = OffsetResponse_v3
-    SCHEMA = OffsetRequest_v2.SCHEMA
+    RESPONSE_TYPE = ListOffsetsResponse_v3
+    SCHEMA = ListOffsetsRequest_v2.SCHEMA
     DEFAULTS = {
         'replica_id': -1
     }
 
 
-class OffsetRequest_v4(Request):
+class ListOffsetsRequest_v4(Request):
     """
     Add current_leader_epoch to request
     """
     API_KEY = 2
     API_VERSION = 4
-    RESPONSE_TYPE = OffsetResponse_v4
+    RESPONSE_TYPE = ListOffsetsResponse_v4
     SCHEMA = Schema(
         ('replica_id', Int32),
         ('isolation_level', Int8),  # <- added isolation_level
@@ -174,21 +174,21 @@ class OffsetRequest_v4(Request):
     }
 
 
-class OffsetRequest_v5(Request):
+class ListOffsetsRequest_v5(Request):
     API_KEY = 2
     API_VERSION = 5
-    RESPONSE_TYPE = OffsetResponse_v5
-    SCHEMA = OffsetRequest_v4.SCHEMA
+    RESPONSE_TYPE = ListOffsetsResponse_v5
+    SCHEMA = ListOffsetsRequest_v4.SCHEMA
     DEFAULTS = {
         'replica_id': -1
     }
 
 
-OffsetRequest = [
-    OffsetRequest_v0, OffsetRequest_v1, OffsetRequest_v2,
-    OffsetRequest_v3, OffsetRequest_v4, OffsetRequest_v5,
+ListOffsetsRequest = [
+    ListOffsetsRequest_v0, ListOffsetsRequest_v1, ListOffsetsRequest_v2,
+    ListOffsetsRequest_v3, ListOffsetsRequest_v4, ListOffsetsRequest_v5,
 ]
-OffsetResponse = [
-    OffsetResponse_v0, OffsetResponse_v1, OffsetResponse_v2,
-    OffsetResponse_v3, OffsetResponse_v4, OffsetResponse_v5,
+ListOffsetsResponse = [
+    ListOffsetsResponse_v0, ListOffsetsResponse_v1, ListOffsetsResponse_v2,
+    ListOffsetsResponse_v3, ListOffsetsResponse_v4, ListOffsetsResponse_v5,
 ]