refactor Baggage with Context interaction

Author: gruebel

opentelemetry-sdk/src/propagation/baggage.rs

@@ -1,5 +1,5 @@
 use opentelemetry::{
-    baggage::{BaggageExt, KeyValueMetadata},
+    baggage::{Baggage, BaggageExt, KeyValueMetadata},
     otel_warn,
     propagation::{text_map_propagator::FieldIter, Extractor, Injector, TextMapPropagator},
     Context,
@@ -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>>()
@@ -148,7 +151,7 @@ impl TextMapPropagator for BaggagePropagator {
                     None
                 }
             });
-            cx.with_baggage(baggage)
+            cx.with_baggage(Baggage::from_iter(baggage))
         } else {
             cx.clone()
         }
@@ -275,7 +278,7 @@ mod tests {
 
         for (kvm, header_parts) in valid_inject_data() {
             let mut injector = HashMap::new();
-            let cx = Context::current_with_baggage(kvm);
+            let cx = Context::current_with_baggage(Baggage::from_iter(kvm));
             propagator.inject_context(&cx, &mut injector);
             let header_value = injector.get(BAGGAGE_HEADER).unwrap();
             assert_eq!(header_parts.join(",").len(), header_value.len(),);
@@ -307,7 +310,7 @@ mod tests {
 
         for (kvm, header_parts) in valid_inject_data_metadata() {
             let mut injector = HashMap::new();
-            let cx = Context::current_with_baggage(kvm);
+            let cx = Context::current_with_baggage(Baggage::from_iter(kvm));
             propagator.inject_context(&cx, &mut injector);
             let header_value = injector.get(BAGGAGE_HEADER).unwrap();
 

opentelemetry/CHANGELOG.md

@@ -7,9 +7,12 @@
 - *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()` only accept a `Baggage` instance and override any existing `Baggage` in the `Context`
 
 ## 0.28.0
 

opentelemetry/src/baggage.rs

@@ -308,39 +308,40 @@ pub trait BaggageExt {
     /// # Examples
     ///
     /// ```
-    /// use opentelemetry::{baggage::BaggageExt, Context, KeyValue, StringValue};
+    /// use opentelemetry::{baggage::{Baggage, BaggageExt}, Context, StringValue};
+    ///
+    /// let mut baggage = Baggage::new();
+    /// let _ = baggage.insert("my-name", "my-value");
     ///
     /// let cx = Context::map_current(|cx| {
-    ///     cx.with_baggage(vec![KeyValue::new("my-name", "my-value")])
+    ///     cx.with_baggage(baggage)
     /// });
     ///
     /// 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(&self, baggage: Baggage) -> 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(baggage: Baggage) -> Self;
 
     /// Returns a clone of the given context with no baggage.
     ///
@@ -360,33 +361,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(&self, baggage: Baggage) -> Self {
+        self.with_value(BaggageContextValue(baggage))
     }
 
-    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(baggage: Baggage) -> 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)
     }
 }
 

opentelemetry/src/propagation/composite.rs

@@ -23,7 +23,7 @@ use std::collections::HashSet;
 ///
 /// ```
 /// use opentelemetry::{
-///     baggage::BaggageExt,
+///     baggage::{Baggage, BaggageExt},
 ///     propagation::{TextMapPropagator, TextMapCompositePropagator},
 ///
 ///     trace::{TraceContextExt, Tracer, TracerProvider},
@@ -56,7 +56,7 @@ use std::collections::HashSet;
 /// // with the current context, call inject to add the headers
 /// composite_propagator.inject_context(
 ///     &Context::current_with_span(example_span)
-///         .with_baggage(vec![KeyValue::new("test", "example")]),
+///         .with_baggage(Baggage::from_iter([KeyValue::new("test", "example")])),
 ///     &mut injector,
 /// );
 ///
@@ -115,7 +115,7 @@ impl TextMapPropagator for TextMapCompositePropagator {
 
 #[cfg(all(test, feature = "trace"))]
 mod tests {
-    use crate::baggage::BaggageExt;
+    use crate::baggage::{Baggage, BaggageExt};
     use crate::propagation::TextMapCompositePropagator;
     use crate::testing::trace::TestSpan;
     use crate::{
@@ -162,7 +162,9 @@ mod tests {
                     false,
                     TraceState::default(),
                 )),
-                ("baggage", Some(_)) => cx.with_baggage(vec![KeyValue::new("baggagekey", "value")]),
+                ("baggage", Some(_)) => {
+                    cx.with_baggage(Baggage::from_iter([KeyValue::new("baggagekey", "value")]))
+                }
                 _ => cx.clone(),
             }
         }
@@ -182,7 +184,7 @@ mod tests {
             TraceState::default(),
         )));
         // setup for baggage propagator
-        cx.with_baggage(vec![KeyValue::new("baggagekey", "value")])
+        cx.with_baggage(Baggage::from_iter([KeyValue::new("baggagekey", "value")]))
     }
 
     fn test_data() -> Vec<(&'static str, &'static str)> {