Kafka partitions, keys and ordering

import threading
import time
import uuid

from confluent_kafka import (
    Consumer,
    KafkaError,
    Message,
    Producer,
    TopicPartition,
)
from confluent_kafka.admin import AdminClient
from confluent_kafka.cimpl import NewTopic

Conf = dict[str, str | int | float | bool]
BOOT: Conf = {"bootstrap.servers": "127.0.0.1:9092"}
topic = f"orders-{uuid.uuid4().hex[:8]}"
admin = AdminClient(BOOT)
admin.create_topics([NewTopic(topic, num_partitions=3)])[topic].result()

# 1. The same key always lands on the same partition.
seen: dict[str, set[int]] = {}


def report(err: KafkaError | None, msg: Message) -> None:
    key = str(msg.key() or b"", "utf-8")
    seen.setdefault(key, set()).add(msg.partition() or 0)


p = Producer(BOOT)
for _ in range(2):
    for user in ("user-1", "user-2", "user-3"):
        p.produce(topic, key=user, value=b"{}", on_delivery=report)
p.flush(10)
sticky = all(len(v) == 1 for v in seen.values())
print("Same key, same partition:", sticky)

# 2. Four consumers, three partitions: one consumer gets nothing.
owned: dict[str, list[int]] = {}


def member(name: str, stop: threading.Event) -> None:
    def on_assign(c: Consumer, parts: list[TopicPartition]) -> None:
        owned[name] = sorted(tp.partition for tp in parts)

    conf: Conf = {**BOOT, "group.id": topic}
    c = Consumer(conf)
    c.subscribe([topic], on_assign=on_assign)
    while not stop.is_set():
        c.poll(0.2)
    c.close()


stop = threading.Event()
names = ["C1", "C2", "C3", "C4"]
ts = [threading.Thread(target=member, args=(n, stop)) for n in names]
for t in ts:
    t.start()
deadline = time.time() + 40
while time.time() < deadline:
    if len(owned) == 4 and sum(map(len, owned.values())) == 3:
        break
    time.sleep(0.5)
stop.set()
for t in ts:
    t.join()
counts = sorted(len(v) for v in owned.values())
print("Partitions per consumer:", counts)
print("Idle consumers:", counts.count(0))