KAFKA-4937: Batch offset fetches in the Consumer

Author: dpkp

kafka/consumer/fetcher.py

@@ -161,11 +161,13 @@ def reset_offsets_if_needed(self, partitions, timeout_ms=None):
         Raises:
             KafkaTimeoutError if timeout_ms provided
         """
-        inner_timeout_ms = timeout_ms_fn(timeout_ms, 'Timeout resetting offsets')
+        needs_offset_reset = set()
         for tp in partitions:
-            # TODO: If there are several offsets to reset, we could submit offset requests in parallel
             if self._subscriptions.is_assigned(tp) and self._subscriptions.is_offset_reset_needed(tp):
-                self._reset_offset(tp, timeout_ms=inner_timeout_ms())
+                needs_offset_reset.add(tp)
+
+        if needs_offset_reset:
+            self._reset_offsets(needs_offset_reset, timeout_ms=timeout_ms)
 
     def _clean_done_fetch_futures(self):
         while True:
@@ -191,25 +193,28 @@ def update_fetch_positions(self, partitions, timeout_ms=None):
                 partition and no reset policy is available
             KafkaTimeoutError if timeout_ms provided.
         """
-        inner_timeout_ms = timeout_ms_fn(timeout_ms, 'Timeout updating fetch positions')
+        needs_offset_reset = set()
         # reset the fetch position to the committed position
         for tp in partitions:
             if not self._subscriptions.is_assigned(tp) or self._subscriptions.has_valid_position(tp):
                 continue
 
             if self._subscriptions.is_offset_reset_needed(tp):
-                self._reset_offset(tp, timeout_ms=inner_timeout_ms())
+                needs_offset_reset.add(tp)
             elif self._subscriptions.assignment[tp].committed is None:
                 # there's no committed position, so we need to reset with the
                 # default strategy
                 self._subscriptions.need_offset_reset(tp)
-                self._reset_offset(tp, timeout_ms=inner_timeout_ms())
+                needs_offset_reset.add(tp)
             else:
                 committed = self._subscriptions.assignment[tp].committed.offset
                 log.debug("Resetting offset for partition %s to the committed"
                           " offset %s", tp, committed)
                 self._subscriptions.seek(tp, committed)
 
+        if needs_offset_reset:
+            self._reset_offsets(needs_offset_reset, timeout_ms=timeout_ms)
+
     def get_offsets_by_times(self, timestamps, timeout_ms):
         offsets = self._retrieve_offsets(timestamps, timeout_ms)
         for tp in timestamps:
@@ -232,37 +237,36 @@ def beginning_or_end_offset(self, partitions, timestamp, timeout_ms):
             offsets[tp] = offsets[tp].offset
         return offsets
 
-    def _reset_offset(self, partition, timeout_ms=None):
-        """Reset offsets for the given partition using the offset reset strategy.
+    def _reset_offsets(self, partitions, timeout_ms=None):
+        """Reset offsets for the given partitions using the offset reset strategy.
 
         Arguments:
-            partition (TopicPartition): the partition that needs reset offset
+            partitions ([TopicPartition]): the partitions that need offsets reset
 
         Raises:
             NoOffsetForPartitionError: if no offset reset strategy is defined
             KafkaTimeoutError if timeout_ms provided
         """
-        timestamp = self._subscriptions.assignment[partition].reset_strategy
-        if timestamp is OffsetResetStrategy.EARLIEST:
-            strategy = 'earliest'
-        elif timestamp is OffsetResetStrategy.LATEST:
-            strategy = 'latest'
-        else:
-            raise NoOffsetForPartitionError(partition)
+        offset_resets = dict()
+        for tp in partitions:
+            ts = self._subscriptions.assignment[tp].reset_strategy
+            if not ts:
+                raise NoOffsetForPartitionError(tp)
+            offset_resets[tp] = ts
 
-        log.debug("Resetting offset for partition %s to offset %s.",
-                  partition, strategy)
-        offsets = self._retrieve_offsets({partition: timestamp}, timeout_ms=timeout_ms)
+        offsets = self._retrieve_offsets(offset_resets, timeout_ms=timeout_ms)
 
-        if partition in offsets:
-            offset = offsets[partition].offset
+        for partition in partitions:
+            if partition not in offsets:
+                raise NoOffsetForPartitionError(partition)
 
             # we might lose the assignment while fetching the offset,
             # so check it is still active
             if self._subscriptions.is_assigned(partition):
+                offset = offsets[partition].offset
+                log.debug("Resetting offset for partition %s to offset %s.",
+                          partition, offset)
                 self._subscriptions.seek(partition, offset)
-        else:
-            log.debug("Could not find offset for partition %s since it is probably deleted" % (partition,))
 
     def _retrieve_offsets(self, timestamps, timeout_ms=None):
         """Fetch offset for each partition passed in ``timestamps`` map.

