feat(stats)!: implement additional metric tags in libdd-trace-stats

libdd-data-pipeline/src/trace_exporter/builder.rs

@@ -763,6 +763,7 @@ impl<R: SharedRuntime> TraceExporterBuilder<R> {
                 span_kinds,
                 self.peer_tags.clone(),
                 None,
+                vec![],
                 #[cfg(feature = "stats-obfuscation")]
                 None,
             )));

libdd-data-pipeline/src/trace_exporter/stats.rs

@@ -143,6 +143,7 @@ pub(crate) fn start_stats_computation<
             span_kinds,
             peer_tags,
             ctx.stats_cardinality_limit,
+            vec![],
             #[cfg(feature = "stats-obfuscation")]
             Some(client_side_stats.obfuscation_config.clone()),
         )));

libdd-trace-protobuf/src/pb_test.rs

@@ -72,7 +72,8 @@ mod tests {
                             ],
                             "HTTPMethod": "GET",
                             "HTTPEndpoint": "/test",
-                            "GRPCStatusCode": "0"
+                            "GRPCStatusCode": "0",
+                            "AdditionalMetricTags": []
                         }
                     ]
                 }

libdd-trace-stats/README.md

@@ -76,6 +76,7 @@ let mut concentrator = SpanConcentrator::new(
     SystemTime::now(),
     vec!["client".to_string(), "server".to_string()], // eligible span kinds
     vec!["peer.service".to_string()], // peer tag keys
+    vec!["example.key".to_string()], // additional metric tag keys
 );
 
 // Add spans

libdd-trace-stats/benches/span_concentrator_bench.rs

@@ -45,6 +45,7 @@ pub fn criterion_benchmark(c: &mut Criterion) {
         vec![],
         vec!["db_name".into(), "bucket_s3".into()],
         None,
+        vec![],
         #[cfg(feature = "stats-obfuscation")]
         None,
     );

libdd-trace-stats/src/span_concentrator/aggregation.rs

@@ -10,13 +10,15 @@ use libdd_trace_obfuscation::ip_address::quantize_peer_ip_addresses;
 use libdd_trace_protobuf::pb;
 use libdd_trace_utils::span::SpanText;
 use std::borrow::{Borrow, Cow};
+use tracing::warn;
 
 use crate::span_concentrator::StatSpan;
 
 /// Sentinel value used for cardinality limiting.
 pub const TRACER_BLOCKED_VALUE: &str = "tracer_blocked_value";
 
 const TAG_STATUS_CODE: &str = "http.status_code";
+const ADDITIONAL_METRIC_TAG_VALUE_MAX_LEN: usize = 200;
 const TAG_SYNTHETICS: &str = "synthetics";
 const TAG_SPANKIND: &str = "span.kind";
 const TAG_ORIGIN: &str = "_dd.origin";
@@ -84,6 +86,7 @@ impl<T> FixedAggregationKey<T> {
 pub(super) struct BorrowedAggregationKey<'a> {
     fixed: FixedAggregationKey<&'a str>,
     peer_tags: Vec<(&'a str, Cow<'a, str>)>,
+    additional_metric_tags: Vec<(&'a str, &'a str)>,
 }
 
 impl hashbrown::Equivalent<OwnedAggregationKey> for BorrowedAggregationKey<'_> {
@@ -96,6 +99,12 @@ impl hashbrown::Equivalent<OwnedAggregationKey> for BorrowedAggregationKey<'_> {
                 .iter()
                 .zip(other.peer_tags.iter())
                 .all(|((k1, v1), (k2, v2))| k1 == k2 && v1 == v2)
+            && self.additional_metric_tags.len() == other.additional_metric_tags.len()
+            && self
+                .additional_metric_tags
+                .iter()
+                .zip(other.additional_metric_tags.iter())
+                .all(|((k1, v1), (k2, v2))| k1 == k2 && v1 == v2)
     }
 }
 
@@ -110,6 +119,7 @@ impl hashbrown::Equivalent<OwnedAggregationKey> for BorrowedAggregationKey<'_> {
 pub(super) struct OwnedAggregationKey {
     fixed: FixedAggregationKey<String>,
     peer_tags: Vec<(String, String)>,
+    additional_metric_tags: Vec<(String, String)>,
 }
 
 impl From<&BorrowedAggregationKey<'_>> for OwnedAggregationKey {
@@ -121,6 +131,11 @@ impl From<&BorrowedAggregationKey<'_>> for OwnedAggregationKey {
                 .iter()
                 .map(|(k, v)| (k.to_string(), v.to_string()))
                 .collect(),
