Add OTel baggage support (airtai#1692)

* feat: add current span in faststream context * feat: add custom baggage * test: add baggage test * feat: refactor baggage * test: fix confluent tests * docs: generate API References * feat: refactor * fix: linters * fix: mypy * docs: add baggage and CurrentSpan docs * docs: fix trailing whitespaces --------- Co-authored-by: draincoder <draincoder@users.noreply.github.com>
IvanKirpichnikov · Sep 9, 2024 · e2e3023 · e2e3023
1 parent 7aaafdd
commit e2e3023
Show file tree

Hide file tree

Showing 12 changed files with 480 additions and 23 deletions.
diff --git a/docs/docs/SUMMARY.md b/docs/docs/SUMMARY.md
@@ -177,6 +177,7 @@ search:
                 - [TestApp](public_api/faststream/nats/TestApp.md)
                 - [TestNatsBroker](public_api/faststream/nats/TestNatsBroker.md)
             - opentelemetry
+                - [Baggage](public_api/faststream/opentelemetry/Baggage.md)
                 - [TelemetryMiddleware](public_api/faststream/opentelemetry/TelemetryMiddleware.md)
                 - [TelemetrySettingsProvider](public_api/faststream/opentelemetry/TelemetrySettingsProvider.md)
             - rabbit
@@ -795,8 +796,11 @@ search:
                     - [TestNatsBroker](api/faststream/nats/testing/TestNatsBroker.md)
                     - [build_message](api/faststream/nats/testing/build_message.md)
             - opentelemetry
+                - [Baggage](api/faststream/opentelemetry/Baggage.md)
                 - [TelemetryMiddleware](api/faststream/opentelemetry/TelemetryMiddleware.md)
                 - [TelemetrySettingsProvider](api/faststream/opentelemetry/TelemetrySettingsProvider.md)
+                - baggage
+                    - [Baggage](api/faststream/opentelemetry/baggage/Baggage.md)
                 - consts
                     - [MessageAction](api/faststream/opentelemetry/consts/MessageAction.md)
                 - middleware

diff --git a/docs/docs/en/api/faststream/opentelemetry/Baggage.md b/docs/docs/en/api/faststream/opentelemetry/Baggage.md
@@ -0,0 +1,11 @@
+---
+# 0.5 - API
+# 2 - Release
+# 3 - Contributing
+# 5 - Template Page
+# 10 - Default
+search:
+  boost: 0.5
+---
+
+::: faststream.opentelemetry.Baggage
diff --git a/docs/docs/en/api/faststream/opentelemetry/baggage/Baggage.md b/docs/docs/en/api/faststream/opentelemetry/baggage/Baggage.md
@@ -0,0 +1,11 @@
+---
+# 0.5 - API
+# 2 - Release
+# 3 - Contributing
+# 5 - Template Page
+# 10 - Default
+search:
+  boost: 0.5
+---
+
+::: faststream.opentelemetry.baggage.Baggage
diff --git a/docs/docs/en/getting-started/opentelemetry/index.md b/docs/docs/en/getting-started/opentelemetry/index.md
@@ -41,7 +41,7 @@ async def third_handler(msg: str):
     await asyncio.sleep(0.075)
 ```
 
-## FastStream Tracing
+### FastStream Tracing
 
 **OpenTelemetry** tracing support in **FastStream** adheres to the [semantic conventions for messaging systems](https://opentelemetry.io/docs/specs/semconv/messaging/){.external-link target="_blank"}.
 
@@ -98,7 +98,22 @@ To visualize traces, you can send them to a backend system that supports distrib
 * **Zipkin**: Similar to **Jaeger**, you can run **Zipkin** using **Docker** and configure the **OpenTelemetry** middleware accordingly. For more details, see the [Zipkin documentation](https://zipkin.io/){.external-link target="_blank"}.
 * **Grafana Tempo**: **Grafana Tempo** is a high-scale distributed tracing backend. You can configure **OpenTelemetry** to export traces to **Tempo**, which can then be visualized using **Grafana**. For more details, see the [Grafana Tempo documentation](https://grafana.com/docs/tempo/latest/){.external-link target="_blank"}.
 
-## Example
+### Context propagation
+
+Quite often it is necessary to communicate with **other** services and to propagate the trace context, you can use the **CurrentSpan** object and follow the example:
+
+```python linenums="1" hl_lines="5-7"
+from opentelemetry import trace, propagate
+from faststream.opentelemetry import CurrentSpan
+
+@broker.subscriber("symbol")
+async def handler(msg: str, span: CurrentSpan) -> None:
+    headers = {}
+    propagate.inject(headers, context=trace.set_span_in_context(span))
+    price = await exchange_client.get_symbol_price(msg, headers=headers)
+```
+
+### Full example
 
 To see how to set up, visualize, and configure tracing for **FastStream** services, go to [example](https://github.com/draincoder/faststream-monitoring){.external-link target="_blank"}.
 
@@ -112,3 +127,60 @@ An example includes:
 
 ![HTML-page](../../../assets/img/distributed-trace.png){ loading=lazy }
 `Visualized via Grafana and Tempo`
+
+## Baggage
+
+[OpenTelemetry Baggage](https://opentelemetry.io/docs/concepts/signals/baggage/){.external-link target="_blank"} is a context propagation mechanism that allows you to pass custom metadata or key-value pairs across service boundaries, providing additional context for distributed tracing and observability.
+
+### FastStream Baggage
+
+**FastStream** provides a convenient abstraction over baggage that allows you to:
+
+* **Initialize** the baggage
+* **Propagate** baggage through headers
+* **Modify** the baggage
+* **Stop** propagating baggage
+
+### Example
+
+To initialize the baggage and start distributing it, follow this example:
+
+```python linenums="1" hl_lines="3-4"
+from faststream.opentelemetry import Baggage
+
+headers = Baggage({"hello": "world"}).to_headers({"header-type": "custom"})
+await broker.publish("hello", "first", headers=headers)
+```
+
+All interaction with baggage at the **consumption level** occurs through the **CurrentBaggage** object, which is automatically substituted from the context:
+
+```python linenums="1" hl_lines="6-10 17-18 24"
+from faststream.opentelemetry import CurrentBaggage
+
+@broker.subscriber("first")
+@broker.publisher("second")
+async def response_handler_first(msg: str, baggage: CurrentBaggage):
+    print(baggage.get_all())  # {'hello': 'world'}
+    baggage.remove("hello")
+    baggage.set("user-id", 1)
+    baggage.set("request-id", "UUID")
+    print(baggage.get("user-id"))  # 1
+    return msg
+
+
+@broker.subscriber("second")
+@broker.publisher("third")
+async def response_handler_second(msg: str, baggage: CurrentBaggage):
+    print(baggage.get_all())  # {'user-id': '1', 'request-id': 'UUID'}
+    baggage.clear()
+    return msg
+
+
+@broker.subscriber("third")
+async def response_handler_third(msg: str, baggage: CurrentBaggage):
+    print(baggage.get_all())  # {}
+```
+
+!!! note
+    If you consume messages in **batches**, then the baggage from each message will be merged into the **common baggage** available
+    through the `get_all` method, but you can still get a list of all the baggage from the batch using the `get_all_batch` method.
diff --git a/faststream/opentelemetry/__init__.py b/faststream/opentelemetry/__init__.py
@@ -1,7 +1,12 @@
+from faststream.opentelemetry.annotations import CurrentBaggage, CurrentSpan
+from faststream.opentelemetry.baggage import Baggage
 from faststream.opentelemetry.middleware import TelemetryMiddleware
 from faststream.opentelemetry.provider import TelemetrySettingsProvider
 
 __all__ = (
+    "Baggage",
+    "CurrentBaggage",
+    "CurrentSpan",
     "TelemetryMiddleware",
     "TelemetrySettingsProvider",
 )
diff --git a/faststream/opentelemetry/annotations.py b/faststream/opentelemetry/annotations.py
@@ -0,0 +1,8 @@
+from opentelemetry.trace import Span
+from typing_extensions import Annotated
+
+from faststream import Context
+from faststream.opentelemetry.baggage import Baggage
+
+CurrentSpan = Annotated[Span, Context("span")]
+CurrentBaggage = Annotated[Baggage, Context("baggage")]
diff --git a/faststream/opentelemetry/baggage.py b/faststream/opentelemetry/baggage.py
@@ -0,0 +1,77 @@
+from typing import TYPE_CHECKING, Any, List, Optional, cast
+
+from opentelemetry import baggage, context
+from opentelemetry.baggage.propagation import W3CBaggagePropagator
+from typing_extensions import Self
+
+if TYPE_CHECKING:
+    from faststream.broker.message import StreamMessage
+    from faststream.types import AnyDict
+
+_BAGGAGE_PROPAGATOR = W3CBaggagePropagator()
+
+
+class Baggage:
+    __slots__ = ("_baggage", "_batch_baggage")
+
+    def __init__(
+        self, payload: "AnyDict", batch_payload: Optional[List["AnyDict"]] = None
+    ) -> None:
+        self._baggage = dict(payload)
+        self._batch_baggage = [dict(b) for b in batch_payload] if batch_payload else []
+
+    def get_all(self) -> "AnyDict":
+        """Get a copy of the current baggage."""
+        return self._baggage.copy()
+
+    def get_all_batch(self) -> List["AnyDict"]:
+        """Get a copy of all batch baggage if exists."""
+        return self._batch_baggage.copy()
+
+    def get(self, key: str) -> Optional[Any]:
+        """Get a value from the baggage by key."""
+        return self._baggage.get(key)
+
+    def remove(self, key: str) -> None:
+        """Remove a baggage item by key."""
+        self._baggage.pop(key, None)
+
+    def set(self, key: str, value: Any) -> None:
+        """Set a key-value pair in the baggage."""
+        self._baggage[key] = value
+
+    def clear(self) -> None:
+        """Clear the current baggage."""
+        self._baggage.clear()
+
+    def to_headers(self, headers: Optional["AnyDict"] = None) -> "AnyDict":
+        """Convert baggage items to headers format for propagation."""
+        current_context = context.get_current()
+        if headers is None:
+            headers = {}
+
+        for k, v in self._baggage.items():
+            current_context = baggage.set_baggage(k, v, context=current_context)
+
+        _BAGGAGE_PROPAGATOR.inject(headers, current_context)
+        return headers
+
+    @classmethod
+    def from_msg(cls, msg: "StreamMessage[Any]") -> Self:
+        """Create a Baggage instance from a StreamMessage."""
+        if len(msg.batch_headers) <= 1:
+            payload = baggage.get_all(_BAGGAGE_PROPAGATOR.extract(msg.headers))
+            return cls(cast("AnyDict", payload))
+
+        cumulative_baggage: AnyDict = {}
+        batch_baggage: List[AnyDict] = []
+
+        for headers in msg.batch_headers:
+            payload = baggage.get_all(_BAGGAGE_PROPAGATOR.extract(headers))
+            cumulative_baggage.update(payload)
+            batch_baggage.append(cast("AnyDict", payload))
+
+        return cls(cumulative_baggage, batch_baggage)
+
+    def __repr__(self) -> str:
+        return self._baggage.__repr__()
diff --git a/faststream/opentelemetry/middleware.py b/faststream/opentelemetry/middleware.py
@@ -1,7 +1,7 @@
 import time
 from collections import defaultdict
 from copy import copy
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type, cast
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Type, cast
 
 from opentelemetry import baggage, context, metrics, trace
 from opentelemetry.baggage.propagation import W3CBaggagePropagator
@@ -11,6 +11,8 @@
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 
 from faststream import BaseMiddleware
+from faststream import context as fs_context
+from faststream.opentelemetry.baggage import Baggage
 from faststream.opentelemetry.consts import (
     ERROR_TYPE,
     MESSAGING_DESTINATION_PUBLISH_NAME,
@@ -21,6 +23,7 @@
 from faststream.opentelemetry.provider import TelemetrySettingsProvider
 
 if TYPE_CHECKING:
+    from contextvars import Token
     from types import TracebackType
 
     from opentelemetry.metrics import Meter, MeterProvider
@@ -118,6 +121,7 @@ def __init__(
         self._metrics = metrics_container
         self._current_span: Optional[Span] = None
         self._origin_context: Optional[Context] = None
+        self._scope_tokens: List[Tuple[str, Token[Any]]] = []
         self.__settings_provider = settings_provider_factory(msg)
 
     async def publish_scope(
@@ -134,6 +138,10 @@ async def publish_scope(
         current_context = context.get_current()
         destination_name = provider.get_publish_destination_name(kwargs)
 
+        current_baggage: Optional[Baggage] = fs_context.get_local("baggage")
+        if current_baggage:
+            headers.update(current_baggage.to_headers())
+
         trace_attributes = provider.get_publish_attrs_from_kwargs(kwargs)
         metrics_attributes = {
             SpanAttributes.MESSAGING_SYSTEM: provider.messaging_system,
@@ -143,6 +151,7 @@ async def publish_scope(
         # NOTE: if batch with single message?
         if (msg_count := len((msg, *args))) > 1:
             trace_attributes[SpanAttributes.MESSAGING_BATCH_MESSAGE_COUNT] = msg_count
+            current_context = _BAGGAGE_PROPAGATOR.extract(headers, current_context)
             _BAGGAGE_PROPAGATOR.inject(
                 headers, baggage.set_baggage(WITH_BATCH, True, context=current_context)
             )
@@ -185,6 +194,9 @@ async def publish_scope(
             duration = time.perf_counter() - start_time
             self._metrics.observe_publish(metrics_attributes, duration, msg_count)
 
+        for key, token in self._scope_tokens:
+            fs_context.reset_local(key, token)
+
         return result
 
     async def consume_scope(
@@ -234,6 +246,12 @@ async def consume_scope(
                     SpanAttributes.MESSAGING_OPERATION, MessageAction.PROCESS
                 )
                 self._current_span = span
+
+                self._scope_tokens.append(("span", fs_context.set_local("span", span)))
+                self._scope_tokens.append(
+                    ("baggage", fs_context.set_local("baggage", Baggage.from_msg(msg)))
+                )
+
                 new_context = trace.set_span_in_context(span, current_context)
                 token = context.attach(new_context)
                 result = await call_next(msg)