KAFKA-4547: Avoid resetting paused partitions to committed offsets

dpkp · dpkp · commit 2f262b9f993d · 2025-04-19T09:47:34.000-07:00
diff --git a/kafka/consumer/fetcher.py b/kafka/consumer/fetcher.py
@@ -194,13 +194,7 @@ def update_fetch_positions(self, partitions, timeout_ms=None):
         inner_timeout_ms = timeout_ms_fn(timeout_ms, 'Timeout updating fetch positions')
         # reset the fetch position to the committed position
         for tp in partitions:
-            if not self._subscriptions.is_assigned(tp):
-                log.warning("partition %s is not assigned - skipping offset"
-                            " update", tp)
-                continue
-            elif self._subscriptions.is_fetchable(tp):
-                log.warning("partition %s is still fetchable -- skipping offset"
-                            " update", tp)
+            if not self._subscriptions.is_assigned(tp) or self._subscriptions.has_valid_position(tp):
                 continue
 
             if self._subscriptions.is_offset_reset_needed(tp):
diff --git a/kafka/consumer/group.py b/kafka/consumer/group.py
@@ -1144,9 +1144,9 @@ def _update_fetch_positions(self, partitions, timeout_ms=None):
             # their own offsets).
             self._fetcher.reset_offsets_if_needed(partitions, timeout_ms=inner_timeout_ms())
 
-            if not self._subscription.has_all_fetch_positions():
-                # if we still don't have offsets for all partitions, then we should either seek
-                # to the last committed position or reset using the auto reset policy
+            if not self._subscription.has_all_fetch_positions(partitions):
+                # if we still don't have offsets for the given partitions, then we should either
+                # seek to the last committed position or reset using the auto reset policy
                 if (self.config['api_version'] >= (0, 8, 1) and
                     self.config['group_id'] is not None):
                     # first refresh commits for all assigned partitions
diff --git a/kafka/consumer/subscription_state.py b/kafka/consumer/subscription_state.py
@@ -351,9 +351,11 @@ def has_default_offset_reset_policy(self):
     def is_offset_reset_needed(self, partition):
         return self.assignment[partition].awaiting_reset
 
-    def has_all_fetch_positions(self):
-        for state in self.assignment.values():
-            if not state.has_valid_position:
+    def has_all_fetch_positions(self, partitions=None):
+        if partitions is None:
+            partitions = self.assigned_partitions()
+        for tp in partitions:
+            if not self.has_valid_position(tp):
                 return False
         return True
 
@@ -364,6 +366,9 @@ def missing_fetch_positions(self):
                 missing.add(partition)
         return missing
 
+    def has_valid_position(self, partition):
+        return partition in self.assignment and self.assignment[partition].has_valid_position
+
     def is_assigned(self, partition):
         return partition in self.assignment
 
@@ -387,6 +392,9 @@ def move_partition_to_end(self, partition):
                 state = self.assignment.pop(partition)
                 self.assignment[partition] = state
 
+    def position(self, partition):
+        return self.assignment[partition].position
+
 
 class TopicPartitionState(object):
     def __init__(self):
diff --git a/test/test_fetcher.py b/test/test_fetcher.py
@@ -17,7 +17,7 @@
 from kafka.future import Future
 from kafka.protocol.broker_api_versions import BROKER_API_VERSIONS
 from kafka.protocol.fetch import FetchRequest, FetchResponse
-from kafka.protocol.list_offsets import ListOffsetsResponse
+from kafka.protocol.list_offsets import ListOffsetsResponse, OffsetResetStrategy
 from kafka.errors import (
     StaleMetadata, LeaderNotAvailableError, NotLeaderForPartitionError,
     UnknownTopicOrPartitionError, OffsetOutOfRangeError
@@ -610,3 +610,51 @@ def test_partition_records_compacted_offset(mocker):
     msgs = records.take()
     assert len(msgs) == batch_end - fetch_offset - 1
     assert msgs[0].offset == fetch_offset + 1
+
+
+def test_update_fetch_positions_paused(subscription_state, client, mocker):
+    fetcher = Fetcher(client, subscription_state)
+    tp = TopicPartition('foo', 0)
+    subscription_state.assign_from_user([tp])
+    subscription_state.pause(tp) # paused partition does not have a valid position
+    subscription_state.need_offset_reset(tp, OffsetResetStrategy.LATEST)
+
+    mocker.patch.object(fetcher, '_retrieve_offsets', return_value={tp: OffsetAndTimestamp(10, 1, -1)})
+    fetcher.update_fetch_positions([tp])
+
+    assert not subscription_state.is_offset_reset_needed(tp)
+    assert not subscription_state.is_fetchable(tp) # because tp is paused
+    assert subscription_state.has_valid_position(tp)
+    assert subscription_state.position(tp) == OffsetAndMetadata(10, '', -1)
+
+
+def test_update_fetch_positions_paused_without_valid(subscription_state, client, mocker):
+    fetcher = Fetcher(client, subscription_state)
+    tp = TopicPartition('foo', 0)
+    subscription_state.assign_from_user([tp])
+    subscription_state.pause(tp) # paused partition does not have a valid position
+
+    mocker.patch.object(fetcher, '_retrieve_offsets', return_value={tp: OffsetAndTimestamp(0, 1, -1)})
+    fetcher.update_fetch_positions([tp])
+
+    assert not subscription_state.is_offset_reset_needed(tp)
+    assert not subscription_state.is_fetchable(tp) # because tp is paused
+    assert subscription_state.has_valid_position(tp)
+    assert subscription_state.position(tp) == OffsetAndMetadata(0, '', -1)
+
+
+def test_update_fetch_positions_paused_with_valid(subscription_state, client, mocker):
+    fetcher = Fetcher(client, subscription_state)
+    tp = TopicPartition('foo', 0)
+    subscription_state.assign_from_user([tp])
+    subscription_state.assignment[tp].committed = OffsetAndMetadata(0, '', -1)
+    subscription_state.seek(tp, 10)
+    subscription_state.pause(tp) # paused partition already has a valid position
+
+    mocker.patch.object(fetcher, '_retrieve_offsets', return_value={tp: OffsetAndTimestamp(0, 1, -1)})
+    fetcher.update_fetch_positions([tp])
+
+    assert not subscription_state.is_offset_reset_needed(tp)
+    assert not subscription_state.is_fetchable(tp) # because tp is paused
+    assert subscription_state.has_valid_position(tp)
+    assert subscription_state.position(tp) == OffsetAndMetadata(10, '', -1)