+            additional_metric_tags: value
+                .additional_metric_tags
+                .iter()
+                .map(|(k, v)| (k.to_string(), v.to_string()))
+                .collect(),
         }
     }
 }
@@ -203,16 +218,28 @@ fn grpc_status_str_to_int_value(v: &str) -> Option<u8> {
 impl<'a> BorrowedAggregationKey<'a> {
     /// Return an AggregationKey matching the given span.
     ///
-    /// If `peer_tags_keys` is not empty then the peer tags of the span will be included in the
+    /// If `peer_tag_keys` is not empty then the peer tags of the span will be included in the
     /// key.
-    pub(super) fn from_span<T: StatSpan<'a>>(span: &'a T, peer_tag_keys: &'a [String]) -> Self {
-        Self::from_obfuscated_span(span.resource(), span, peer_tag_keys)
+    /// If `additional_metric_tags` is not empty then matching span tags keys are included in the
+    /// key.
+    pub(super) fn from_span<T: StatSpan<'a>>(
+        span: &'a T,
+        peer_tag_keys: &'a [String],
+        additional_metric_tag_keys: &'a [String],
+    ) -> Self {
+        Self::from_obfuscated_span(
+            span.resource(),
+            span,
+            peer_tag_keys,
+            additional_metric_tag_keys,
+        )
     }
 
     pub(crate) fn from_obfuscated_span<'b, T>(
         resource_name: &'a str,
         span: &'b T,
         peer_tag_keys: &'b [String],
+        additional_metric_tag_keys: &'b [String],
     ) -> BorrowedAggregationKey<'a>
     where
         T: StatSpan<'b>,
@@ -255,6 +282,24 @@ impl<'a> BorrowedAggregationKey<'a> {
         let grpc_status_code = get_grpc_status_code(span);
         let service_source = span.get_meta(TAG_SVC_SRC).unwrap_or_default();
 
+        let additional_metric_tags: Vec<(&'a str, &'a str)> = additional_metric_tag_keys
+            .iter()
+            .filter_map(|key| match span.get_meta(key.as_str()) {
+                Some(v) if !v.is_empty() => {
+                    if v.chars().count() > ADDITIONAL_METRIC_TAG_VALUE_MAX_LEN {
+                        warn!(
+                            "additional_metric_tags: value for key '{}' exceeds {} characters; substituting tracer_blocked_value",
+                            key, ADDITIONAL_METRIC_TAG_VALUE_MAX_LEN,
+                        );
+                        Some((key.as_str(), TRACER_BLOCKED_VALUE))
+                    } else {
+                        Some((key.as_str(), v))
+                    }
+                }
+                _ => None,
+            })
+            .collect();
+
         Self {
             fixed: FixedAggregationKey {
                 resource_name,
@@ -277,6 +322,7 @@ impl<'a> BorrowedAggregationKey<'a> {
                 },
             },
             peer_tags,
+            additional_metric_tags,
         }
     }
 }
@@ -300,6 +346,7 @@ impl OwnedAggregationKey {
                 is_trace_root: pb::Trilean::NotSet,
             },
             peer_tags: vec![],
+            additional_metric_tags: vec![],
         }
     }
 }
@@ -330,6 +377,14 @@ impl From<pb::ClientGroupedStats> for OwnedAggregationKey {
                     Some((key.to_string(), value.to_string()))
                 })
                 .collect(),
+            additional_metric_tags: value
+                .additional_metric_tags
+                .into_iter()
+                .filter_map(|t| {
+                    let (key, value) = t.split_once(':')?;
+                    Some((key.to_string(), value.to_string()))
+                })
+                .collect(),
         }
     }
 }
@@ -462,28 +517,43 @@ impl StatsBucket {
         self.collapsed_count
     }
 
