aio-libs / aio-libs/aiokafka

StickyPartitionAssignor is unsafe when multiple AIOKafkaConsumer instances share a Python process (class-level state clobber)

Abierto
#1,175 0 comentarios 0 reacciones 0 asignados Ver en GitHub
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

Abrir la guía de contribución

Evaluación

Este issue todavía no se ha evaluado.

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.