StickyPartitionAssignor.assign() loops forever when members subscribe to different topics

The hang is in StickyAssignmentExecutor._perform_reassignments in kafka/coordinator/assignors/sticky/sticky_assignor.py, which calls _reassign_partition_to_consumer and PartitionMovements.get_partition_to_be_moved. Run the reproduction script from the issue under faulthandler to confirm it prints 'assignment completed' once fixed. The work is already under way in linked PR #3164, which includes a regression test, so review or extend that PR rather than starting a parallel fix.

Summary

With StickyPartitionAssignor (and, since 3.0, CooperativeStickyAssignor, which uses the same executor), the group leader can hang forever in StickyAssignmentExecutor._perform_reassignments. Members subscribe to different topic sets. The leader never sends SyncGroup, so the whole group stalls.

Reproduced on master (3.1.0.dev), 3.0.11 and 2.3.2. The _reassign_partition_to_consumer code is identical on all three. For context, Apache Kafka's Java client (4.3.1) has the same logic and hangs on the same input.

Reproduction

Every member reports its owned partitions with generation -1, the default when on_assignment has not run yet.

import faulthandler

from kafka.coordinator.assignors.sticky.sticky_assignor import StickyPartitionAssignor
from kafka.protocol.consumer import JoinGroupResponse
from kafka.structs import TopicPartition

PARTITIONS = {'t5': 1, 't8': 12, 't10': 6, 't11': 11, 't12': 24, 't15': 1, 't17': 4}
SUBSCRIPTIONS = {
    'c2-p3': ['t11'],
    'c5-p2': ['t15', 't10'],
    'c8-p0': ['t11', 't10', 't8'],
    'c10-p3': ['t10', 't17'],
    **{'c11-p%d' % i: ['t5', 't12', 't11'] for i in range(4)},
    **{'c12-p%d' % i: ['t5', 't8', 't10'] for i in range(4)},
}
OWNED = {
    'c5-p2': [('t10', 0), ('t10', 1)],
    'c10-p3': [('t10', 4)],
    'c11-p0': [('t11', 4), ('t11', 8)],
    'c11-p1': [('t11', 5), ('t11', 9)],
    'c11-p2': [('t11', 6), ('t11', 10)],
    'c11-p3': [('t11', 3), ('t11', 7)],
    'c12-p0': [('t10', 5), ('t8', 0), ('t8', 4), ('t8', 8)],
    'c12-p1': [('t8', 1), ('t8', 5), ('t8', 9)],
    'c12-p2': [('t10', 2), ('t8', 2), ('t8', 6), ('t8', 10)],
    'c12-p3': [('t10', 3), ('t8', 3), ('t8', 7), ('t8', 11)],
}


class Cluster:
    def topics(self):
        return set(PARTITIONS)

    def partitions_for_topic(self, topic):
        return set(range(PARTITIONS[topic])) if topic in PARTITIONS else None


def member_metadata(member_id):
    assignor = StickyPartitionAssignor()
    assignor.member_assignment = [TopicPartition(t, p) for t, p in OWNED.get(member_id, [])]
    return assignor.metadata(SUBSCRIPTIONS[member_id])


members = [
    JoinGroupResponse.JoinGroupResponseMember(member_id=m, metadata=member_metadata(m))
    for m in SUBSCRIPTIONS
]
faulthandler.dump_traceback_later(10, exit=True)
StickyPartitionAssignor().assign(Cluster(), members)
print('assignment completed')

Expected: assignment completed. Actual: the script never finishes. After 10 seconds, faulthandler shows a stack like this:

  File ".../kafka/coordinator/assignors/sticky/sticky_assignor.py" in _is_balanced
  File ".../kafka/coordinator/assignors/sticky/sticky_assignor.py" in _perform_reassignments
  File ".../kafka/coordinator/assignors/sticky/sticky_assignor.py" in balance
  File ".../kafka/coordinator/assignors/sticky/sticky_assignor.py" in assign

Root cause

_reassign_partition_to_consumer(partition, new_consumer) asks PartitionMovements.get_partition_to_be_moved(partition, consumer, new_consumer) which partition to move. consumer is the current, overloaded owner.

1. Suppose partition already moved from some consumer X to consumer during this rebalance. Then get_partition_to_be_moved replaces consumer with X and looks up the reverse pair new_consumer -> X.
2. If a same-topic partition moved new_consumer -> X, the method returns that partition as the substitute. X owns the substitute, not consumer.
3. _move_partition(substitute, new_consumer) moves it from X back to new_consumer. consumer keeps its load, so _is_balanced() stays false and modified stays true. The while True loop then repeats the same choice forever.

Proposed fix

If X owns the substitute, also move partition back to X. The substitute goes X -> new_consumer, then partition goes consumer -> X. Both movement records cancel, so no stickiness cycle remains. The net effect is that one partition leaves consumer, and new_consumer has at least 2 fewer partitions than consumer. The sum of squared assignment sizes therefore strictly decreases, and the loop terminates.

I have a PR with the fix and a regression test.

I used an AI coding assistant (Claude) to investigate this, write the fix, and draft this report.