diff --git a/README.md b/README.md index cc17eb5..f60b530 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,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`. 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 new file mode 100644 index 0000000..84ec10a --- /dev/null +++ b/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) diff --git a/pyproject.toml b/pyproject.toml index 15ae31d..0a6a59d 100644 --- a/pyproject.toml +++ b/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", 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) diff --git a/tests/test_telemetry.py b/tests/test_telemetry.py new file mode 100644 index 0000000..7721b5f --- /dev/null +++ b/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()