Fix stuck partitions retained after rebalance when using listener'

kafka/consumer/subscription_state.py

@@ -290,6 +290,9 @@ def _apply_assignment(self, partitions, randomize=False):
         for tp in list(self.assignment.keys()):
             if tp not in new_set:
                 del self.assignment[tp]
+            else:
+                # Retained partitions are no longer pending_revocation
+                self.assignment[tp].mark_pending_revocation(False)
         # Add new partitions; kept partitions retain their existing
         # TopicPartitionState (positions, paused flag, KIP-392 cache,
         # etc.).
@@ -651,12 +654,12 @@ def pause(self):
     def resume(self):
         self.paused = False
 
-    def mark_pending_revocation(self):
+    def mark_pending_revocation(self, val=True):
         """KIP-429: gate fetches while an on_partitions_revoked /
         on_partitions_lost listener is in progress for this partition.
         Single-shot: the surrounding ``assign_from_subscribed`` drops
         the state object once the listener returns."""
-        self._pending_revocation = True
+        self._pending_revocation = val
 
     def is_fetchable(self):
         return (not self.paused

test/consumer/test_subscription_state.py

@@ -141,6 +141,49 @@ def test_mark_pending_revocation_ignores_unknown_partition(self):
         # Still gates the assigned one.
         assert s.assignment[TopicPartition('foo', 0)]._pending_revocation is True
 
+    def test_pending_revocation_cleared_for_partition_retained_across_rebalance(self):
+        """Regression test for #3131.
+
+        Under the EAGER rebalance protocol the coordinator marks *every*
+        currently-owned partition pending-revocation before JoinGroup (see
+        ConsumerCoordinator._on_join_prepare_async). When the rebalance
+        completes, ``assign_from_subscribed`` preserves the
+        ``TopicPartitionState`` for any partition the member is given back,
+        so a *retained* partition carries a stale ``_pending_revocation``
+        flag. Because the flag is only ever cleared in
+        ``TopicPartitionState.__init__``, such a partition reports
+        not-fetchable forever - the fetcher positions it (ListOffsets) but
+        never issues a FetchRequest - until a later rebalance drops it and
+        creates fresh state.
+
+        A partition retained across a rebalance must be fetchable again.
+        """
+        s = self._subscribed()
+        foo0 = TopicPartition('foo', 0)
+        foo1 = TopicPartition('foo', 1)
+        s.assign_from_subscribed([foo0, foo1])
+        for tp in (foo0, foo1):
+            s.assignment[tp].seek(
+                OffsetAndMetadata(offset=0, metadata='', leader_epoch=-1))
+        assert s.is_fetchable(foo0)
+        assert s.is_fetchable(foo1)
+
+        # EAGER pre-rebalance: mark the full current assignment for revocation
+        # (mirrors mark_pending_revocation(assigned_partitions())).
+        s.mark_pending_revocation({foo0, foo1})
+        assert not s.is_fetchable(foo0)
+        assert not s.is_fetchable(foo1)
+
+        # Rebalance completes and the member is assigned the same partitions
+        # back. Their state (and valid position) is preserved.
+        s.assign_from_subscribed([foo0, foo1])
+        assert s.assignment[foo0].has_valid_position
+        assert s.assignment[foo1].has_valid_position
+
+        # They must be fetchable again - the revocation window is over.
+        assert s.is_fetchable(foo0)
+        assert s.is_fetchable(foo1)
+
     def test_preserves_preferred_read_replica_on_kept_partition(self):
         s = self._subscribed()
         s.assign_from_subscribed([TopicPartition('foo', 0)])