From ca5e1492b1fe304b2da08caf349a3ac3e188aa43 Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Thu, 24 Sep 2026 21:48:01 +0300 Subject: [PATCH 1/3] feat: build the rebalance listener from the broker context --- README.md | 14 ++++- faststream_concurrent_aiokafka/processing.py | 4 ++ faststream_concurrent_aiokafka/rebalance.py | 62 +++++++++++++++---- tests/test_integration.py | 63 ++++++++++++++++++++ tests/test_rebalance.py | 39 ++++++++++++ 5 files changed, 166 insertions(+), 16 deletions(-) diff --git a/README.md b/README.md index 37ea5c0..8195f53 100644 --- a/README.md +++ b/README.md @@ -55,6 +55,7 @@ from faststream.asgi import AsgiFastStream from faststream.kafka import KafkaBroker from faststream.middlewares import AckPolicy from faststream_concurrent_aiokafka import ( + ConsumerRebalanceListener, KafkaConcurrentProcessingMiddleware, initialize_concurrent_processing, stop_concurrent_processing, @@ -82,7 +83,14 @@ async def lifespan(_context: ContextRepo): app = AsgiFastStream(broker, lifespan=lifespan) -@broker.subscriber("my-topic", group_id="my-group", ack_policy=AckPolicy.MANUAL) +@broker.subscriber( + "my-topic", + group_id="my-group", + ack_policy=AckPolicy.MANUAL, + # Flush finished work when partitions are revoked; the handler is looked up in the + # context on each revocation, so it may be created later by the lifespan. + listener=ConsumerRebalanceListener.from_context(broker.context), +) async def handle(msg: str) -> None: ... @@ -105,7 +113,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` that flushes pending commits when partitions are revoked. Pass `ConsumerRebalanceListener.from_context(broker.context)` as each concurrent subscriber's `listener=`; it finds the running handler at revocation time, so it can be declared before `initialize_concurrent_processing` runs. `handler.create_rebalance_listener()` builds one bound to an existing handler. Without a listener, finished work on revoked partitions is not committed and is redelivered to the new owner after every rebalance This library does **not** install signal handlers — shutdown is driven by your lifespan / process manager calling `stop_concurrent_processing`. @@ -166,7 +174,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**: When Kafka revokes a partition, the `ConsumerRebalanceListener` (from `ConsumerRebalanceListener.from_context(context, flush_timeout_sec=...)` or `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. 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. diff --git a/faststream_concurrent_aiokafka/processing.py b/faststream_concurrent_aiokafka/processing.py index deccaad..695b4ee 100644 --- a/faststream_concurrent_aiokafka/processing.py +++ b/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 diff --git a/faststream_concurrent_aiokafka/rebalance.py b/faststream_concurrent_aiokafka/rebalance.py index 3bb5c61..0998a84 100644 --- a/faststream_concurrent_aiokafka/rebalance.py +++ b/faststream_concurrent_aiokafka/rebalance.py @@ -8,6 +8,9 @@ if typing.TYPE_CHECKING: from aiokafka.structs import TopicPartition + from faststream import ContextRepo + + from faststream_concurrent_aiokafka.processing import KafkaConcurrentHandler class ConsumerRebalanceListener(BaseConsumerRebalanceListener): @@ -17,19 +20,21 @@ 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() + Subscribers are usually declared before the lifespan creates the concurrent handler, + so build the listener from the broker's context; the handler is looked up on each + revocation:: - @broker.subscriber("my-topic", listener=listener) - async def handle(msg: str) -> None: - ... + broker = KafkaBroker(...) + broker.add_middleware(KafkaConcurrentProcessingMiddleware) - 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 +46,42 @@ 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) -> None: + self._context = context + self._flush_timeout_sec = flush_timeout_sec + + 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 diff --git a/tests/test_integration.py b/tests/test_integration.py index a1d3474..5004af7 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -11,6 +11,7 @@ from faststream.middlewares import AckPolicy from faststream_concurrent_aiokafka import ( + ConsumerRebalanceListener, KafkaConcurrentProcessingMiddleware, initialize_concurrent_processing, stop_concurrent_processing, @@ -566,3 +567,65 @@ 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 test_real_kafka_listener_from_context_commits_on_rebalance(kafka_bootstrap_servers: str) -> None: + """A listener declared before the handler exists flushes finished work when a new 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() + broker1: typing.Final = _broker(kafka_bootstrap_servers) + + @broker1.subscriber( + topic, + group_id=group, + auto_offset_reset="earliest", + ack_policy=AckPolicy.MANUAL, + listener=ConsumerRebalanceListener.from_context(broker1.context), + ) + async def handler1(_msg: dict[str, int]) -> None: + processed.set() + + broker2: typing.Final = KafkaBroker(kafka_bootstrap_servers) + + @broker2.subscriber(topic, group_id=group, auto_offset_reset="earliest", ack_policy=AckPolicy.MANUAL) + async def handler2(_msg: dict[str, int]) -> None: ... + + await _create_topic(kafka_bootstrap_servers, topic) + async with broker1: + await initialize_concurrent_processing( + context=broker1.context, commit_batch_size=100, commit_batch_timeout_sec=600, concurrency_limit=5 + ) + try: + await broker1.start() + await asyncio.sleep(CONSUMER_READY_SLEEP) + await broker1.publish({"id": 1}, topic=topic) + await asyncio.wait_for(processed.wait(), timeout=POLL_SLEEP) + assert await _committed_offsets(kafka_bootstrap_servers, group) == {} + + async with broker2: + await broker2.start() + deadline: typing.Final = asyncio.get_running_loop().time() + 20 + committed: dict[int, int] = {} + while asyncio.get_running_loop().time() < deadline: + committed = await _committed_offsets(kafka_bootstrap_servers, group) + if committed: + break + await asyncio.sleep(0.5) + assert committed == {0: 1} + finally: + await stop_concurrent_processing(broker1.context) diff --git a/tests/test_rebalance.py b/tests/test_rebalance.py index 6210508..08dba99 100644 --- a/tests/test_rebalance.py +++ b/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() From 6877821683415e5b440cf11a635f94f01c0b379b Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Thu, 24 Sep 2026 22:57:25 +0300 Subject: [PATCH 2/3] feat: attach the rebalance listener to concurrent subscribers automatically --- README.md | 17 +- faststream_concurrent_aiokafka/middleware.py | 4 +- faststream_concurrent_aiokafka/rebalance.py | 80 +++++++- tests/test_integration.py | 23 +-- tests/test_rebalance_attachment.py | 196 +++++++++++++++++++ 5 files changed, 285 insertions(+), 35 deletions(-) create mode 100644 tests/test_rebalance_attachment.py diff --git a/README.md b/README.md index 8195f53..818f186 100644 --- a/README.md +++ b/README.md @@ -55,7 +55,6 @@ from faststream.asgi import AsgiFastStream from faststream.kafka import KafkaBroker from faststream.middlewares import AckPolicy from faststream_concurrent_aiokafka import ( - ConsumerRebalanceListener, KafkaConcurrentProcessingMiddleware, initialize_concurrent_processing, stop_concurrent_processing, @@ -83,14 +82,7 @@ async def lifespan(_context: ContextRepo): app = AsgiFastStream(broker, lifespan=lifespan) -@broker.subscriber( - "my-topic", - group_id="my-group", - ack_policy=AckPolicy.MANUAL, - # Flush finished work when partitions are revoked; the handler is looked up in the - # context on each revocation, so it may be created later by the lifespan. - listener=ConsumerRebalanceListener.from_context(broker.context), -) +@broker.subscriber("my-topic", group_id="my-group", ack_policy=AckPolicy.MANUAL) async def handle(msg: str) -> None: ... @@ -113,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 -- A `ConsumerRebalanceListener` that flushes pending commits when partitions are revoked. Pass `ConsumerRebalanceListener.from_context(broker.context)` as each concurrent subscriber's `listener=`; it finds the running handler at revocation time, so it can be declared before `initialize_concurrent_processing` runs. `handler.create_rebalance_listener()` builds one bound to an existing handler. Without a listener, finished work on revoked partitions is not committed and is redelivered to the new owner after every rebalance +- 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`. @@ -125,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 | |---|---|---| @@ -135,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. @@ -174,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` (from `ConsumerRebalanceListener.from_context(context, flush_timeout_sec=...)` or `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. diff --git a/faststream_concurrent_aiokafka/middleware.py b/faststream_concurrent_aiokafka/middleware.py index 36cef5a..3663d55 100644 --- a/faststream_concurrent_aiokafka/middleware.py +++ b/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 diff --git a/faststream_concurrent_aiokafka/rebalance.py b/faststream_concurrent_aiokafka/rebalance.py index 0998a84..0d6dd1d 100644 --- a/faststream_concurrent_aiokafka/rebalance.py +++ b/faststream_concurrent_aiokafka/rebalance.py @@ -1,6 +1,10 @@ +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 @@ -13,6 +17,9 @@ from faststream_concurrent_aiokafka.processing import KafkaConcurrentHandler +logger = logging.getLogger(__name__) + + class ConsumerRebalanceListener(BaseConsumerRebalanceListener): """Commits all pending offsets when Kafka revokes partitions during rebalance. @@ -20,12 +27,10 @@ class ConsumerRebalanceListener(BaseConsumerRebalanceListener): batch-committed will be redelivered to another consumer after a rebalance, causing duplicate processing. - Subscribers are usually declared before the lifespan creates the concurrent handler, - so build the listener from the broker's context; the handler is looked up on each - revocation:: - - broker = KafkaBroker(...) - broker.add_middleware(KafkaConcurrentProcessingMiddleware) + ``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", @@ -76,12 +81,73 @@ async def on_partitions_revoked(self, revoked: object) -> None: class _ContextConsumerRebalanceListener(ConsumerRebalanceListener): - def __init__(self, context: "ContextRepo", flush_timeout_sec: float) -> None: + 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 + ) diff --git a/tests/test_integration.py b/tests/test_integration.py index 5004af7..02f043b 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -4,14 +4,13 @@ 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 from faststream.middlewares import AckPolicy from faststream_concurrent_aiokafka import ( - ConsumerRebalanceListener, KafkaConcurrentProcessingMiddleware, initialize_concurrent_processing, stop_concurrent_processing, @@ -579,8 +578,8 @@ async def _committed_offsets(bootstrap_servers: str, group_id: str) -> dict[int, return {tp.partition: meta.offset for tp, meta in offsets.items() if meta.offset >= 0} -async def test_real_kafka_listener_from_context_commits_on_rebalance(kafka_bootstrap_servers: str) -> None: - """A listener declared before the handler exists flushes finished work when a new member joins. +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. @@ -590,13 +589,9 @@ async def test_real_kafka_listener_from_context_commits_on_rebalance(kafka_boots processed: typing.Final = asyncio.Event() broker1: typing.Final = _broker(kafka_bootstrap_servers) - @broker1.subscriber( - topic, - group_id=group, - auto_offset_reset="earliest", - ack_policy=AckPolicy.MANUAL, - listener=ConsumerRebalanceListener.from_context(broker1.context), - ) + FastStream(broker1) + + @broker1.subscriber(topic, group_id=group, auto_offset_reset="earliest", ack_policy=AckPolicy.MANUAL) async def handler1(_msg: dict[str, int]) -> None: processed.set() @@ -621,11 +616,9 @@ async def handler2(_msg: dict[str, int]) -> None: ... await broker2.start() deadline: typing.Final = asyncio.get_running_loop().time() + 20 committed: dict[int, int] = {} - while asyncio.get_running_loop().time() < deadline: - committed = await _committed_offsets(kafka_bootstrap_servers, group) - if committed: - break + while not committed and asyncio.get_running_loop().time() < deadline: await asyncio.sleep(0.5) + committed = await _committed_offsets(kafka_bootstrap_servers, group) assert committed == {0: 1} finally: await stop_concurrent_processing(broker1.context) diff --git a/tests/test_rebalance_attachment.py b/tests/test_rebalance_attachment.py new file mode 100644 index 0000000..82b36d0 --- /dev/null +++ b/tests/test_rebalance_attachment.py @@ -0,0 +1,196 @@ +# ruff: noqa: SLF001 +import inspect +import logging +import typing +from unittest.mock import AsyncMock, Mock + +import pytest +from aiokafka import ConsumerRebalanceListener as BaseConsumerRebalanceListener +from aiokafka.structs import TopicPartition +from faststream import FastStream +from faststream.kafka import KafkaBroker, KafkaRouter +from faststream.kafka.subscriber.usecase import ConcurrentBetweenPartitionsSubscriber, LogicSubscriber +from faststream.middlewares import AckPolicy + +from faststream_concurrent_aiokafka import ( + ConsumerRebalanceListener, + initialize_concurrent_processing, + stop_concurrent_processing, +) + + +async def _handle(_msg: str) -> None: ... + + +def _subscribe(registrator: KafkaBroker | KafkaRouter, *args: typing.Any, **kwargs: typing.Any) -> typing.Any: # noqa: ANN401 + subscriber: typing.Final = registrator.subscriber(*args, **kwargs) + subscriber(_handle) + return subscriber + + +def _broker_in_app() -> KafkaBroker: + broker: typing.Final = KafkaBroker() + FastStream(broker) + return broker + + +async def _initialize(broker: KafkaBroker) -> None: + await initialize_concurrent_processing(context=broker.context, rebalance_flush_timeout_sec=2.5) + + +async def test_listener_is_attached_to_manual_subscribers_on_the_broker_and_its_routers() -> None: + broker: typing.Final = _broker_in_app() + on_broker: typing.Final = _subscribe(broker, "a", group_id="g", ack_policy=AckPolicy.MANUAL) + router: typing.Final = KafkaRouter() + on_router: typing.Final = _subscribe(router, "b", group_id="g", ack_policy=AckPolicy.MANUAL) + broker.include_router(router) + + await _initialize(broker) + try: + for subscriber in (on_broker, on_router): + listener = subscriber._listener + assert isinstance(listener, ConsumerRebalanceListener) + assert listener._flush_timeout_sec == 2.5 + finally: + await stop_concurrent_processing(broker.context) + + +async def test_attached_listener_flushes_the_running_handler() -> None: + broker: typing.Final = _broker_in_app() + subscriber: typing.Final = _subscribe(broker, "a", group_id="g", ack_policy=AckPolicy.MANUAL) + handler: typing.Final = await initialize_concurrent_processing(context=broker.context) + try: + handler.committer.commit_all = AsyncMock() + revoked: typing.Final = {TopicPartition(topic="a", partition=0)} + + await subscriber._listener.on_partitions_revoked(revoked) + + handler.committer.commit_all.assert_awaited_once() + finally: + await stop_concurrent_processing(broker.context) + + +@pytest.mark.parametrize( + "subscriber_kwargs", + [ + pytest.param({"ack_policy": AckPolicy.ACK_FIRST}, id="not-manual"), + pytest.param({"ack_policy": AckPolicy.MANUAL, "batch": True}, id="batch"), + ], +) +async def test_listener_is_not_attached_to_subscribers_this_library_does_not_process( + subscriber_kwargs: dict[str, typing.Any], +) -> None: + broker: typing.Final = _broker_in_app() + subscriber: typing.Final = _subscribe(broker, "a", group_id="g", **subscriber_kwargs) + + await _initialize(broker) + try: + assert subscriber._listener is None + finally: + await stop_concurrent_processing(broker.context) + + +async def test_listener_is_not_attached_to_manually_assigned_partitions() -> None: + """Partitions assigned with `partitions=` never rebalance, so there is nothing to flush.""" + broker: typing.Final = _broker_in_app() + subscriber: typing.Final = _subscribe( + broker, partitions=[TopicPartition(topic="a", partition=0)], ack_policy=AckPolicy.MANUAL + ) + + await _initialize(broker) + try: + assert subscriber._listener is None + finally: + await stop_concurrent_processing(broker.context) + + +@pytest.mark.parametrize("is_async", [True, False]) +async def test_a_user_listener_is_kept_and_still_called(*, is_async: bool) -> None: + user_listener: typing.Final = Mock(spec=BaseConsumerRebalanceListener) + if is_async: + user_listener.on_partitions_revoked = AsyncMock() + user_listener.on_partitions_assigned = AsyncMock() + broker: typing.Final = _broker_in_app() + subscriber: typing.Final = _subscribe( + broker, "a", group_id="g", ack_policy=AckPolicy.MANUAL, listener=user_listener + ) + partitions: typing.Final = {TopicPartition(topic="a", partition=0)} + + handler: typing.Final = await initialize_concurrent_processing(context=broker.context) + try: + handler.committer.commit_all = AsyncMock() + await subscriber._listener.on_partitions_revoked(partitions) + await subscriber._listener.on_partitions_assigned(partitions) + finally: + await stop_concurrent_processing(broker.context) + + handler.committer.commit_all.assert_awaited_once() + user_listener.on_partitions_revoked.assert_called_once_with(partitions) + user_listener.on_partitions_assigned.assert_called_once_with(partitions) + + +async def test_an_explicit_concurrent_listener_is_left_alone() -> None: + broker: typing.Final = _broker_in_app() + explicit: typing.Final = ConsumerRebalanceListener.from_context(broker.context) + subscriber: typing.Final = _subscribe(broker, "a", group_id="g", ack_policy=AckPolicy.MANUAL, listener=explicit) + + await _initialize(broker) + try: + assert subscriber._listener is explicit + finally: + await stop_concurrent_processing(broker.context) + + +async def test_restarting_processing_does_not_wrap_the_listener_twice() -> None: + broker: typing.Final = _broker_in_app() + subscriber: typing.Final = _subscribe(broker, "a", group_id="g", ack_policy=AckPolicy.MANUAL) + await _initialize(broker) + await stop_concurrent_processing(broker.context) + first: typing.Final = subscriber._listener + + await _initialize(broker) + try: + assert subscriber._listener is first + finally: + await stop_concurrent_processing(broker.context) + + +async def test_missing_application_is_logged_as_error(caplog: pytest.LogCaptureFixture) -> None: + broker: typing.Final = KafkaBroker() + subscriber: typing.Final = _subscribe(broker, "a", group_id="g", ack_policy=AckPolicy.MANUAL) + + await _initialize(broker) + try: + assert subscriber._listener is None + assert [r.levelno for r in caplog.records if "rebalance listener" in r.getMessage()] == [logging.ERROR] + finally: + await stop_concurrent_processing(broker.context) + + +async def test_already_started_broker_is_logged_as_error(caplog: pytest.LogCaptureFixture) -> None: + """The listener is fixed when the consumer subscribes, so attaching after start would do nothing.""" + broker: typing.Final = _broker_in_app() + subscriber: typing.Final = _subscribe(broker, "late-topic", group_id="g", ack_policy=AckPolicy.MANUAL) + broker.running = True + + await _initialize(broker) + try: + assert subscriber._listener is None + errors: typing.Final = [r for r in caplog.records if r.levelno == logging.ERROR] + assert len(errors) == 1 + assert "late-topic" in errors[0].getMessage() + finally: + await stop_concurrent_processing(broker.context) + + +@pytest.mark.parametrize("subscriber_class", [LogicSubscriber, ConcurrentBetweenPartitionsSubscriber]) +def test_faststream_still_reads_the_listener_when_the_consumer_subscribes( + subscriber_class: type[LogicSubscriber[typing.Any]], +) -> None: + """INVARIANT: attaching relies on FastStream reading the private `_listener` in `start()`. + + FastStream offers no public way to add a rebalance listener after a subscriber is declared. + If a FastStream release renames the attribute or stops reading it at start, attaching would + silently do nothing; this fails first. + """ + assert "listener=self._listener" in inspect.getsource(subscriber_class.start) From 498a788b5614a792ba05c9a3a09ecb23176debb7 Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Thu, 24 Sep 2026 23:28:15 +0300 Subject: [PATCH 3/3] test: keep the rebalance integration test fully covered on Python 3.11 --- tests/test_integration.py | 54 ++++++++++++++++++++++----------------- 1 file changed, 31 insertions(+), 23 deletions(-) diff --git a/tests/test_integration.py b/tests/test_integration.py index 02f043b..9e1eb8f 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -578,6 +578,23 @@ async def _committed_offsets(bootstrap_servers: str, group_id: str) -> dict[int, 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. @@ -587,38 +604,29 @@ async def test_real_kafka_attached_listener_commits_on_rebalance(kafka_bootstrap topic: typing.Final = _topic("rebalance") group: typing.Final = f"rebalance-group-{uuid.uuid4().hex[:6]}" processed: typing.Final = asyncio.Event() - broker1: typing.Final = _broker(kafka_bootstrap_servers) - - FastStream(broker1) + broker: typing.Final = _broker(kafka_bootstrap_servers) + FastStream(broker) - @broker1.subscriber(topic, group_id=group, auto_offset_reset="earliest", ack_policy=AckPolicy.MANUAL) - async def handler1(_msg: dict[str, int]) -> None: + @broker.subscriber(topic, group_id=group, auto_offset_reset="earliest", ack_policy=AckPolicy.MANUAL) + async def handler(_msg: dict[str, int]) -> None: processed.set() - broker2: typing.Final = KafkaBroker(kafka_bootstrap_servers) - - @broker2.subscriber(topic, group_id=group, auto_offset_reset="earliest", ack_policy=AckPolicy.MANUAL) - async def handler2(_msg: dict[str, int]) -> None: ... - await _create_topic(kafka_bootstrap_servers, topic) - async with broker1: + async with broker: await initialize_concurrent_processing( - context=broker1.context, commit_batch_size=100, commit_batch_timeout_sec=600, concurrency_limit=5 + context=broker.context, commit_batch_size=100, commit_batch_timeout_sec=600, concurrency_limit=5 ) try: - await broker1.start() + await broker.start() await asyncio.sleep(CONSUMER_READY_SLEEP) - await broker1.publish({"id": 1}, topic=topic) + await broker.publish({"id": 1}, topic=topic) await asyncio.wait_for(processed.wait(), timeout=POLL_SLEEP) assert await _committed_offsets(kafka_bootstrap_servers, group) == {} - async with broker2: - await broker2.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(kafka_bootstrap_servers, group) - assert committed == {0: 1} + # 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(broker1.context) + await stop_concurrent_processing(broker.context)