fix: preserve Kafka telemetry spans during concurrent processing

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,14 @@ modern_di_faststream.setup_di(app, container=container)  # registered after →
 
 ## How it works
 
+### OpenTelemetry
+
+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
 
 On each incoming message, `consume_scope` calls `handle_task()`, which acquires a semaphore slot then fires the handler coroutine as a background `asyncio.Task`.

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)

faststream_concurrent_aiokafka/telemetry.py

@@ -0,0 +1,76 @@
+import sys
+from collections.abc import Awaitable, Callable
+from types import TracebackType
+from typing import Any
+
+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.message import StreamMessage
+from faststream.response import PublishCommand
+
+
+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
+        self._cleanup_requested = False
+
+    async def on_receive(self) -> None:
+        await self._scope.on_receive()
+
+    async def publish_scope(
+        self,
+        call_next: Callable[[PublishCommand], Awaitable[Any]],
+        cmd: PublishCommand,
+    ) -> Any:  # noqa: ANN401
+        return await self._scope.publish_scope(call_next, cmd)
+
+    async def consume_scope(
+        self,
+        call_next: Callable[[StreamMessage[Any]], Awaitable[Any]],
+        msg: StreamMessage[Any],
+    ) -> Any:  # noqa: ANN401
+        try:
+            return await self._scope.consume_scope(call_next, msg)
+        finally:
+            self._consume_finished = True
+            if self._cleanup_requested:
+                await self._scope.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 self._scope.after_processed(exc_type, exc_val, exc_tb)
+        return False
+
+
+class _ConcurrentTelemetryFactory:
+    def __init__(self, factory: BrokerMiddleware[Any, PublishCommand]) -> None:
+        self._factory = factory
+
+    def __call__(self, msg: object, /, *, context: ContextRepo) -> _ConcurrentTelemetryScope:
+        return _ConcurrentTelemetryScope(self._factory(msg, 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)

pyproject.toml

@@ -42,6 +42,7 @@ Changelog = "https://github.com/modern-python/faststream-concurrent-aiokafka/rel
 
 [dependency-groups]
 dev = [
+    "faststream[otel]",
     "pytest>=8",
     "pytest-cov>=5",
     "pytest-xdist>=3",

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)

tests/test_telemetry.py

@@ -0,0 +1,108 @@
+import asyncio
+import time
+from collections.abc import Sequence
+from uuid import uuid4
+
+import pytest
+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
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExportResult
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+
+from faststream_concurrent_aiokafka import (
+    KafkaConcurrentProcessingMiddleware,
+    initialize_concurrent_processing,
+    stop_concurrent_processing,
+)
+
+
+class _ProcessSpanExporter(InMemorySpanExporter):
+    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.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(topic)
+    provider = TracerProvider()
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
+    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
+
+    @registrator.subscriber(
+        topic,
+        group_id=topic,
+        auto_offset_reset="earliest",
+        ack_policy=AckPolicy.ACK if outcome == "inline" else AckPolicy.MANUAL,
+    )
+    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:
+            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)
+                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.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)
+        if outcome == "cancel":
+            assert handler_task.cancelled()
+    finally:
+        provider.shutdown()