kafka/consumer/group.py

@@ -760,7 +760,8 @@ def position(self, partition, timeout_ms=None):
         assert self._subscription.is_assigned(partition), 'Partition is not assigned'
         position = self._subscription.assignment[partition].position
         if position is None:
-            self._update_fetch_positions([partition], timeout_ms=timeout_ms)
+            # batch update fetch positions for any partitions without a valid position
+            self._update_fetch_positions(self._subscription.assigned_partitions(), timeout_ms=timeout_ms)
             position = self._subscription.assignment[partition].position
         return position.offset if position else None
 

kafka/consumer/subscription_state.py

@@ -402,7 +402,7 @@ def __init__(self):
         self.has_valid_position = False # whether we have valid position
         self.paused = False # whether this partition has been paused by the user
         self.awaiting_reset = False # whether we are awaiting reset
-        self.reset_strategy = None # the reset strategy if awaitingReset is set
+        self.reset_strategy = None # the reset strategy if awaiting_reset is set
         self._position = None # OffsetAndMetadata exposed to the user
         self.highwater = None
         self.drop_pending_record_batch = False

test/test_fetcher.py

@@ -109,46 +109,46 @@ def test_create_fetch_requests(fetcher, mocker, api_version, fetch_version):
 
 
 def test_update_fetch_positions(fetcher, topic, mocker):
-    mocker.patch.object(fetcher, '_reset_offset')
+    mocker.patch.object(fetcher, '_reset_offsets')
     partition = TopicPartition(topic, 0)
 
     # unassigned partition
     fetcher.update_fetch_positions([TopicPartition('fizzbuzz', 0)])
-    assert fetcher._reset_offset.call_count == 0
+    assert fetcher._reset_offsets.call_count == 0
 
     # fetchable partition (has offset, not paused)
     fetcher.update_fetch_positions([partition])
-    assert fetcher._reset_offset.call_count == 0
+    assert fetcher._reset_offsets.call_count == 0
 
     # partition needs reset, no committed offset
     fetcher._subscriptions.need_offset_reset(partition)
     fetcher._subscriptions.assignment[partition].awaiting_reset = False
     fetcher.update_fetch_positions([partition])
-    fetcher._reset_offset.assert_called_with(partition, timeout_ms=None)
+    fetcher._reset_offsets.assert_called_with(set([partition]), timeout_ms=None)
     assert fetcher._subscriptions.assignment[partition].awaiting_reset is True
     fetcher.update_fetch_positions([partition])
-    fetcher._reset_offset.assert_called_with(partition, timeout_ms=None)
+    fetcher._reset_offsets.assert_called_with(set([partition]), timeout_ms=None)
 
     # partition needs reset, has committed offset
-    fetcher._reset_offset.reset_mock()
+    fetcher._reset_offsets.reset_mock()
     fetcher._subscriptions.need_offset_reset(partition)
     fetcher._subscriptions.assignment[partition].awaiting_reset = False
     fetcher._subscriptions.assignment[partition].committed = OffsetAndMetadata(123, '', -1)
     mocker.patch.object(fetcher._subscriptions, 'seek')
     fetcher.update_fetch_positions([partition])
-    assert fetcher._reset_offset.call_count == 0
+    assert fetcher._reset_offsets.call_count == 0
     fetcher._subscriptions.seek.assert_called_with(partition, 123)
 
 
-def test__reset_offset(fetcher, mocker):
+def test__reset_offsets(fetcher, mocker):
     tp = TopicPartition("topic", 0)
     fetcher._subscriptions.subscribe(topics=["topic"])
     fetcher._subscriptions.assign_from_subscribed([tp])
     fetcher._subscriptions.need_offset_reset(tp)
     mocked = mocker.patch.object(fetcher, '_retrieve_offsets')
 
     mocked.return_value = {tp: OffsetAndTimestamp(1001, None, -1)}
-    fetcher._reset_offset(tp)
+    fetcher._reset_offsets([tp])
     assert not fetcher._subscriptions.assignment[tp].awaiting_reset
     assert fetcher._subscriptions.assignment[tp].position.offset == 1001