perf: Baggage insert & BaggageExt::with_baggage & updated constants to latest standard

opentelemetry/CHANGELOG.md

@@ -8,6 +8,7 @@
 - {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).
 
 ## 0.28.0
 

opentelemetry/src/baggage.rs

@@ -15,14 +15,14 @@
 //!
 //! [W3C Baggage]: https://w3c.github.io/baggage
 use crate::{Context, Key, KeyValue, StringValue};
+use std::collections::hash_map::Entry;
 use std::collections::{hash_map, HashMap};
 use std::fmt;
 use std::sync::OnceLock;
 
 static DEFAULT_BAGGAGE: OnceLock<Baggage> = OnceLock::new();
 
-const MAX_KEY_VALUE_PAIRS: usize = 180;
-const MAX_BYTES_FOR_ONE_PAIR: usize = 4096;
+const MAX_KEY_VALUE_PAIRS: usize = 64;
 const MAX_LEN_OF_ALL_PAIRS: usize = 8192;
 
 /// Returns the default baggage, ensuring it is initialized only once.
@@ -49,11 +49,10 @@ fn get_default_baggage() -> &'static Baggage {
 ///
 /// ### Limits
 ///
-/// * Maximum number of name/value pairs: `180`.
-/// * Maximum number of bytes per a single name/value pair: `4096`.
+/// * Maximum number of name/value pairs: `64`.
 /// * Maximum total length of all name/value pairs: `8192`.
 ///
