fix: instrument every FastStream broker, with per-type middleware classes (#299)

Author: lesnik512

docs/integrations/faststream.md

@@ -51,3 +51,41 @@ application = bootstrapper.bootstrap()
 ```
 
 Read more about available configuration options [here](../introduction/configuration.md).
+
+## Multiple brokers
+
+Tracing, metrics and logging cover every broker the application has when `bootstrap()` runs. Add
+brokers with `add_broker` before calling `bootstrap()`: a broker added later is not instrumented.
+
+Each FastStream middleware class handles one broker type. A single class in
+`opentelemetry_middleware_cls` or `prometheus_middleware_cls` goes on every broker of the first
+broker's type. For an application with brokers of different types, pass a mapping from broker class
+to middleware class:
+
+```python
+from faststream.kafka import KafkaBroker
+from faststream.kafka.opentelemetry import KafkaTelemetryMiddleware
+from faststream.kafka.prometheus import KafkaPrometheusMiddleware
+from faststream.redis import RedisBroker
+from faststream.redis.opentelemetry import RedisTelemetryMiddleware
+from faststream.redis.prometheus import RedisPrometheusMiddleware
+
+
+application = AsgiFastStream(KafkaBroker(), RedisBroker())
+bootstrapper_config = FastStreamConfig(
+    service_name="microservice",
+    opentelemetry_endpoint="otl",
+    opentelemetry_middleware_cls={
+        KafkaBroker: KafkaTelemetryMiddleware,
+        RedisBroker: RedisTelemetryMiddleware,
+    },
+    prometheus_middleware_cls={
+        KafkaBroker: KafkaPrometheusMiddleware,
+        RedisBroker: RedisPrometheusMiddleware,
+    },
+    application=application,
+)
+```
+
+A broker whose type has no middleware class is left uninstrumented, and `bootstrap()` warns with its
+type name. All brokers share one Prometheus registry; each metric's `broker` label tells them apart.

docs/introduction/configuration.md

@@ -76,7 +76,9 @@ To bootstrap Prometheus for FastStream, you must provide at least one of:
 
 - `prometheus_middleware_cls` - the broker metrics middleware, e.g.
   `faststream.redis.prometheus.RedisPrometheusMiddleware`. It is constructed with the instrument's
-  registry and added to the broker, when the application has one.
+  registry and added to every broker of the first broker's type. For an application with brokers of
+  different types, pass a mapping from broker class to middleware class instead; see
+  [multiple brokers](../integrations/faststream.md#multiple-brokers).
 - `prometheus_collector_registry` - a `prometheus_client.CollectorRegistry` of your own, used in
   place of the fresh one the instrument would otherwise build.
 
@@ -183,8 +185,9 @@ For FastAPI and Litestar there is additionally:
 For FastStream there is additionally:
 
 - `opentelemetry_middleware_cls` - the broker telemetry middleware, e.g.
-  `faststream.redis.opentelemetry.RedisTelemetryMiddleware`. You must provide it to trace broker
-  messages. Without it the rest of the configuration above still applies: the exporter is built, the
+  `faststream.redis.opentelemetry.RedisTelemetryMiddleware`, or a mapping from broker class to
+  middleware class (see [multiple brokers](../integrations/faststream.md#multiple-brokers)). You must
+  provide it to trace broker messages. Without it the rest of the configuration above still applies: the exporter is built, the
   health-check span is recorded unless `opentelemetry_generate_health_check_spans` is False, and any
   `opentelemetry_instrumentors` are applied.
 

lite_bootstrap/bootstrappers/faststream_bootstrapper.py

@@ -3,9 +3,11 @@
 import json
 import logging
 import typing
+from collections.abc import Mapping, Sequence
 
 from lite_bootstrap import import_checker
 from lite_bootstrap.bootstrappers.base import BaseBootstrapper
+from lite_bootstrap.helpers.warn import warn_at_caller
 from lite_bootstrap.instruments.healthchecks_instrument import HealthChecksConfig, HealthChecksInstrument
 from lite_bootstrap.instruments.logging_instrument import LoggingConfig, LoggingInstrument
 from lite_bootstrap.instruments.opentelemetry_instrument import OpenTelemetryConfig, OpenTelemetryInstrument
@@ -65,6 +67,34 @@ def __init__(
     def __call__(self, msg: typing.Any, /, *, context: typing.Any) -> typing.Any: ...  # noqa: ANN401
 
 
+_MiddlewareT = typing.TypeVar("_MiddlewareT")
+
+
+def _middleware_targets(
+    brokers: Sequence[typing.Any],
+    configured: "type[_MiddlewareT] | Mapping[type, type[_MiddlewareT]]",
+    field_name: str,
+) -> list[tuple[typing.Any, "type[_MiddlewareT]"]]:
+    by_broker_type = configured if isinstance(configured, Mapping) else {type(brokers[0]): configured}
+    targets: list[tuple[typing.Any, type[_MiddlewareT]]] = []
+    skipped: dict[str, None] = {}
+    for broker in brokers:
+        middleware_cls = next(
+            (by_broker_type[one] for one in type(broker).__mro__ if one in by_broker_type),
+            None,
+        )
+        if middleware_cls is None:
+            skipped[repr(type(broker).__name__)] = None
+        else:
+            targets.append((broker, middleware_cls))
+    if skipped:
+        warn_at_caller(
+            f"FastStream brokers {', '.join(skipped)} are not instrumented: {field_name} has no middleware "
+            "class for their type. Pass a mapping from broker class to middleware class to cover them."
+        )
+    return targets
+
+
 def _make_asgi_faststream() -> "AsgiFastStream":
     return AsgiFastStream()
 
@@ -74,8 +104,14 @@ class FastStreamConfig(
     HealthChecksConfig, LoggingConfig, OpenTelemetryConfig, PrometheusConfig, PyroscopeConfig, SentryConfig
 ):
     application: "AsgiFastStream" = dataclasses.field(default_factory=_make_asgi_faststream)
-    opentelemetry_middleware_cls: type[FastStreamTelemetryMiddlewareProtocol] | None = None
-    prometheus_middleware_cls: type[FastStreamPrometheusMiddlewareProtocol] | None = None
+    opentelemetry_middleware_cls: (
+        type[FastStreamTelemetryMiddlewareProtocol] | Mapping[type, type[FastStreamTelemetryMiddlewareProtocol]] | None
+    ) = None
+    prometheus_middleware_cls: (
+        type[FastStreamPrometheusMiddlewareProtocol]
+        | Mapping[type, type[FastStreamPrometheusMiddlewareProtocol]]
+        | None
+    ) = None
     prometheus_collector_registry: "prometheus_client.CollectorRegistry | None" = None
     faststream_log_level: int = logging.WARNING
     faststream_health_check_broker_timeout: float = 5.0
@@ -112,27 +148,25 @@ async def _define_framework_health_status(self) -> bool:
 
 @dataclasses.dataclass(kw_only=True)
 class FastStreamLoggingInstrument(LoggingInstrument[FastStreamConfig]):
-    _prior_broker_params_storage: typing.Any = dataclasses.field(default=None, init=False, repr=False, compare=False)
-    _broker_logger_replaced: bool = dataclasses.field(default=False, init=False, repr=False, compare=False)
+    _prior_broker_params_storages: list[tuple[typing.Any, typing.Any]] = dataclasses.field(
+        default_factory=list, init=False, repr=False, compare=False
+    )
 
     def bootstrap(self) -> None:
         super().bootstrap()
-        broker = self.bootstrap_config.application.broker
-        if broker is not None and import_checker.is_structlog_installed and import_checker.is_faststream_installed:
+        brokers = self.bootstrap_config.application.brokers
+        if brokers and import_checker.is_structlog_installed and import_checker.is_faststream_installed:
             logger = structlog.get_logger("faststream")
             logger.setLevel(self.bootstrap_config.faststream_log_level)
-            self._prior_broker_params_storage = broker.config.logger.params_storage
-            broker.config.logger.params_storage = ManualLoggerStorage(logger)
-            self._broker_logger_replaced = True
+            for broker in brokers:
+                self._prior_broker_params_storages.append((broker, broker.config.logger.params_storage))
+                broker.config.logger.params_storage = ManualLoggerStorage(logger)
 
     def teardown(self) -> None:
         try:
-            if self._broker_logger_replaced:
-                broker = self.bootstrap_config.application.broker
-                if broker is not None:
-                    broker.config.logger.params_storage = self._prior_broker_params_storage
-                self._broker_logger_replaced = False
-                self._prior_broker_params_storage = None
+            while self._prior_broker_params_storages:
+                broker, prior_storage = self._prior_broker_params_storages.pop()
+                broker.config.logger.params_storage = prior_storage
         finally:
             super().teardown()
 
@@ -142,12 +176,13 @@ class FastStreamOpenTelemetryInstrument(OpenTelemetryInstrument[FastStreamConfig
     def bootstrap(self) -> None:
         super().bootstrap()
         config = self.bootstrap_config
-        if config.opentelemetry_middleware_cls and config.application.broker:
-            config.application.broker.add_middleware(
-                config.opentelemetry_middleware_cls(
-                    tracer_provider=get_tracer_provider(), meter_provider=get_meter_provider()
+        if config.opentelemetry_middleware_cls and config.application.brokers:
+            for broker, middleware_cls in _middleware_targets(
+                config.application.brokers, config.opentelemetry_middleware_cls, "opentelemetry_middleware_cls"
+            ):
+                broker.add_middleware(
+                    middleware_cls(tracer_provider=get_tracer_provider(), meter_provider=get_meter_provider())
                 )
-            )
 
 
 def _make_collector_registry() -> "prometheus_client.CollectorRegistry":
@@ -183,8 +218,11 @@ def bootstrap(self) -> None:
         config.application.mount(
             config.prometheus_metrics_path, prometheus_client.make_asgi_app(self.collector_registry)
         )
-        if config.prometheus_middleware_cls and config.application.broker:
-            config.application.broker.add_middleware(config.prometheus_middleware_cls(registry=self.collector_registry))
+        if config.prometheus_middleware_cls and config.application.brokers:
+            for broker, middleware_cls in _middleware_targets(
+                config.application.brokers, config.prometheus_middleware_cls, "prometheus_middleware_cls"
+            ):
+                broker.add_middleware(middleware_cls(registry=self.collector_registry))
 
 
 class FastStreamBootstrapper(BaseBootstrapper["AsgiFastStream"]):

tests/test_faststream_bootstrap.py

@@ -511,3 +511,107 @@ async def test_faststream_health_check_fails_on_broker_ping_despite_passing_prob
 
     assert response.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR
     assert response.json()["health_status"] is False
+
+
+class OtherRedisBroker(RedisBroker): ...
+
+
+class OtherTelemetryMiddleware(RedisTelemetryMiddleware): ...
+
+
+class OtherPrometheusMiddleware(RedisPrometheusMiddleware): ...
+
+
+def _middleware_types(broker: BrokerUsecase[typing.Any, typing.Any, typing.Any]) -> list[type]:
+    return [type(middleware) for middleware in broker.config.broker_middlewares]
+
+
+def _bootstrap_with_brokers(
+    brokers: list[BrokerUsecase[typing.Any, typing.Any, typing.Any]], **overrides: object
+) -> FastStreamBootstrapper:
+    config = dataclasses.replace(build_faststream_config(broker=brokers[0]), **overrides)
+    for extra in brokers[1:]:
+        config.application.add_broker(extra)
+    bootstrapper = FastStreamBootstrapper(bootstrap_config=config)
+    bootstrapper.bootstrap()
+    return bootstrapper
+
+
+def test_faststream_middlewares_are_added_to_every_broker(
+    broker: RedisBroker, recwarn: pytest.WarningsRecorder
+) -> None:
+    second_broker = RedisBroker()
+    bootstrapper = _bootstrap_with_brokers([broker, second_broker])
+    try:
+        for one in (broker, second_broker):
+            assert _middleware_types(one) == [RedisTelemetryMiddleware, RedisPrometheusMiddleware]
+        assert not [one for one in recwarn if "not instrumented" in str(one.message)]
+    finally:
+        bootstrapper.teardown()
+
+
+def test_faststream_single_middleware_class_skips_brokers_of_another_type_with_a_warning(broker: RedisBroker) -> None:
+    first_broker = OtherRedisBroker()
+    with pytest.warns(UserWarning, match=r"not instrumented") as recorded:
+        bootstrapper = _bootstrap_with_brokers(
+            [first_broker, broker],
+            opentelemetry_middleware_cls=OtherTelemetryMiddleware,
+            prometheus_middleware_cls=OtherPrometheusMiddleware,
+        )
+    try:
+        assert _middleware_types(first_broker) == [OtherTelemetryMiddleware, OtherPrometheusMiddleware]
+        assert _middleware_types(broker) == []
+        skipped_messages = [str(one.message) for one in recorded if "not instrumented" in str(one.message)]
+        expected_warning_count = 2
+        assert len(skipped_messages) == expected_warning_count
+        assert all("'RedisBroker'" in message for message in skipped_messages)
+    finally:
+        bootstrapper.teardown()
+
+
+def test_faststream_middleware_mapping_picks_the_most_specific_class_per_broker(broker: RedisBroker) -> None:
+    other_broker = OtherRedisBroker()
+    bootstrapper = _bootstrap_with_brokers(
+        [broker, other_broker],
+        opentelemetry_middleware_cls={
+            RedisBroker: RedisTelemetryMiddleware,
+            OtherRedisBroker: OtherTelemetryMiddleware,
+        },
+        prometheus_middleware_cls={
+            RedisBroker: RedisPrometheusMiddleware,
+            OtherRedisBroker: OtherPrometheusMiddleware,
+        },
+    )
+    try:
+        assert _middleware_types(broker) == [RedisTelemetryMiddleware, RedisPrometheusMiddleware]
+        assert _middleware_types(other_broker) == [OtherTelemetryMiddleware, OtherPrometheusMiddleware]
+    finally:
+        bootstrapper.teardown()
+
+
+def test_faststream_middleware_mapping_without_an_entry_skips_the_broker_with_a_warning(broker: RedisBroker) -> None:
+    other_broker = OtherRedisBroker()
+    with pytest.warns(UserWarning, match=r"'RedisBroker'.*not instrumented"):
+        bootstrapper = _bootstrap_with_brokers(
+            [other_broker, broker],
+            opentelemetry_middleware_cls={OtherRedisBroker: OtherTelemetryMiddleware},
+            prometheus_middleware_cls=None,
+        )
+    try:
+        assert _middleware_types(other_broker) == [OtherTelemetryMiddleware]
+        assert _middleware_types(broker) == []
+    finally:
+        bootstrapper.teardown()
+
+
+def test_faststream_logging_covers_every_broker_and_teardown_restores_each(broker: RedisBroker) -> None:
+    second_broker = RedisBroker()
+    original_storages = [broker.config.logger.params_storage, second_broker.config.logger.params_storage]
+    bootstrapper = _bootstrap_with_brokers([broker, second_broker])
+    try:
+        for one in (broker, second_broker):
+            assert isinstance(one.config.logger.params_storage, ManualLoggerStorage)
+    finally:
+        bootstrapper.teardown()
+
+    assert [broker.config.logger.params_storage, second_broker.config.logger.params_storage] == original_storages