-    /// Insert a value as stats in the group corresponding to the aggregation key. If the key is new
-    /// and the `max_entries` limit has not been reached, a new entry is created, else the span is
-    /// instead merged into the overflow sentinel key.
+    /// Insert a value as stats in the group corresponding to the aggregation key, if it does not
+    /// exist it creates it.
+    ///
+    /// Keys that already exist in this bucket always merge normally. A new key is subject to the
+    /// `max_entries` limit, which collapses it into the overflow sentinel key.
     pub(super) fn insert(
         &mut self,
         key: BorrowedAggregationKey<'_>,
         duration: i64,
         is_error: bool,
         is_top_level: bool,
     ) {
-        if self.data.len() >= self.max_entries && !self.data.contains_key(&key) {
-            self.collapsed_count += 1;
-            self.data
-                .entry(OwnedAggregationKey::overflow_key())
-                .or_default()
-                .insert(duration, is_error, is_top_level);
-            return;
+        // The map can't change size before the entry below is resolved, so this single read
+        // covers the `max_entries` check in the vacant branch without a further lookup.
+        let len_before_insert = self.data.len();
+
+        match self.data.entry_ref(&key) {
+            // Existing key, merge
+            hashbrown::hash_map::EntryRef::Occupied(mut e) => {
+                e.get_mut().insert(duration, is_error, is_top_level);
+            }
+            hashbrown::hash_map::EntryRef::Vacant(e) => {
+                // New key over the max entry limit, collapse into the overflow
+                // sentinel.
+                if len_before_insert >= self.max_entries {
+                    self.collapsed_count += 1;
+                    self.data
+                        .entry(OwnedAggregationKey::overflow_key())
+                        .or_default()
+                        .insert(duration, is_error, is_top_level);
+                    return;
+                }
+                // Within the max entry limit, admit key as a new distinct entry.
+                e.insert(GroupedStats::default())
+                    .insert(duration, is_error, is_top_level);
+            }
         }
-        self.data
-            .entry_ref(&key)
-            .or_default()
-            .insert(duration, is_error, is_top_level);
     }
 
     /// Consume the bucket and return a ClientStatsBucket containing the bucket stats.
@@ -561,7 +631,11 @@ fn encode_grouped_stats(key: OwnedAggregationKey, group: GroupedStats) -> pb::Cl
             .unwrap_or_default(),
         service_source: f.service_source,
         span_derived_primary_tags: vec![],
-        additional_metric_tags: vec![],
+        additional_metric_tags: key
+            .additional_metric_tags
+            .into_iter()
+            .map(|(k, v)| format!("{k}:{v}"))
+            .collect(),
     }
 }
 
@@ -585,12 +659,14 @@ mod tests {
             OwnedAggregationKey {
                 fixed: self,
                 peer_tags: vec![],
+                additional_metric_tags: vec![],
             }
         }
         fn into_key_with_peers(self, peer_tags: Vec<(String, String)>) -> OwnedAggregationKey {
             OwnedAggregationKey {
                 fixed: self,
                 peer_tags,
+                additional_metric_tags: vec![],
             }
         }
     }
@@ -1103,7 +1179,7 @@ mod tests {
         ];
 
         for (span, expected_key) in test_cases {
-            let borrowed_key = BorrowedAggregationKey::from_span(&span, &[]);
+            let borrowed_key = BorrowedAggregationKey::from_span(&span, &[], &[]);
             assert_eq!(
                 OwnedAggregationKey::from(&borrowed_key),
                 expected_key,
@@ -1116,7 +1192,8 @@ mod tests {
         }
 
         for (span, expected_key) in test_cases_with_peer_tags {
-            let borrowed_key = BorrowedAggregationKey::from_span(&span, test_peer_tags.as_slice());
+            let borrowed_key =
+                BorrowedAggregationKey::from_span(&span, test_peer_tags.as_slice(), &[]);
             assert_eq!(OwnedAggregationKey::from(&borrowed_key), expected_key);
             assert_eq!(
                 get_hash(&borrowed_key),
@@ -1144,7 +1221,7 @@ mod tests {
             .into(),
             ..Default::default()
         };
-        let key = BorrowedAggregationKey::from_span(&span_ipv4, &peer_tag_keys);
+        let key = BorrowedAggregationKey::from_span(&span_ipv4, &peer_tag_keys, &[]);
         let owned = OwnedAggregationKey::from(&key);
         assert_eq!(
             owned.peer_tags,
@@ -1172,7 +1249,7 @@ mod tests {
             ..Default::default()
         };
         let ipv6_keys = vec!["peer.hostname".to_string()];
-        let key = BorrowedAggregationKey::from_span(&span_ipv6, &ipv6_keys);
+        let key = BorrowedAggregationKey::from_span(&span_ipv6, &ipv6_keys, &[]);
         let owned = OwnedAggregationKey::from(&key);
         assert_eq!(
             owned.peer_tags,
@@ -1193,7 +1270,7 @@ mod tests {
             ..Default::default()
         };
         let non_ip_keys = vec!["db.instance".to_string()];
-        let key = BorrowedAggregationKey::from_span(&span_non_ip, &non_ip_keys);
+        let key = BorrowedAggregationKey::from_span(&span_non_ip, &non_ip_keys, &[]);
         let owned = OwnedAggregationKey::from(&key);
         assert_eq!(
             owned.peer_tags,