StickyPartitionAssignor is unsafe when multiple AIOKafkaConsumer instances share a Python process (class-level state clobber)
- Lenguaje dominante
- Python
- Estrellas
- 1.4k
- Forks
- 269
- Merge medio
- 1 d 1 h
- PR fusionados (30 d)
- 6
Descripción
### Description
`StickyPartitionAssignor` stores `member_assignment` and `generation` as **class** attributes ([`sticky_assignor.py:746-747`](https://github.com/aio-libs/aiokafka/blob/master/aiokafka/coordinator/assignors/sticky/sticky_assignor.py)):
```python
class StickyPartitionAssignor(AbstractPartitionAssignor):
...
member_assignment: list[TopicPartition] | None = None
generation: int = DEFAULT_GENERATION_ID
@classmethod
def on_assignment(cls, assignment: ConsumerProtocolMemberAssignment) -> None:
cls.member_assignment = assignment.partitions()
@classmethod
def metadata(cls, topics):
return cls._metadata(topics, cls.member_assignment, cls.generation)
```
Because both slots live on the class, every `AIOKafkaConsumer` in the same Python process that uses this assignor shares them. `on_assignment` fired by one consumer clobbers the state observed by every other consumer.
On the next rebalance, each consumer's `metadata()` reads `cls.member_assignment` — which contains only the *last writer's* partitions — and encodes those partitions as its own `previous_assignment` in the user data sent to the group coordinator. The coordinator receives N members all claiming the same partitions at `generation=-1`. `_init_current_assignments` sees the conflict and logs
```
Partition TopicPartition(topic='...', partition=...) is assigned to multiple consumers following sticky assignment generation -1
```
(line 264) for every collision. The sticky assignment is corrupt for at least N-1 members, partitions get dropped, and consumers stall or duplicate work.
This bites structurally when a framework instantiates multiple `AIOKafkaConsumer` in one process — for example [FastStream](https://github.com/ag2ai/faststream) creates one `AIOKafkaConsumer` per `@broker.subscriber` decorator. Any application that instantiates two consumers in the same process (multiple subscriber patterns, one process per topic, two consumer groups in the same worker, etc.) is exposed.
### Reproducer
Self-contained, no Kafka broker needed:
```python
from aiokafka.coordinator.assignors.sticky.sticky_assignor import StickyPartitionAssignor
from aiokafka.coordinator.protocol import ConsumerProtocolMemberAssignment
# Two AIOKafkaConsumer in one process, same group, disjoint topic subscriptions.
# Simulate the coordinator handing each its assignment via on_assignment.
a = ConsumerProtocolMemberAssignment(version=0,
assignment=[("topic-a", [0, 1, 2])],
user_data=b"")
b = ConsumerProtocolMemberAssignment(version=0,
assignment=[("topic-b", [0, 1, 2])],
user_data=b"")
StickyPartitionAssignor.on_assignment(a) # Consumer 1 gets its assignment
after_a = list(StickyPartitionAssignor.member_assignment)
StickyPartitionAssignor.on_assignment(b) # Consumer 2 gets its assignment
after_b = list(StickyPartitionAssignor.member_assignment)
assert after_a != after_b # Consumer 1's state was overwritten
# Next rebalance: each consumer calls metadata() with its OWN topic list,
# but they all read the same cls.member_assignment.
md1 = StickyPartitionAssignor.metadata(topics=["topic-a"])
md2 = StickyPartitionAssignor.metadata(topics=["topic-b"])
# Both consumers' previous_assignment user_data is identical -- both claim
# topic-b partitions, even though consumer 1 subscribes to topic-a.
assert md1.user_data == md2.user_data # The bug.
```
Output:
```
After A: [TopicPartition(topic='topic-a', partition=0), ..., partition=2)]
After B: [TopicPartition(topic='topic-b', partition=0), ..., partition=2)]
A leaked/lost? True
Consumer 1 metadata.user_data (claims): 000000010007746f7069632d6200000003...
Consumer 2 metadata.user_data (claims): 000000010007746f7069632d6200000003...
Same user_data leaks between two consumers? True
```
### Expected behavior
Each `AIOKafkaConsumer` maintains its own view of previously-assigned partitions, so its `metadata()` output correctly reflects what *it* previously owned. Two consumers subscribed to disjoint topic sets should send disjoint `previous_assignment` user data.
### Actual behavior
Both consumers send the same `previous_assignment` (whichever consumer's `on_assignment` fired last), and the coordinator sees them claim the same partitions.
### Environment
- aiokafka 0.13.0
- Python 3.12
- Kafka broker version irrelevant — the bug is entirely on the client side, in the assignor's Python-level state.
### Suggested fix
Two options:
1. **Move `member_assignment` and `generation` to instance attributes.** Each `AbstractPartitionAssignor` subclass instance the coordinator manages would then have its own slots. This mirrors the reference Java client (`kafka-clients` stores this state on the instance in `AbstractStickyAssignor`).
2. **Mint per-consumer subclasses at coordinator init time.** When `_lookup_assignor` first resolves an assignor for a consumer, wrap the class in a fresh subclass so `cls.member_assignment` and `cls.generation` are per-consumer. Non-invasive relative to option 1.
Callers can work around this today by minting a subclass in application code (using `type(...)` or a small factory) and passing the subclass to each `AIOKafkaConsumer`'s `partition_assignment_strategy`, but that pushes a subtle memory model concern onto every downstream user.
### Related
- `StickyPartitionAssignor.on_generation_assignment` is defined but never invoked anywhere in aiokafka — filed separately.
- KIP-54 defines the sticky assignor semantics.
Guía de contribución
Evaluación
Este issue todavía no se ha evaluado.