feat: attach the rebalance listener to concurrent subscribers automatically

README.md

@@ -105,7 +105,7 @@ The processing engine. Manages:
 - In-flight task tracking via a `set[asyncio.Task]`; each task's done-callback releases the semaphore, removes the task from the set, and logs any non-cancellation exception at ERROR with a traceback
 - FastStream control signals raised by a middleware registered *after* this one are absorbed before they can end the task, so they neither pin the message body via a traceback nor reach error reporters that wrap asyncio tasks. See [Limitations](#faststream-control-signals-from-a-middleware-registered-after-this-one) for which are honoured and which only log
 - A `KafkaBatchCommitter` for offset commits
-- An optional `ConsumerRebalanceListener` (via `handler.create_rebalance_listener()`) that flushes pending commits when partitions are revoked
+- A `ConsumerRebalanceListener` on every concurrent subscriber that flushes pending commits when partitions are revoked. `initialize_concurrent_processing` attaches it automatically, including to subscribers declared on routers; see [Rebalance handling](#how-it-works)
 
 This library does **not** install signal handlers — shutdown is driven by your lifespan / process manager calling `stop_concurrent_processing`.
 
@@ -117,7 +117,7 @@ Runs as a background asyncio task. A streaming loop absorbs `KafkaCommitTask` ob
 
 ### `initialize_concurrent_processing(context, ...)`
 
-Create and start the concurrent processing handler; store it in FastStream's context.
+Create and start the concurrent processing handler; store it in FastStream's context; attach a rebalance listener to every concurrent subscriber. Call it before the broker starts.
 
 | Parameter | Default | Description |
 |---|---|---|
@@ -127,6 +127,7 @@ Create and start the concurrent processing handler; store it in FastStream's con
 | `commit_batch_timeout_sec` | `10.0` | Max seconds before flushing a batch |
 | `shutdown_timeout_sec` | `20.0` | Max seconds the batch committer waits for its background task to drain before forcing cancellation |
 | `max_uncommitted_tasks` | `10000` | Max tasks accepted but not yet committed before the consume path blocks (backpressure). `None` disables the bound. |
+| `rebalance_flush_timeout_sec` | `10.0` | Max seconds the attached rebalance listener waits for in-flight tasks when partitions are revoked |
 
 Returns the `KafkaConcurrentHandler` instance.
 
@@ -166,7 +167,7 @@ modern_di_faststream.setup_di(app, container=container)  # registered after →
 
 3. **Offset committing**: Each dispatched task is paired with its Kafka offset and consumer reference and enqueued in `KafkaBatchCommitter`. Once the task completes, the committer groups offsets by partition and calls `consumer.commit(partitions_to_offsets)` with `offset + 1` (Kafka's "next offset to fetch" convention).
 
-4. **Rebalance handling**: When Kafka revokes a partition, the `ConsumerRebalanceListener` (returned by `handler.create_rebalance_listener(flush_timeout_sec=...)`) calls `committer.commit_all()` to flush pending offsets before the partition is reassigned. The flush waits for in-flight tasks up to `flush_timeout_sec` (default 10 s) so a slow handler cannot stall the rebalance past `max.poll.interval.ms`; on timeout, the remaining in-flight messages are redelivered after reassignment (at-least-once). A future optimization may scope the wait to only the revoked partitions.
+4. **Rebalance handling**: `initialize_concurrent_processing` attaches a `ConsumerRebalanceListener` to every subscriber it processes concurrently: `AckPolicy.MANUAL`, not `batch=True`, subscribed by topic or pattern, on any Kafka broker of the FastStream application in the context, including subscribers from included routers. It must run before the broker starts (in the lifespan, as above), because FastStream hands the listener to aiokafka when the consumer subscribes; if the broker has already started, or the context holds no FastStream application, it logs an ERROR and attaches nothing. A `listener=` you pass yourself is kept and called after the flush. For setups where it cannot attach, pass `listener=ConsumerRebalanceListener.from_context(broker.context)` to the subscriber. When Kafka revokes a partition, the listener calls `committer.commit_all()` to flush pending offsets before the partition is reassigned. The flush waits for in-flight tasks up to `flush_timeout_sec` (default 10 s) so a slow handler cannot stall the rebalance past `max.poll.interval.ms`; on timeout, the remaining in-flight messages are redelivered after reassignment (at-least-once). A future optimization may scope the wait to only the revoked partitions.
 
 5. **Shutdown**: `stop_concurrent_processing` cancels every in-flight asyncio task, then awaits `committer.close()`. The committer treats cancelled tasks as a hard offset boundary — cancelled-and-after offsets stay uncommitted and get redelivered on restart. Total wall-clock is sub-second in normal conditions and bounded by `shutdown_timeout_sec` only as a safety net for stuck network commits.
 

faststream_concurrent_aiokafka/middleware.py

@@ -10,7 +10,7 @@
 from faststream.kafka.message import KafkaAckableMessage
 from faststream.middlewares import AckPolicy
 
-from faststream_concurrent_aiokafka import consts
+from faststream_concurrent_aiokafka import consts, rebalance
 from faststream_concurrent_aiokafka.batch_committer import CommitterIsDeadError, KafkaBatchCommitter
 from faststream_concurrent_aiokafka.processing import KafkaConcurrentHandler
 
@@ -259,6 +259,7 @@ async def initialize_concurrent_processing(  # noqa: PLR0913, PLR0917
     commit_batch_timeout_sec: float = consts.DEFAULT_COMMIT_BATCH_TIMEOUT_SEC,
     shutdown_timeout_sec: float = consts.DEFAULT_SHUTDOWN_TIMEOUT_SEC,
     max_uncommitted_tasks: int | None = consts.DEFAULT_MAX_UNCOMMITTED_TASKS,
+    rebalance_flush_timeout_sec: float = consts.DEFAULT_REBALANCE_FLUSH_TIMEOUT_SEC,
 ) -> KafkaConcurrentHandler:
     existing: KafkaConcurrentHandler | None = context.get(consts.PROCESSING_CONTEXT_KEY)
     if existing and existing.is_running:
@@ -276,6 +277,7 @@ async def initialize_concurrent_processing(  # noqa: PLR0913, PLR0917
     )
     await concurrent_processing.start()
     context.set_global(consts.PROCESSING_CONTEXT_KEY, concurrent_processing)
+    rebalance.attach_rebalance_listeners(context, rebalance_flush_timeout_sec)
     logger.info("Kafka middleware. Concurrent processing is active")
     return concurrent_processing
 

faststream_concurrent_aiokafka/processing.py

@@ -165,6 +165,10 @@ def create_rebalance_listener(
         """
         return ConsumerRebalanceListener(self._committer, flush_timeout_sec)
 
+    @property
+    def committer(self) -> KafkaBatchCommitter:
+        return self._committer
+
     @property
     def is_healthy(self) -> bool:
         return self._is_running and self._committer.is_healthy

faststream_concurrent_aiokafka/rebalance.py

@@ -1,13 +1,23 @@
+import inspect
+import logging
 import typing
 
 from aiokafka import ConsumerRebalanceListener as BaseConsumerRebalanceListener
+from faststream.kafka.subscriber.usecase import BatchSubscriber, LogicSubscriber
+from faststream.middlewares import AckPolicy
 
 from faststream_concurrent_aiokafka import consts
 from faststream_concurrent_aiokafka.batch_committer import KafkaBatchCommitter
 
 
 if typing.TYPE_CHECKING:
     from aiokafka.structs import TopicPartition
+    from faststream import ContextRepo
+
+    from faststream_concurrent_aiokafka.processing import KafkaConcurrentHandler
+
+
+logger = logging.getLogger(__name__)
 
 
 class ConsumerRebalanceListener(BaseConsumerRebalanceListener):
@@ -17,19 +27,19 @@ class ConsumerRebalanceListener(BaseConsumerRebalanceListener):
     batch-committed will be redelivered to another consumer after a rebalance, causing
     duplicate processing.
 
-    Usage::
-
-        @asynccontextmanager
-        async def lifespan(context: ContextRepo) -> AsyncIterator[None]:
-            handler = await initialize_concurrent_processing(context, ...)
-            listener = handler.create_rebalance_listener()
+    ``initialize_concurrent_processing`` attaches one to every concurrent subscriber (see
+    ``attach_rebalance_listeners``). Where that is not possible, pass one explicitly; the
+    context form looks the handler up on each revocation, so it can be declared before the
+    lifespan creates the handler::
 
-            @broker.subscriber("my-topic", listener=listener)
-            async def handle(msg: str) -> None:
-                ...
-
-    Yield:
-            await stop_concurrent_processing(context)
+        @broker.subscriber(
+            "my-topic",
+            group_id="my-group",
+            ack_policy=AckPolicy.MANUAL,
+            listener=ConsumerRebalanceListener.from_context(broker.context),
+        )
+        async def handle(msg: str) -> None:
+            ...
 
     """
 
@@ -41,11 +51,103 @@ def __init__(
         self._committer = committer
         self._flush_timeout_sec = flush_timeout_sec
 
+    @classmethod
+    def from_context(
+        cls,
+        context: "ContextRepo",
+        flush_timeout_sec: float = consts.DEFAULT_REBALANCE_FLUSH_TIMEOUT_SEC,
+    ) -> "ConsumerRebalanceListener":
+        """Return a listener that finds the running concurrent handler in ``context`` on each revocation.
+
+        A revocation while concurrent processing is not running (before
+        ``initialize_concurrent_processing`` or after ``stop_concurrent_processing``) is a no-op.
+        """
+        return _ContextConsumerRebalanceListener(context, flush_timeout_sec)
+
+    def _resolve_committer(self) -> KafkaBatchCommitter | None:
+        return self._committer
+
     async def on_partitions_assigned(self, _assigned: object) -> None:  # ty: ignore[invalid-method-override]
         pass
 
     async def on_partitions_revoked(self, revoked: object) -> None:
-        await self._committer.commit_all(self._flush_timeout_sec)
+        committer: typing.Final = self._resolve_committer()
+        if committer is None:
+            return
+        await committer.commit_all(self._flush_timeout_sec)
         # The revoked partitions' next assignment (possibly to another consumer) starts
         # fresh, so the cancellation floor — if any was set — must not carry over.
-        self._committer.clear_cancellation_watermarks(typing.cast("typing.Iterable[TopicPartition]", revoked))
+        committer.clear_cancellation_watermarks(typing.cast("typing.Iterable[TopicPartition]", revoked))
+
+
+class _ContextConsumerRebalanceListener(ConsumerRebalanceListener):
+    def __init__(
+        self,
+        context: "ContextRepo",
+        flush_timeout_sec: float,
+        chained: BaseConsumerRebalanceListener | None = None,
+    ) -> None:
+        self._context = context
+        self._flush_timeout_sec = flush_timeout_sec
+        self._chained = chained
+
+    def _resolve_committer(self) -> KafkaBatchCommitter | None:
+        handler: typing.Final[KafkaConcurrentHandler | None] = self._context.get(consts.PROCESSING_CONTEXT_KEY)
+        if handler is None or not handler.is_running:
+            return None
+        return handler.committer
+
+    async def on_partitions_assigned(self, assigned: object) -> None:  # ty: ignore[invalid-method-override]
+        if self._chained is not None:
+            await _call_or_await(self._chained.on_partitions_assigned(assigned))
+
+    async def on_partitions_revoked(self, revoked: object) -> None:
+        await super().on_partitions_revoked(revoked)
+        if self._chained is not None:
+            await _call_or_await(self._chained.on_partitions_revoked(revoked))
+
+
+async def _call_or_await(result: object) -> None:
+    if inspect.isawaitable(result):
+        await result
+
+
+def attach_rebalance_listeners(context: "ContextRepo", flush_timeout_sec: float) -> None:
+    """Give every subscriber this library processes a listener that flushes on revocation.
+
+    Must run before the broker starts: FastStream hands the listener to aiokafka when the
+    consumer subscribes. A listener the user passed is kept and called after the flush.
+    """
+    brokers: typing.Final = getattr(context.get("app"), "brokers", None)
+    if not brokers:
+        logger.error(
+            "Kafka middleware. No FastStream application in the context, so no rebalance listener was attached; "
+            "pass listener=ConsumerRebalanceListener.from_context(broker.context) to each concurrent subscriber"
+        )
+        return
+    for broker in brokers:
+        for subscriber in broker.subscribers:
+            _attach(subscriber, context, flush_timeout_sec, broker_running=broker.running)
+
+
+def _attach(subscriber: object, context: "ContextRepo", flush_timeout_sec: float, *, broker_running: bool) -> None:
+    if (
+        not isinstance(subscriber, LogicSubscriber)
+        or isinstance(subscriber, BatchSubscriber)
+        or subscriber.ack_policy is not AckPolicy.MANUAL
+        or not (subscriber.topics or subscriber.pattern)
+        or isinstance(subscriber._listener, ConsumerRebalanceListener)  # noqa: SLF001
+    ):
+        return
+    if broker_running:
+        logger.error(
+            "Kafka middleware. Broker already started, so no rebalance listener was attached to the subscriber "
+            "for topics %s; call initialize_concurrent_processing before the broker starts (in the lifespan)",
+            subscriber.topics or subscriber.pattern,
+        )
+        return
+    subscriber._listener = _ContextConsumerRebalanceListener(  # noqa: SLF001
+        context,
+        flush_timeout_sec,
+        chained=subscriber._listener,  # noqa: SLF001
+    )

tests/test_integration.py

@@ -4,7 +4,7 @@
 import uuid
 
 from aiokafka.admin import AIOKafkaAdminClient, NewTopic
-from faststream import BaseMiddleware, ContextRepo
+from faststream import BaseMiddleware, ContextRepo, FastStream
 from faststream.asgi import AsgiFastStream
 from faststream.exceptions import StopApplication
 from faststream.kafka import KafkaBroker, KafkaMessage, KafkaRouter
@@ -566,3 +566,67 @@ async def handler(msg: dict[str, int], message: KafkaMessage) -> None:
     assert len(errors) == 1
     assert "Do not call `message.ack()`" in errors[0]
     assert [m["id"] for m in processed] == [1, 2]
+
+
+async def _committed_offsets(bootstrap_servers: str, group_id: str) -> dict[int, int]:
+    admin: typing.Final = AIOKafkaAdminClient(bootstrap_servers=bootstrap_servers)
+    await admin.start()
+    try:
+        offsets: typing.Final = await admin.list_consumer_group_offsets(group_id)
+    finally:
+        await admin.close()
+    return {tp.partition: meta.offset for tp, meta in offsets.items() if meta.offset >= 0}
+
+
+async def _join_group_and_wait_for_commit(bootstrap_servers: str, topic: str, group: str) -> dict[int, int]:
+    """Join `group` with a second member, forcing a rebalance, and poll until an offset is committed."""
+    broker: typing.Final = KafkaBroker(bootstrap_servers)
+
+    @broker.subscriber(topic, group_id=group, auto_offset_reset="earliest", ack_policy=AckPolicy.MANUAL)
+    async def handler(_msg: dict[str, int]) -> None: ...
+
+    async with broker:
+        await broker.start()
+        deadline: typing.Final = asyncio.get_running_loop().time() + 20
+        committed: dict[int, int] = {}
+        while not committed and asyncio.get_running_loop().time() < deadline:
+            await asyncio.sleep(0.5)
+            committed = await _committed_offsets(bootstrap_servers, group)
+        return committed
+
+
+async def test_real_kafka_attached_listener_commits_on_rebalance(kafka_bootstrap_servers: str) -> None:
+    """The listener attached by initialize_concurrent_processing flushes finished work when a member joins.
+
+    The batch size and timeout are far out of reach, so the only thing that can commit the offset
+    while both brokers are running is the revoke callback.
+    """
+    topic: typing.Final = _topic("rebalance")
+    group: typing.Final = f"rebalance-group-{uuid.uuid4().hex[:6]}"
+    processed: typing.Final = asyncio.Event()
+    broker: typing.Final = _broker(kafka_bootstrap_servers)
+    FastStream(broker)
+
+    @broker.subscriber(topic, group_id=group, auto_offset_reset="earliest", ack_policy=AckPolicy.MANUAL)
+    async def handler(_msg: dict[str, int]) -> None:
+        processed.set()
+
+    await _create_topic(kafka_bootstrap_servers, topic)
+    async with broker:
+        await initialize_concurrent_processing(
+            context=broker.context, commit_batch_size=100, commit_batch_timeout_sec=600, concurrency_limit=5
+        )
+        try:
+            await broker.start()
+            await asyncio.sleep(CONSUMER_READY_SLEEP)
+            await broker.publish({"id": 1}, topic=topic)
+            await asyncio.wait_for(processed.wait(), timeout=POLL_SLEEP)
+            assert await _committed_offsets(kafka_bootstrap_servers, group) == {}
+
+            # A separate task keeps Python 3.11's coverage tracer on this frame once the second broker stops.
+            joined: typing.Final = asyncio.create_task(
+                _join_group_and_wait_for_commit(kafka_bootstrap_servers, topic, group)
+            )
+            assert await joined == {0: 1}
+        finally:
+            await stop_concurrent_processing(broker.context)

tests/test_rebalance.py

@@ -5,8 +5,10 @@
 import aiokafka
 import pytest
 from aiokafka.structs import TopicPartition
+from faststream._internal.context import ContextRepo
 
 from faststream_concurrent_aiokafka import consts
+from faststream_concurrent_aiokafka.processing import KafkaConcurrentHandler
 from faststream_concurrent_aiokafka.rebalance import ConsumerRebalanceListener
 from tests.mocks import MockKafkaBatchCommitter
 
@@ -107,3 +109,40 @@ def test_the_rebalance_flush_default_stays_under_aiokafkas_max_poll_interval() -
         inspect.signature(aiokafka.AIOKafkaConsumer.__init__).parameters["max_poll_interval_ms"].default / 1000
     )
     assert aiokafka_default_sec > consts.DEFAULT_REBALANCE_FLUSH_TIMEOUT_SEC
+
+
+async def test_from_context_resolves_a_handler_registered_after_the_listener(
+    committer: MockKafkaBatchCommitter,
+) -> None:
+    """Subscribers are declared before the lifespan creates the handler, so lookup happens on revoke."""
+    context: typing.Final = ContextRepo()
+    listener: typing.Final = ConsumerRebalanceListener.from_context(context, flush_timeout_sec=2.5)
+    handler: typing.Final = KafkaConcurrentHandler(committer=committer)  # ty: ignore[invalid-argument-type]
+    await handler.start()
+    context.set_global(consts.PROCESSING_CONTEXT_KEY, handler)
+    revoked: typing.Final = {TopicPartition(topic="t", partition=0)}
+
+    await listener.on_partitions_revoked(revoked)
+
+    committer.commit_all.assert_called_once_with(2.5)
+    committer.clear_cancellation_watermarks.assert_called_once_with(revoked)
+
+
+async def test_from_context_is_a_noop_without_a_handler() -> None:
+    listener: typing.Final = ConsumerRebalanceListener.from_context(ContextRepo())
+
+    await listener.on_partitions_revoked({TopicPartition(topic="t", partition=0)})
+
+
+async def test_from_context_skips_a_stopped_handler(committer: MockKafkaBatchCommitter) -> None:
+    context: typing.Final = ContextRepo()
+    listener: typing.Final = ConsumerRebalanceListener.from_context(context)
+    context.set_global(
+        consts.PROCESSING_CONTEXT_KEY,
+        KafkaConcurrentHandler(committer=committer),  # ty: ignore[invalid-argument-type]
+    )
+
+    await listener.on_partitions_revoked({TopicPartition(topic="t", partition=0)})
+
+    committer.commit_all.assert_not_called()
+    committer.clear_cancellation_watermarks.assert_not_called()