diff --git a/kafka/consumer/subscription_state.py b/kafka/consumer/subscription_state.py index d8283e349..74e4a24f1 100644 --- a/kafka/consumer/subscription_state.py +++ b/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 diff --git a/test/consumer/test_subscription_state.py b/test/consumer/test_subscription_state.py index 1855cfb9e..7b5c8f1d5 100644 --- a/test/consumer/test_subscription_state.py +++ b/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)])