refactor: refactor Baggage with Context interaction (#2748)

Co-authored-by: Zhongyang Wu <[email protected]>
Co-authored-by: Cijo Thomas <[email protected]>

Author: 3 people

opentelemetry-sdk/src/propagation/baggage.rs

@@ -30,7 +30,7 @@ fn baggage_fields() -> &'static [String; 1] {
 /// # Examples
 ///
 /// ```
-/// use opentelemetry::{baggage::BaggageExt, KeyValue, propagation::TextMapPropagator};
+/// use opentelemetry::{baggage::{Baggage, BaggageExt}, propagation::TextMapPropagator};
 /// use opentelemetry_sdk::propagation::BaggagePropagator;
 /// use std::collections::HashMap;
 ///
@@ -48,14 +48,17 @@ fn baggage_fields() -> &'static [String; 1] {
 /// }
 ///
 /// // Add new baggage
-/// let cx_with_additions = cx.with_baggage(vec![KeyValue::new("server_id", 42)]);
+/// let mut baggage = Baggage::new();
+/// let _ = baggage.insert("server_id", "42");
+///
+/// let cx_with_additions = cx.with_baggage(baggage);
 ///
 /// // Inject baggage into http request
 /// propagator.inject_context(&cx_with_additions, &mut headers);
 ///
 /// let header_value = headers.get("baggage").expect("header is injected");
-/// assert!(header_value.contains("user_id=1"), "still contains previous name-value");
-/// assert!(header_value.contains("server_id=42"), "contains new name-value pair");
+/// assert!(!header_value.contains("user_id=1"), "still contains previous name-value");
+/// assert!(header_value.contains("server_id=42"), "does not contain new name-value pair");
 /// ```
 ///
 /// [W3C Baggage]: https://w3c.github.io/baggage
@@ -98,7 +101,7 @@ impl TextMapPropagator for BaggagePropagator {
     /// Extracts a `Context` with baggage values from a `Extractor`.
     fn extract_with_context(&self, cx: &Context, extractor: &dyn Extractor) -> Context {
         if let Some(header_value) = extractor.get(BAGGAGE_HEADER) {
-            let baggage = header_value.split(',').flat_map(|context_value| {
+            let baggage = header_value.split(',').filter_map(|context_value| {
                 if let Some((name_and_value, props)) = context_value
                     .split(';')
                     .collect::<Vec<&str>>()

opentelemetry/CHANGELOG.md

@@ -7,11 +7,13 @@
 - *Breaking* Moved `TraceResult` type alias from `opentelemetry::trace::TraceResult` to `opentelemetry_sdk::trace::TraceResult`
 - {PLACEHOLDER} - Remove the above completely. // TODO fill this when changes are actually in.
 - Bug Fix: `InstrumentationScope` implementation for `PartialEq` and `Hash` fixed to include Attributes also.
-- *Breaking* Changed value type of `Baggage` from `Value` to `StringValue`
-- Updated `Baggage` constants to reflect latest standard (`MAX_KEY_VALUE_PAIRS` - 180 -> 64, `MAX_BYTES_FOR_ONE_PAIR` - removed) and increased insert performance see #[2284](https://github.com/open-telemetry/opentelemetry-rust/pull/2284).
-- *Breaking* Align `Baggage.remove()` signature with `.get()` to take the key as a reference
+- **Breaking changes for baggage users**: [#2717](https://github.com/open-telemetry/opentelemetry-rust/issues/2717)
+  - Changed value type of `Baggage` from `Value` to `StringValue`
+  - Updated `Baggage` constants to reflect latest standard (`MAX_KEY_VALUE_PAIRS` - 180 -> 64, `MAX_BYTES_FOR_ONE_PAIR` - removed) and increased insert performance see #[2284](https://github.com/open-telemetry/opentelemetry-rust/pull/2284).
+  - Align `Baggage.remove()` signature with `.get()` to take the key as a reference
+  - `Baggage` can't be retrieved from the `Context` directly anymore and needs to be accessed via `context.baggage()`
+  - `with_baggage()` and `current_with_baggage()` override any existing `Baggage` in the `Context`
 - Changed `Context` to use a stack to properly handle out of order dropping of `ContextGuard`. This imposes a limit of `65535` nested contexts on a single thread. See #[2378](https://github.com/open-telemetry/opentelemetry-rust/pull/2284) and #[1887](https://github.com/open-telemetry/opentelemetry-rust/issues/1887).
-
 - Added additional `name: Option<&str>` parameter to the `event_enabled` method
   on the `Logger` trait. This allows implementations (SDK, processor, exporters)
   to leverage this additional information to determine if an event is enabled.

opentelemetry/src/baggage.rs

@@ -273,6 +273,16 @@ impl FromIterator<KeyValueMetadata> for Baggage {
     }
 }
 
+impl<I> From<I> for Baggage
+where
+    I: IntoIterator,
+    I::Item: Into<KeyValueMetadata>,
+{
+    fn from(value: I) -> Self {
+        value.into_iter().map(Into::into).collect()
+    }
+}
+
 fn encode(s: &str) -> String {
     let mut encoded_string = String::with_capacity(s.len());
 
@@ -312,39 +322,46 @@ pub trait BaggageExt {
     /// # Examples
     ///
     /// ```
-    /// use opentelemetry::{baggage::BaggageExt, Context, KeyValue, StringValue};
+    /// use opentelemetry::{baggage::{Baggage, BaggageExt}, Context, KeyValue, StringValue};
+    ///
+    /// // Explicit `Baggage` creation
+    /// let mut baggage = Baggage::new();
+    /// let _ = baggage.insert("my-name", "my-value");
+    ///
+    /// let cx = Context::map_current(|cx| {
+    ///     cx.with_baggage(baggage)
+    /// });
     ///
+    /// // Passing an iterator
     /// let cx = Context::map_current(|cx| {
-    ///     cx.with_baggage(vec![KeyValue::new("my-name", "my-value")])
+    ///     cx.with_baggage([KeyValue::new("my-name", "my-value")])
     /// });
     ///
     /// assert_eq!(
     ///     cx.baggage().get("my-name"),
     ///     Some(&StringValue::from("my-value")),
     /// )
     /// ```
-    fn with_baggage<T: IntoIterator<Item = I>, I: Into<KeyValueMetadata>>(
-        &self,
-        baggage: T,
-    ) -> Self;
+    fn with_baggage<T: Into<Baggage>>(&self, baggage: T) -> Self;
 
     /// Returns a clone of the current context with the included name/value pairs.
     ///
     /// # Examples
     ///
     /// ```
-    /// use opentelemetry::{baggage::BaggageExt, Context, KeyValue, StringValue};
+    /// use opentelemetry::{baggage::{Baggage, BaggageExt}, Context, StringValue};
     ///
-    /// let cx = Context::current_with_baggage(vec![KeyValue::new("my-name", "my-value")]);
+    /// let mut baggage = Baggage::new();
+    /// let _ = baggage.insert("my-name", "my-value");
+    ///
+    /// let cx = Context::current_with_baggage(baggage);
     ///
     /// assert_eq!(
     ///     cx.baggage().get("my-name"),
     ///     Some(&StringValue::from("my-value")),
     /// )
     /// ```
-    fn current_with_baggage<T: IntoIterator<Item = I>, I: Into<KeyValueMetadata>>(
-        baggage: T,
-    ) -> Self;
+    fn current_with_baggage<T: Into<Baggage>>(baggage: T) -> Self;
 
     /// Returns a clone of the given context with no baggage.
     ///
@@ -364,33 +381,26 @@ pub trait BaggageExt {
     fn baggage(&self) -> &Baggage;
 }
 
-impl BaggageExt for Context {
-    fn with_baggage<T: IntoIterator<Item = I>, I: Into<KeyValueMetadata>>(
-        &self,
-        baggage: T,
-    ) -> Self {
-        let old = self.baggage();
-        let mut merged = Baggage {
-            inner: old.inner.clone(),
-            kv_content_len: old.kv_content_len,
-        };
-        for kvm in baggage.into_iter().map(|kv| kv.into()) {
-            merged.insert_with_metadata(kvm.key, kvm.value, kvm.metadata);
-        }
+/// Solely used to store `Baggage` in the `Context` without allowing direct access
+#[derive(Debug)]
+struct BaggageContextValue(Baggage);
 
-        self.with_value(merged)
+impl BaggageExt for Context {
+    fn with_baggage<T: Into<Baggage>>(&self, baggage: T) -> Self {
+        self.with_value(BaggageContextValue(baggage.into()))
     }
 
-    fn current_with_baggage<T: IntoIterator<Item = I>, I: Into<KeyValueMetadata>>(kvs: T) -> Self {
-        Context::map_current(|cx| cx.with_baggage(kvs))
+    fn current_with_baggage<T: Into<Baggage>>(baggage: T) -> Self {
+        Context::map_current(|cx| cx.with_baggage(baggage))
     }
 
     fn with_cleared_baggage(&self) -> Self {
         self.with_value(Baggage::new())
     }
 
     fn baggage(&self) -> &Baggage {
-        self.get::<Baggage>().unwrap_or(get_default_baggage())
+        self.get::<BaggageContextValue>()
+            .map_or(get_default_baggage(), |b| &b.0)
     }
 }