From 1404751b99596467bb8a0fb6b209fd446f612036 Mon Sep 17 00:00:00 2001 From: Marcus Stojcevich <871961+mstojcevich@users.noreply.github.com> Date: Fri, 9 Oct 2026 12:15:48 -0400 Subject: [PATCH] Fix StickyPartitionAssignor infinite loop with mixed subscriptions Suppose a partition already moved from X to C during this rebalance and is then reassigned from C to N. PartitionMovements.get_partition_to_be_moved may return a same-topic substitute that moved from N to X. X owns that substitute, so moving only the substitute leaves C's load unchanged, and _perform_reassignments never terminates. Also move the original partition back to X. Both earlier movements are undone, and C gives up one partition to N. Co-Authored-By: Claude Opus 5.5 --- .../assignors/sticky/sticky_assignor.py | 5 ++ test/consumer/test_assignors.py | 47 ++++++++++++++++++- 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/kafka/coordinator/assignors/sticky/sticky_assignor.py b/kafka/coordinator/assignors/sticky/sticky_assignor.py index 9e1407046..67e333ac6 100644 --- a/kafka/coordinator/assignors/sticky/sticky_assignor.py +++ b/kafka/coordinator/assignors/sticky/sticky_assignor.py @@ -450,6 +450,11 @@ def _reassign_partition_to_consumer(self, partition, new_consumer): consumer = self.current_partition_consumer[partition] # find the correct partition movement considering the stickiness requirement partition_to_be_moved = self.partition_movements.get_partition_to_be_moved(partition, consumer, new_consumer) + substitute_owner = self.current_partition_consumer[partition_to_be_moved] + if substitute_owner != consumer: + # partition and the substitute both moved earlier in this rebalance; + # undo both moves, otherwise consumer stays overloaded forever + self._move_partition(partition, substitute_owner) self._move_partition(partition_to_be_moved, new_consumer) def _move_partition(self, partition, new_consumer): diff --git a/test/consumer/test_assignors.py b/test/consumer/test_assignors.py index 04a07c928..210ad93d0 100644 --- a/test/consumer/test_assignors.py +++ b/test/consumer/test_assignors.py @@ -9,7 +9,7 @@ from kafka.coordinator.assignors.abstract import ConsumerProtocolAssignment from kafka.coordinator.assignors.range import RangePartitionAssignor from kafka.coordinator.assignors.roundrobin import RoundRobinPartitionAssignor -from kafka.coordinator.assignors.sticky.sticky_assignor import StickyPartitionAssignor +from kafka.coordinator.assignors.sticky.sticky_assignor import StickyAssignmentExecutor, StickyPartitionAssignor from kafka.protocol.consumer import JoinGroupResponse @@ -802,6 +802,51 @@ def test_assignment_with_conflicting_previous_generations(mocker, sticky_assigno assert sticky_assignor._latest_partition_movements.are_sticky() +def test_reassignment_with_mixed_subscriptions_terminates(mocker, sticky_assignor): + partitions = {'t5': 1, 't8': 12, 't10': 6, 't11': 11, 't12': 24, 't15': 1, 't17': 4} + cluster = create_cluster(mocker, topics=set(partitions), topic_partitions_lambda=lambda t: set(range(partitions[t]))) + + subscriptions = { + 'C2': {'t11'}, + 'C5': {'t15', 't10'}, + 'C8': {'t11', 't10', 't8'}, + 'C10': {'t10', 't17'}, + } + for i in range(4): + subscriptions['C11-{}'.format(i)] = {'t5', 't12', 't11'} + subscriptions['C12-{}'.format(i)] = {'t5', 't8', 't10'} + member_assignments = { + 'C5': [('t10', 0), ('t10', 1)], + 'C10': [('t10', 4)], + 'C11-0': [('t11', 4), ('t11', 8)], + 'C11-1': [('t11', 5), ('t11', 9)], + 'C11-2': [('t11', 6), ('t11', 10)], + 'C11-3': [('t11', 3), ('t11', 7)], + 'C12-0': [('t10', 5), ('t8', 0), ('t8', 4), ('t8', 8)], + 'C12-1': [('t8', 1), ('t8', 5), ('t8', 9)], + 'C12-2': [('t10', 2), ('t8', 2), ('t8', 6), ('t8', 10)], + 'C12-3': [('t10', 3), ('t8', 3), ('t8', 7), ('t8', 11)], + } + member_metadata = {} + for member, topics in subscriptions.items(): + owned = [TopicPartition(t, p) for t, p in member_assignments.get(member, [])] + member_metadata[member] = sticky_assignor._metadata(topics, owned) + + # this input used to loop forever in _perform_reassignments; fail fast instead of hanging + move_partition = StickyAssignmentExecutor._move_partition + moves = [] + + def capped_move_partition(self, partition, new_consumer): + moves.append(partition) + assert len(moves) < 1000, 'Too many partition movements; reassignment is not converging' + move_partition(self, partition, new_consumer) + mocker.patch.object(StickyAssignmentExecutor, '_move_partition', capped_move_partition) + + assignment = sticky_assignor.assign(cluster, make_join_group_response_members(member_metadata)) + verify_validity_and_balance(subscriptions, assignment) + assert sticky_assignor._latest_partition_movements.are_sticky() + + def assert_assignment(result_assignment, expected_assignment): assert result_assignment == expected_assignment assert set(result_assignment) == set(expected_assignment)