-/// [RFC2616, Section 2.2]: https://tools.ietf.org/html/rfc2616#section-2.2
+/// <https://www.w3.org/TR/baggage/#limits>
 #[derive(Debug, Default)]
 pub struct Baggage {
     inner: HashMap<Key, (StringValue, BaggageMetadata)>,
@@ -133,6 +132,8 @@ impl Baggage {
     /// Same with `insert`, if the name was not present, [`None`] will be returned.
     /// If the name is present, the old value and metadata will be returned.
     ///
+    /// Also checks for [limits](https://w3c.github.io/baggage/#limits).
+    ///
     /// # Examples
     ///
     /// ```
@@ -155,10 +156,38 @@ impl Baggage {
         S: Into<BaggageMetadata>,
     {
         let (key, value, metadata) = (key.into(), value.into(), metadata.into());
-        if self.insertable(&key, &value, &metadata) {
-            self.inner.insert(key, (value, metadata))
-        } else {
-            None
+        if !key.as_str().is_ascii() {
+            return None;
+        }
+        let entry_content_len =
+            key_value_metadata_bytes_size(key.as_str(), value.as_str(), metadata.as_str());
+        let entries_count = self.inner.len();
+        match self.inner.entry(key) {
+            Entry::Occupied(mut occupied_entry) => {
+                let prev_content_len = key_value_metadata_bytes_size(
+                    occupied_entry.key().as_str(),
+                    occupied_entry.get().0.as_str(),
+                    occupied_entry.get().1.as_str(),
+                );
+                let new_content_len = self.kv_content_len + entry_content_len - prev_content_len;
+                if new_content_len > MAX_LEN_OF_ALL_PAIRS {
+                    return None;
+                }
+                self.kv_content_len = new_content_len;
+                Some(occupied_entry.insert((value, metadata)))
+            }
+            Entry::Vacant(vacant_entry) => {
+                if entries_count == MAX_KEY_VALUE_PAIRS {
+                    return None;
+                }
+                let new_content_len = self.kv_content_len + entry_content_len;
+                if new_content_len > MAX_LEN_OF_ALL_PAIRS {
+                    return None;
+                }
+                self.kv_content_len = new_content_len;
+                vacant_entry.insert((value, metadata));
+                None
+            }
         }
     }
 
@@ -178,59 +207,10 @@ impl Baggage {
         self.inner.is_empty()
     }
 
-    /// Gets an iterator over the baggage items, sorted by name.
+    /// Gets an iterator over the baggage items, in any order.
     pub fn iter(&self) -> Iter<'_> {
         self.into_iter()
     }
-
-    /// Determine whether the key value pair exceed one of the [limits](https://w3c.github.io/baggage/#limits).
-    /// If not, update the total length of key values
-    fn insertable(&mut self, key: &Key, value: &StringValue, metadata: &BaggageMetadata) -> bool {
-        if !key.as_str().is_ascii() {
-            return false;
-        }
-        let value = value.as_str();
-        if key_value_metadata_bytes_size(key.as_str(), value, metadata.as_str())
-            < MAX_BYTES_FOR_ONE_PAIR
-        {
-            match self.inner.get(key) {
-                None => {
-                    // check total length
-                    if self.kv_content_len
-                        + metadata.as_str().len()
-                        + value.len()
-                        + key.as_str().len()
-                        > MAX_LEN_OF_ALL_PAIRS
-                    {
-                        return false;
-                    }
-                    // check number of pairs
-                    if self.inner.len() + 1 > MAX_KEY_VALUE_PAIRS {
-                        return false;
-                    }
-                    self.kv_content_len +=
-                        metadata.as_str().len() + value.len() + key.as_str().len()
-                }
-                Some((old_value, old_metadata)) => {
-                    let old_value = old_value.as_str();
-                    if self.kv_content_len - old_metadata.as_str().len() - old_value.len()
-                        + metadata.as_str().len()
-                        + value.len()
-                        > MAX_LEN_OF_ALL_PAIRS
-                    {
-                        return false;
-                    }
-                    self.kv_content_len =
-                        self.kv_content_len - old_metadata.as_str().len() - old_value.len()
-                            + metadata.as_str().len()
-                            + value.len()
-                }
-            }
-            true
-        } else {
-            false
-        }
-    }
 }
 
 /// Get the number of bytes for one key-value pair
@@ -385,13 +365,11 @@ impl BaggageExt for Context {
         &self,
         baggage: T,
     ) -> Self {
-        let mut merged: Baggage = self
-            .baggage()
-            .iter()
-            .map(|(key, (value, metadata))| {
-                KeyValueMetadata::new(key.clone(), value.clone(), metadata.clone())
-            })
-            .collect();
+        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);
         }
@@ -531,29 +509,6 @@ mod tests {
         assert_eq!(baggage.len(), MAX_KEY_VALUE_PAIRS)
     }
 
-    #[test]
-    fn insert_too_long_pair() {
-        let pair = KeyValue::new(
-            "test",
-            String::from_utf8_lossy(vec![12u8; MAX_BYTES_FOR_ONE_PAIR].as_slice()).to_string(),
-        );
-        let mut baggage = Baggage::default();
-        baggage.insert(pair.key.clone(), pair.value.clone());
-        assert_eq!(
-            baggage.len(),
-            0,
-            "The input pair is too long to insert into baggage"
-        );
-
-        baggage.insert("test", "value");
-        baggage.insert(pair.key.clone(), pair.value);
-        assert_eq!(
-            baggage.get(pair.key),
-            Some(&StringValue::from("value")),
-            "If the input pair is too long, then don't replace entry with same key"
-        )
-    }
-
     #[test]
     fn insert_pairs_length_exceed() {
         let mut data = vec![];
@@ -611,4 +566,22 @@ mod tests {
         assert!(b.to_string().contains("bar=2;yellow"));
         assert!(b.to_string().contains("foo=1;red;state=on"));
     }
+
+    #[test]
+    fn replace_existing_key() {
+        let half_minus2: StringValue = (0..MAX_LEN_OF_ALL_PAIRS / 2 - 2)
+            .map(|_| 'x')
+            .collect::<String>()
+            .into();
+
+        let mut b = Baggage::default();
+        b.insert("a", half_minus2.clone()); // +1 for key
+        b.insert("b", half_minus2); // +1 for key
+        b.insert("c", StringValue::from(".")); // total of 2 bytes
+        assert!(b.get("a").is_some());
+        assert!(b.get("b").is_some());
+        assert!(b.get("c").is_some());
+        assert!(b.insert("c", StringValue::from("..")).is_none()); // exceeds MAX_LEN_OF_ALL_PAIRS
+        assert_eq!(b.insert("c", StringValue::from("!")).unwrap(), ".".into()); // replaces existing
+    }
 }