From 5465555e77eb786214c5f393b268fa7295cbb606 Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Fri, 6 Nov 2026 11:13:10 +0300 Subject: [PATCH 1/3] fix: preserve Kafka telemetry spans during concurrent processing --- README.md | 16 +++- faststream_concurrent_aiokafka/telemetry.py | 50 +++++++++++ pyproject.toml | 4 + tests/test_telemetry.py | 96 +++++++++++++++++++++ 4 files changed, 165 insertions(+), 1 deletion(-) create mode 100644 faststream_concurrent_aiokafka/telemetry.py create mode 100644 tests/test_telemetry.py diff --git a/README.md b/README.md index cc17eb5..c8c0d7a 100644 --- a/README.md +++ b/README.md @@ -152,7 +152,7 @@ Returns `True` if the running handler's offset commits have been failing for lon FastStream middleware class. Register it via `broker.add_middleware(...)`. See [Quick start](#quick-start) for usage examples. -> This middleware must be the outermost one. `consume_scope` fires the handler as a background task and returns `None` immediately, so any middleware that wraps it on the outside sees that premature return and misfires: wrong timing, early cleanup, or missed exceptions. Middlewares added after it (inner in the chain) run correctly inside the background task. +> This middleware must be the outermost one. `consume_scope` fires the handler as a background task and returns `None` immediately, so any middleware that wraps it on the outside sees that premature return and misfires: wrong timing, early cleanup, or missed exceptions. The `consume_scope` of a middleware registered after it runs inside the background task. FastStream still controls the timing of `after_processed`. #### DI framework compatibility (`modern-di-faststream` and similar) @@ -168,6 +168,20 @@ modern_di_faststream.setup_di(app, container=container) # registered after → ## How it works +### OpenTelemetry + +Install `faststream-concurrent-aiokafka[otel]` and use +`KafkaConcurrentTelemetryMiddleware` in place of FastStream's `KafkaTelemetryMiddleware`. +Register it after the concurrency middleware. It accepts the same telemetry options and +keeps the consumer process span open until the task finishes. + +```python +from faststream_concurrent_aiokafka.telemetry import KafkaConcurrentTelemetryMiddleware + +broker.add_middleware(KafkaConcurrentProcessingMiddleware) +broker.add_middleware(KafkaConcurrentTelemetryMiddleware()) +``` + ### Message dispatch On each incoming message, `consume_scope` calls `handle_task()`, which acquires a semaphore slot then fires the handler coroutine as a background `asyncio.Task`. diff --git a/faststream_concurrent_aiokafka/telemetry.py b/faststream_concurrent_aiokafka/telemetry.py new file mode 100644 index 0000000..9eff4ea --- /dev/null +++ b/faststream_concurrent_aiokafka/telemetry.py @@ -0,0 +1,50 @@ +from collections.abc import Awaitable, Callable +from types import TracebackType +from typing import Any + +from faststream import ContextRepo +from faststream.kafka.opentelemetry import KafkaTelemetryMiddleware +from faststream.kafka.response import KafkaPublishCommand +from faststream.message import StreamMessage +from faststream.opentelemetry.middleware import BaseTelemetryMiddleware + + +class _ConcurrentTelemetryScope(BaseTelemetryMiddleware[KafkaPublishCommand]): + _consume_finished = False + _cleanup_requested = False + + async def consume_scope( + self, + call_next: Callable[[StreamMessage[Any]], Awaitable[Any]], + msg: StreamMessage[Any], + ) -> Any: # noqa: ANN401 + try: + return await super().consume_scope(call_next, msg) + finally: + self._consume_finished = True + if self._cleanup_requested: + await super().after_processed() + + async def after_processed( + self, + exc_type: type[BaseException] | None = None, + exc_val: BaseException | None = None, + exc_tb: TracebackType | None = None, + ) -> bool | None: + self._cleanup_requested = True + if self._consume_finished: + return await super().after_processed(exc_type, exc_val, exc_tb) + return False + + +class KafkaConcurrentTelemetryMiddleware(KafkaTelemetryMiddleware): + """Kafka telemetry that keeps the process span open through concurrent processing.""" + + def __call__(self, msg: object, /, *, context: ContextRepo) -> BaseTelemetryMiddleware[KafkaPublishCommand]: + return _ConcurrentTelemetryScope( + msg, + tracer=self._tracer, + metrics_container=self._metrics, + settings_provider_factory=self._settings_provider_factory, + context=context, + ) diff --git a/pyproject.toml b/pyproject.toml index 15ae31d..d0d2589 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,9 @@ dependencies = [ "faststream[kafka]>=0.7.1,<0.8", ] +[project.optional-dependencies] +otel = ["faststream[otel]>=0.7.8,<0.8"] + [project.urls] Homepage = "https://modern-python.org" Repository = "https://github.com/modern-python/faststream-concurrent-aiokafka" @@ -42,6 +45,7 @@ Changelog = "https://github.com/modern-python/faststream-concurrent-aiokafka/rel [dependency-groups] dev = [ + "faststream[otel]>=0.7.8,<0.8", "pytest>=8", "pytest-cov>=5", "pytest-xdist>=3", diff --git a/tests/test_telemetry.py b/tests/test_telemetry.py new file mode 100644 index 0000000..91d8ae8 --- /dev/null +++ b/tests/test_telemetry.py @@ -0,0 +1,96 @@ +import asyncio +import time +from collections.abc import Sequence +from uuid import uuid4 + +import pytest +from faststream.kafka import KafkaBroker +from faststream.middlewares import AckPolicy +from opentelemetry import trace +from opentelemetry.sdk.trace import ReadableSpan, TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExportResult +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from faststream_concurrent_aiokafka import ( + ConsumerRebalanceListener, + KafkaConcurrentProcessingMiddleware, + initialize_concurrent_processing, + stop_concurrent_processing, +) +from faststream_concurrent_aiokafka.telemetry import KafkaConcurrentTelemetryMiddleware + + +class _ProcessSpanExporter(InMemorySpanExporter): + def __init__(self) -> None: + super().__init__() + self.process_exported = asyncio.Event() + + def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: + result = super().export(spans) + if any(span.kind is trace.SpanKind.CONSUMER for span in spans): + self.process_exported.set() + return result + + +@pytest.mark.parametrize("outcome", ["success", "error", "cancel", "inline"]) +async def test_process_span_follows_concurrent_task(kafka_bootstrap_servers: str, outcome: str) -> None: + topic = f"telemetry-{uuid4().hex}" + exporter = _ProcessSpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + broker = KafkaBroker( + kafka_bootstrap_servers, + middlewares=(KafkaConcurrentProcessingMiddleware, KafkaConcurrentTelemetryMiddleware(tracer_provider=provider)), + ) + started = asyncio.Event() + release = asyncio.Event() + completed: list[int] = [] + handler_task: asyncio.Task | None = None + + @broker.subscriber( + topic, + group_id=topic, + auto_offset_reset="earliest", + ack_policy=AckPolicy.ACK if outcome == "inline" else AckPolicy.MANUAL, + listener=ConsumerRebalanceListener.from_context(broker.context), + ) + async def handler(body: str) -> None: + nonlocal handler_task + handler_task = asyncio.current_task() + started.set() + try: + await release.wait() + if outcome == "error": + raise ValueError(body) + finally: + completed.append(time.time_ns()) + + try: + async with broker: + await initialize_concurrent_processing(broker.context) + try: + await broker.start() + await broker.publish("test exception", topic) + await asyncio.wait_for(started.wait(), 15) + assert handler_task is not None + assert not exporter.process_exported.is_set() + if outcome == "cancel": + handler_task.cancel() + else: + release.set() + await asyncio.wait_for(exporter.process_exported.wait(), 5) + finally: + release.set() + await stop_concurrent_processing(broker.context) + + (span,) = [span for span in exporter.get_finished_spans() if span.kind is trace.SpanKind.CONSUMER] + assert span.name == f"{topic} process" + assert span.end_time is not None + assert span.end_time >= completed[0] + if outcome == "error": + assert span.status.status_code is trace.StatusCode.ERROR + assert any(event.name == "exception" for event in span.events) + if outcome == "cancel": + assert handler_task.cancelled() + finally: + provider.shutdown() From d9e2789173782ba44f176d7e3cb7998e639e4648 Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Fri, 6 Nov 2026 11:37:30 +0300 Subject: [PATCH 2/3] fix: adapt existing Kafka telemetry during initialization --- README.md | 16 ++--- faststream_concurrent_aiokafka/middleware.py | 5 ++ faststream_concurrent_aiokafka/telemetry.py | 64 ++++++++++++++------ pyproject.toml | 5 +- tests/test_telemetry.py | 46 ++++++++------ 5 files changed, 85 insertions(+), 51 deletions(-) diff --git a/README.md b/README.md index c8c0d7a..f60b530 100644 --- a/README.md +++ b/README.md @@ -170,17 +170,11 @@ modern_di_faststream.setup_di(app, container=container) # registered after → ### OpenTelemetry -Install `faststream-concurrent-aiokafka[otel]` and use -`KafkaConcurrentTelemetryMiddleware` in place of FastStream's `KafkaTelemetryMiddleware`. -Register it after the concurrency middleware. It accepts the same telemetry options and -keeps the consumer process span open until the task finishes. - -```python -from faststream_concurrent_aiokafka.telemetry import KafkaConcurrentTelemetryMiddleware - -broker.add_middleware(KafkaConcurrentProcessingMiddleware) -broker.add_middleware(KafkaConcurrentTelemetryMiddleware()) -``` +If your application already registers FastStream's `KafkaTelemetryMiddleware` after +`KafkaConcurrentProcessingMiddleware`, `initialize_concurrent_processing` automatically +keeps its consumer process spans open through background processing. Initialize in the +FastStream application lifespan before the broker starts, as shown in the quick start. +Applications without telemetry can use the base package. ### Message dispatch diff --git a/faststream_concurrent_aiokafka/middleware.py b/faststream_concurrent_aiokafka/middleware.py index 873da7f..6b61e7f 100644 --- a/faststream_concurrent_aiokafka/middleware.py +++ b/faststream_concurrent_aiokafka/middleware.py @@ -7,12 +7,14 @@ import weakref from faststream import BaseMiddleware, ContextRepo +from faststream.kafka import KafkaBroker from faststream.kafka.message import KafkaAckableMessage from faststream.middlewares import AckPolicy from faststream_concurrent_aiokafka import consts, rebalance from faststream_concurrent_aiokafka.batch_committer import CommitterIsDeadError, KafkaBatchCommitter from faststream_concurrent_aiokafka.processing import KafkaConcurrentHandler +from faststream_concurrent_aiokafka.telemetry import adapt_kafka_telemetry if typing.TYPE_CHECKING: @@ -277,6 +279,9 @@ async def initialize_concurrent_processing( # noqa: PLR0913, PLR0917 ), concurrency_limit=concurrency_limit, ) + for broker in getattr(context.get("app"), "brokers", ()): + if isinstance(broker, KafkaBroker) and KafkaConcurrentProcessingMiddleware in broker.middlewares: + adapt_kafka_telemetry(broker) await concurrent_processing.start() context.set_global(consts.PROCESSING_CONTEXT_KEY, concurrent_processing) rebalance.attach_rebalance_listeners(context, rebalance_flush_timeout_sec) diff --git a/faststream_concurrent_aiokafka/telemetry.py b/faststream_concurrent_aiokafka/telemetry.py index 9eff4ea..45829e2 100644 --- a/faststream_concurrent_aiokafka/telemetry.py +++ b/faststream_concurrent_aiokafka/telemetry.py @@ -1,17 +1,32 @@ +import sys from collections.abc import Awaitable, Callable from types import TracebackType from typing import Any -from faststream import ContextRepo -from faststream.kafka.opentelemetry import KafkaTelemetryMiddleware +from faststream import BaseMiddleware, ContextRepo +from faststream._internal.broker.registrator import Registrator +from faststream._internal.configs.broker import BrokerConfig +from faststream._internal.types import BrokerMiddleware from faststream.kafka.response import KafkaPublishCommand from faststream.message import StreamMessage -from faststream.opentelemetry.middleware import BaseTelemetryMiddleware -class _ConcurrentTelemetryScope(BaseTelemetryMiddleware[KafkaPublishCommand]): - _consume_finished = False - _cleanup_requested = False +class _ConcurrentTelemetryScope(BaseMiddleware[KafkaPublishCommand]): + def __init__(self, scope: BaseMiddleware[KafkaPublishCommand]) -> None: + super().__init__(scope.msg, context=scope.context) + self._scope = scope + self._consume_finished = False + self._cleanup_requested = False + + async def on_receive(self) -> None: + await self._scope.on_receive() + + async def publish_scope( + self, + call_next: Callable[[KafkaPublishCommand], Awaitable[Any]], + cmd: KafkaPublishCommand, + ) -> Any: # noqa: ANN401 + return await self._scope.publish_scope(call_next, cmd) async def consume_scope( self, @@ -19,11 +34,11 @@ async def consume_scope( msg: StreamMessage[Any], ) -> Any: # noqa: ANN401 try: - return await super().consume_scope(call_next, msg) + return await self._scope.consume_scope(call_next, msg) finally: self._consume_finished = True if self._cleanup_requested: - await super().after_processed() + await self._scope.after_processed() async def after_processed( self, @@ -33,18 +48,29 @@ async def after_processed( ) -> bool | None: self._cleanup_requested = True if self._consume_finished: - return await super().after_processed(exc_type, exc_val, exc_tb) + return await self._scope.after_processed(exc_type, exc_val, exc_tb) return False -class KafkaConcurrentTelemetryMiddleware(KafkaTelemetryMiddleware): - """Kafka telemetry that keeps the process span open through concurrent processing.""" +class _ConcurrentTelemetryFactory: + def __init__(self, factory: BrokerMiddleware[Any, KafkaPublishCommand]) -> None: + self._factory = factory + + def __call__(self, msg: object, /, *, context: ContextRepo) -> _ConcurrentTelemetryScope: + return _ConcurrentTelemetryScope(self._factory(msg, context=context)) + - def __call__(self, msg: object, /, *, context: ContextRepo) -> BaseTelemetryMiddleware[KafkaPublishCommand]: - return _ConcurrentTelemetryScope( - msg, - tracer=self._tracer, - metrics_container=self._metrics, - settings_provider_factory=self._settings_provider_factory, - context=context, - ) +def adapt_kafka_telemetry(registrator: Registrator) -> None: + # A registered FastStream telemetry factory already loaded this module. + module = sys.modules.get("faststream.kafka.opentelemetry.middleware") + stock_factory = getattr(module, "KafkaTelemetryMiddleware", None) + for config in registrator.config.configs: + if isinstance(config, BrokerConfig) and any( + type(factory) is stock_factory for factory in config.broker_middlewares + ): + config.broker_middlewares = tuple( + _ConcurrentTelemetryFactory(factory) if type(factory) is stock_factory else factory + for factory in config.broker_middlewares + ) + for router in registrator.routers: + adapt_kafka_telemetry(router) diff --git a/pyproject.toml b/pyproject.toml index d0d2589..0a6a59d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,9 +34,6 @@ dependencies = [ "faststream[kafka]>=0.7.1,<0.8", ] -[project.optional-dependencies] -otel = ["faststream[otel]>=0.7.8,<0.8"] - [project.urls] Homepage = "https://modern-python.org" Repository = "https://github.com/modern-python/faststream-concurrent-aiokafka" @@ -45,7 +42,7 @@ Changelog = "https://github.com/modern-python/faststream-concurrent-aiokafka/rel [dependency-groups] dev = [ - "faststream[otel]>=0.7.8,<0.8", + "faststream[otel]", "pytest>=8", "pytest-cov>=5", "pytest-xdist>=3", diff --git a/tests/test_telemetry.py b/tests/test_telemetry.py index 91d8ae8..7721b5f 100644 --- a/tests/test_telemetry.py +++ b/tests/test_telemetry.py @@ -4,7 +4,9 @@ from uuid import uuid4 import pytest -from faststream.kafka import KafkaBroker +from faststream import FastStream +from faststream.kafka import KafkaBroker, KafkaRouter +from faststream.kafka.opentelemetry import KafkaTelemetryMiddleware from faststream.middlewares import AckPolicy from opentelemetry import trace from opentelemetry.sdk.trace import ReadableSpan, TracerProvider @@ -12,47 +14,57 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from faststream_concurrent_aiokafka import ( - ConsumerRebalanceListener, KafkaConcurrentProcessingMiddleware, initialize_concurrent_processing, stop_concurrent_processing, ) -from faststream_concurrent_aiokafka.telemetry import KafkaConcurrentTelemetryMiddleware class _ProcessSpanExporter(InMemorySpanExporter): - def __init__(self) -> None: + def __init__(self, topic: str) -> None: super().__init__() + self.process_name = f"{topic} process" self.process_exported = asyncio.Event() def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: result = super().export(spans) - if any(span.kind is trace.SpanKind.CONSUMER for span in spans): + if any(span.name == self.process_name for span in spans): self.process_exported.set() return result +def _broker_with_telemetry( + bootstrap_servers: str, + provider: TracerProvider, + *, + router: bool, +) -> tuple[KafkaBroker, KafkaBroker | KafkaRouter]: + broker = KafkaBroker(bootstrap_servers, middlewares=(KafkaConcurrentProcessingMiddleware,)) + registrator = KafkaRouter() if router else broker + registrator.add_middleware(KafkaTelemetryMiddleware(tracer_provider=provider)) + if isinstance(registrator, KafkaRouter): + broker.include_router(registrator) + FastStream(broker) + return broker, registrator + + @pytest.mark.parametrize("outcome", ["success", "error", "cancel", "inline"]) async def test_process_span_follows_concurrent_task(kafka_bootstrap_servers: str, outcome: str) -> None: topic = f"telemetry-{uuid4().hex}" - exporter = _ProcessSpanExporter() + exporter = _ProcessSpanExporter(topic) provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) - broker = KafkaBroker( - kafka_bootstrap_servers, - middlewares=(KafkaConcurrentProcessingMiddleware, KafkaConcurrentTelemetryMiddleware(tracer_provider=provider)), - ) + broker, registrator = _broker_with_telemetry(kafka_bootstrap_servers, provider, router=outcome == "error") started = asyncio.Event() release = asyncio.Event() completed: list[int] = [] handler_task: asyncio.Task | None = None - @broker.subscriber( + @registrator.subscriber( topic, group_id=topic, auto_offset_reset="earliest", ack_policy=AckPolicy.ACK if outcome == "inline" else AckPolicy.MANUAL, - listener=ConsumerRebalanceListener.from_context(broker.context), ) async def handler(body: str) -> None: nonlocal handler_task @@ -67,8 +79,9 @@ async def handler(body: str) -> None: try: async with broker: - await initialize_concurrent_processing(broker.context) + processing = await initialize_concurrent_processing(broker.context) try: + assert await initialize_concurrent_processing(broker.context) is processing await broker.start() await broker.publish("test exception", topic) await asyncio.wait_for(started.wait(), 15) @@ -83,10 +96,9 @@ async def handler(body: str) -> None: release.set() await stop_concurrent_processing(broker.context) - (span,) = [span for span in exporter.get_finished_spans() if span.kind is trace.SpanKind.CONSUMER] - assert span.name == f"{topic} process" - assert span.end_time is not None - assert span.end_time >= completed[0] + (span,) = [span for span in exporter.get_finished_spans() if span.name == exporter.process_name] + assert span.kind is trace.SpanKind.CONSUMER + assert (span.end_time or 0) >= completed[0] if outcome == "error": assert span.status.status_code is trace.StatusCode.ERROR assert any(event.name == "exception" for event in span.events) From 3b5c74aec09f7eded0130fc8d4951b9643c41d72 Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Fri, 6 Nov 2026 11:50:38 +0300 Subject: [PATCH 3/3] fix: align telemetry wrapper with FastStream middleware types --- faststream_concurrent_aiokafka/telemetry.py | 12 ++++++------ tests/test_middleware.py | 6 +++--- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/faststream_concurrent_aiokafka/telemetry.py b/faststream_concurrent_aiokafka/telemetry.py index 45829e2..84ec10a 100644 --- a/faststream_concurrent_aiokafka/telemetry.py +++ b/faststream_concurrent_aiokafka/telemetry.py @@ -7,12 +7,12 @@ from faststream._internal.broker.registrator import Registrator from faststream._internal.configs.broker import BrokerConfig from faststream._internal.types import BrokerMiddleware -from faststream.kafka.response import KafkaPublishCommand from faststream.message import StreamMessage +from faststream.response import PublishCommand -class _ConcurrentTelemetryScope(BaseMiddleware[KafkaPublishCommand]): - def __init__(self, scope: BaseMiddleware[KafkaPublishCommand]) -> None: +class _ConcurrentTelemetryScope(BaseMiddleware[PublishCommand]): + def __init__(self, scope: BaseMiddleware[PublishCommand]) -> None: super().__init__(scope.msg, context=scope.context) self._scope = scope self._consume_finished = False @@ -23,8 +23,8 @@ async def on_receive(self) -> None: async def publish_scope( self, - call_next: Callable[[KafkaPublishCommand], Awaitable[Any]], - cmd: KafkaPublishCommand, + call_next: Callable[[PublishCommand], Awaitable[Any]], + cmd: PublishCommand, ) -> Any: # noqa: ANN401 return await self._scope.publish_scope(call_next, cmd) @@ -53,7 +53,7 @@ async def after_processed( class _ConcurrentTelemetryFactory: - def __init__(self, factory: BrokerMiddleware[Any, KafkaPublishCommand]) -> None: + def __init__(self, factory: BrokerMiddleware[Any, PublishCommand]) -> None: self._factory = factory def __call__(self, msg: object, /, *, context: ContextRepo) -> _ConcurrentTelemetryScope: diff --git a/tests/test_middleware.py b/tests/test_middleware.py index 06fde48..e398e7e 100644 --- a/tests/test_middleware.py +++ b/tests/test_middleware.py @@ -274,7 +274,7 @@ async def raising_handle_task(coro: typing.Any, *_args: typing.Any) -> None: await test_broker.publish({"id": 1}, topic="cancel-topic") assert "Task cancelled during shutdown" in caplog.text - handler_instance.handle_task = original_handle_task # ty: ignore[invalid-assignment] + handler_instance.handle_task = original_handle_task await stop_concurrent_processing(test_broker.context) @@ -497,14 +497,14 @@ def mock_get(key: str, default: typing.Any = None) -> typing.Any: return mock_msg return original_get(key, default) - test_broker.context.get = mock_get # ty: ignore[invalid-assignment] + test_broker.context.get = mock_get middleware = KafkaConcurrentProcessingMiddleware(MagicMock(), context=test_broker.context) middleware.msg = ({"a": 1}, {"b": 2}) try: with pytest.raises(RuntimeError, match="does not support batch subscribers"): await middleware.consume_scope(AsyncMock(), mock_msg) finally: - test_broker.context.get = original_get # ty: ignore[invalid-assignment] + test_broker.context.get = original_get await asyncio.sleep(0) await stop_concurrent_processing(test_broker.context)