refactor how collected metrics are stored to avoid parsing labels

alex · alex · commit fb715209007e · 2025-09-11T20:22:23.000-04:00
currently the prootbuf implementation takes the stringified-labels and parses them back into their original form, which is silly. refactored to avoid that, and in the process also avoid some copies and allocations.

this changes public APIs around these, so should only be landed with an appropriate version bump
diff --git a/metrics-exporter-prometheus/src/common.rs b/metrics-exporter-prometheus/src/common.rs
@@ -80,9 +80,42 @@ pub enum BuildError {
     ZeroBucketDuration,
 }
 
+/// Represents a set of labels as structured key-value pairs
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct LabelSet {
+    pub labels: Vec<(String, String)>,
+}
+
+impl LabelSet {
+    pub fn from_key_and_global(
+        key: &metrics::Key,
+        global_labels: &IndexMap<String, String>,
+    ) -> Self {
+        let mut labels = global_labels.clone();
+        key.labels().for_each(|label| {
+            labels.insert(label.key().to_string(), label.value().to_string());
+        });
+        Self { labels: labels.into_iter().collect() }
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.labels.is_empty()
+    }
+
+    pub fn to_strings(&self) -> impl Iterator<Item = String> + '_ {
+        self.labels.iter().map(|(k, v)| {
+            format!(
+                "{}=\"{}\"",
+                crate::formatting::sanitize_label_key(k),
+                crate::formatting::sanitize_label_value(v)
+            )
+        })
+    }
+}
+
 #[derive(Debug)]
 pub struct Snapshot {
-    pub counters: HashMap<String, HashMap<Vec<String>, u64>>,
-    pub gauges: HashMap<String, HashMap<Vec<String>, f64>>,
-    pub distributions: HashMap<String, IndexMap<Vec<String>, Distribution>>,
+    pub counters: HashMap<String, HashMap<LabelSet, u64>>,
+    pub gauges: HashMap<String, HashMap<LabelSet, f64>>,
+    pub distributions: HashMap<String, IndexMap<LabelSet, Distribution>>,
 }
diff --git a/metrics-exporter-prometheus/src/formatting.rs b/metrics-exporter-prometheus/src/formatting.rs
@@ -1,30 +1,8 @@
 //! Helpers for rendering metrics in the Prometheus exposition format.
 
-use indexmap::IndexMap;
-use metrics::{Key, Unit};
+use metrics::Unit;
 
-/// Breaks a key into the name and label components, with optional default labels.
-///
-/// If any of the default labels are not already present, they will be added to the overall list of labels.
-///
-/// Both the metric name, and labels, are sanitized. See [`sanitize_metric_name`], [`sanitize_label_key`],
-/// and [`sanitize_label_value`] for more information.
-pub fn key_to_parts(
-    key: &Key,
-    default_labels: Option<&IndexMap<String, String>>,
-) -> (String, Vec<String>) {
-    let name = sanitize_metric_name(key.name());
-    let mut values = default_labels.cloned().unwrap_or_default();
-    key.labels().for_each(|label| {
-        values.insert(label.key().to_string(), label.value().to_string());
-    });
-    let labels = values
-        .iter()
-        .map(|(k, v)| format!("{}=\"{}\"", sanitize_label_key(k), sanitize_label_value(v)))
-        .collect();
-
-    (name, labels)
-}
+use crate::common::LabelSet;
 
 /// Writes a help (description) line in the Prometheus [exposition format].
 ///
@@ -73,7 +51,7 @@ pub fn write_metric_line<T, T2>(
     buffer: &mut String,
     name: &str,
     suffix: Option<&'static str>,
-    labels: &[String],
+    labels: &LabelSet,
     additional_label: Option<(&'static str, T)>,
     value: T2,
     unit: Option<Unit>,
@@ -87,13 +65,13 @@ pub fn write_metric_line<T, T2>(
         buffer.push('{');
 
         let mut first = true;
-        for label in labels {
+        for label in labels.to_strings() {
             if first {
                 first = false;
             } else {
                 buffer.push(',');
             }
-            buffer.push_str(label);
+            buffer.push_str(&label);
         }
 
         if let Some((name, value)) = additional_label {
diff --git a/metrics-exporter-prometheus/src/protobuf.rs b/metrics-exporter-prometheus/src/protobuf.rs
@@ -1,11 +1,10 @@
 //! Protobuf serialization support for Prometheus metrics.
 
-use indexmap::IndexMap;
 use metrics::Unit;
 use prost::Message;
 use std::collections::HashMap;
 
-use crate::common::Snapshot;
+use crate::common::{LabelSet, Snapshot};
 use crate::distribution::Distribution;
 use crate::formatting::sanitize_metric_name;
 
@@ -26,28 +25,27 @@ pub(crate) const PROTOBUF_CONTENT_TYPE: &str =
 /// length header.
 #[allow(clippy::too_many_lines)]
 pub(crate) fn render_protobuf(
-    snapshot: &Snapshot,
+    snapshot: Snapshot,
     descriptions: &HashMap<String, (metrics::SharedString, Option<Unit>)>,
-    global_labels: &IndexMap<String, String>,
     counter_suffix: Option<&'static str>,
 ) -> Vec<u8> {
     let mut output = Vec::new();
 
     // Process counters
-    for (name, by_labels) in &snapshot.counters {
-        let sanitized_name = sanitize_metric_name(name);
+    for (name, by_labels) in snapshot.counters {
+        let sanitized_name = sanitize_metric_name(&name);
         let help =
             descriptions.get(name.as_str()).map(|(desc, _)| desc.to_string()).unwrap_or_default();
 
         let mut metrics = Vec::new();
         for (labels, value) in by_labels {
-            let label_pairs = parse_labels(labels, global_labels);
+            let label_pairs = label_set_to_protobuf(labels);
 
             metrics.push(pb::Metric {
                 label: label_pairs,
                 counter: Some(pb::Counter {
                     #[allow(clippy::cast_precision_loss)]
-                    value: Some(*value as f64),
+                    value: Some(value as f64),
 
                     ..Default::default()
                 }),
@@ -68,18 +66,18 @@ pub(crate) fn render_protobuf(
     }
 
     // Process gauges
-    for (name, by_labels) in &snapshot.gauges {
-        let sanitized_name = sanitize_metric_name(name);
+    for (name, by_labels) in snapshot.gauges {
+        let sanitized_name = sanitize_metric_name(&name);
         let help =
             descriptions.get(name.as_str()).map(|(desc, _)| desc.to_string()).unwrap_or_default();
 
         let mut metrics = Vec::new();
         for (labels, value) in by_labels {
-            let label_pairs = parse_labels(labels, global_labels);
+            let label_pairs = label_set_to_protobuf(labels);
 
             metrics.push(pb::Metric {
                 label: label_pairs,
-                gauge: Some(pb::Gauge { value: Some(*value) }),
+                gauge: Some(pb::Gauge { value: Some(value) }),
 
                 ..Default::default()
             });
@@ -97,18 +95,20 @@ pub(crate) fn render_protobuf(
     }
 
     // Process distributions (histograms and summaries)
-    for (name, by_labels) in &snapshot.distributions {
-        let sanitized_name = sanitize_metric_name(name);
+    for (name, by_labels) in snapshot.distributions {
+        let sanitized_name = sanitize_metric_name(&name);
         let help =
             descriptions.get(name.as_str()).map(|(desc, _)| desc.to_string()).unwrap_or_default();
 
         let mut metrics = Vec::new();
+        let mut metric_type = None;
         for (labels, distribution) in by_labels {
-            let label_pairs = parse_labels(labels, global_labels);
+            let label_pairs = label_set_to_protobuf(labels);
 
             let metric = match distribution {
                 Distribution::Summary(summary, quantiles, sum) => {
                     use quanta::Instant;
+                    metric_type = Some(pb::MetricType::Summary);
                     let snapshot = summary.snapshot(Instant::now());
                     let quantile_values: Vec<pb::Quantile> = quantiles
                         .iter()
@@ -122,7 +122,7 @@ pub(crate) fn render_protobuf(
                         label: label_pairs,
                         summary: Some(pb::Summary {
                             sample_count: Some(summary.count() as u64),
-                            sample_sum: Some(*sum),
+                            sample_sum: Some(sum),
                             quantile: quantile_values,
 
                             created_timestamp: None,
@@ -132,6 +132,7 @@ pub(crate) fn render_protobuf(
                     }
                 }
                 Distribution::Histogram(histogram) => {
+                    metric_type = Some(pb::MetricType::Histogram);
                     let mut buckets = Vec::new();
                     for (le, count) in histogram.buckets() {
                         buckets.push(pb::Bucket {
@@ -167,10 +168,9 @@ pub(crate) fn render_protobuf(
             metrics.push(metric);
         }
 
-        let metric_type = match by_labels.values().next() {
-            Some(Distribution::Summary(_, _, _)) => pb::MetricType::Summary,
-            Some(Distribution::Histogram(_)) => pb::MetricType::Histogram,
-            None => continue, // Skip empty metric families
+        let Some(metric_type) = metric_type else {
+            // Skip empty metric families
+            continue;
         };
 
         let metric_family = pb::MetricFamily {
@@ -187,29 +187,11 @@ pub(crate) fn render_protobuf(
     output
 }
 
-fn parse_labels(labels: &[String], global_labels: &IndexMap<String, String>) -> Vec<pb::LabelPair> {
+fn label_set_to_protobuf(labels: LabelSet) -> Vec<pb::LabelPair> {
     let mut label_pairs = Vec::new();
 
-    // Add global labels first
-    for (key, value) in global_labels {
-        label_pairs.push(pb::LabelPair { name: Some(key.clone()), value: Some(value.clone()) });
-    }
-
-    // Add metric-specific labels
-    for label_str in labels {
-        if let Some(eq_pos) = label_str.find('=') {
-            let key = &label_str[..eq_pos];
-            let value = &label_str[eq_pos + 1..];
-            let value = value.trim_matches('"');
-
-            // Skip if this label key already exists from global labels
-            if !global_labels.contains_key(key) {
-                label_pairs.push(pb::LabelPair {
-                    name: Some(key.to_string()),
-                    value: Some(value.to_string()),
-                });
-            }
-        }
+    for (key, value) in labels.labels {
+        label_pairs.push(pb::LabelPair { name: Some(key), value: Some(value) });
     }
 
     label_pairs
@@ -235,16 +217,18 @@ mod tests {
     fn test_render_protobuf_counters() {
         let mut counters = HashMap::new();
         let mut counter_labels = HashMap::new();
-        counter_labels.insert(vec!["method=\"GET\"".to_string()], 42u64);
+        let labels = LabelSet::from_key_and_global(
+            &metrics::Key::from_parts("", vec![metrics::Label::new("method", "GET")]),
+            &IndexMap::new(),
+        );
+        counter_labels.insert(labels, 42u64);
         counters.insert("http_requests".to_string(), counter_labels);
 
         let snapshot = Snapshot { counters, gauges: HashMap::new(), distributions: HashMap::new() };
 
         let descriptions = HashMap::new();
-        let global_labels = IndexMap::new();
 
-        let protobuf_data =
-            render_protobuf(&snapshot, &descriptions, &global_labels, Some("total"));
+        let protobuf_data = render_protobuf(snapshot, &descriptions, Some("total"));
 
         assert!(!protobuf_data.is_empty(), "Protobuf data should not be empty");
 
@@ -264,7 +248,11 @@ mod tests {
     fn test_render_protobuf_gauges() {
         let mut gauges = HashMap::new();
         let mut gauge_labels = HashMap::new();
-        gauge_labels.insert(vec!["instance=\"localhost\"".to_string()], 0.75f64);
+        let labels = LabelSet::from_key_and_global(
+            &metrics::Key::from_parts("", vec![metrics::Label::new("instance", "localhost")]),
+            &IndexMap::new(),
+        );
+        gauge_labels.insert(labels, 0.75f64);
         gauges.insert("cpu_usage".to_string(), gauge_labels);
 
         let snapshot = Snapshot { counters: HashMap::new(), gauges, distributions: HashMap::new() };
@@ -274,9 +262,8 @@ mod tests {
             "cpu_usage".to_string(),
             (SharedString::const_str("CPU usage percentage"), None),
         );
-        let global_labels = IndexMap::new();
 
-        let protobuf_data = render_protobuf(&snapshot, &descriptions, &global_labels, None);
+        let protobuf_data = render_protobuf(snapshot, &descriptions, None);
 
         assert!(!protobuf_data.is_empty(), "Protobuf data should not be empty");
 
diff --git a/metrics-exporter-prometheus/src/recorder.rs b/metrics-exporter-prometheus/src/recorder.rs
@@ -8,18 +8,18 @@ use metrics::{Counter, Gauge, Histogram, Key, KeyName, Metadata, Recorder, Share
 use metrics_util::registry::{Recency, Registry};
 use quanta::Instant;
 
-use crate::common::Snapshot;
+use crate::common::{LabelSet, Snapshot};
 use crate::distribution::{Distribution, DistributionBuilder};
 use crate::formatting::{
-    key_to_parts, sanitize_metric_name, write_help_line, write_metric_line, write_type_line,
+    sanitize_metric_name, write_help_line, write_metric_line, write_type_line,
 };
 use crate::registry::GenerationalAtomicStorage;
 
 #[derive(Debug)]
 pub(crate) struct Inner {
     pub registry: Registry<Key, GenerationalAtomicStorage>,
     pub recency: Recency<Key>,
-    pub distributions: RwLock<HashMap<String, IndexMap<Vec<String>, Distribution>>>,
+    pub distributions: RwLock<HashMap<String, IndexMap<LabelSet, Distribution>>>,
     pub distribution_builder: DistributionBuilder,
     pub descriptions: RwLock<HashMap<String, (SharedString, Option<Unit>)>>,
     pub global_labels: IndexMap<String, String>,
@@ -37,7 +37,8 @@ impl Inner {
                 continue;
             }
 
-            let (name, labels) = key_to_parts(&key, Some(&self.global_labels));
+            let name = sanitize_metric_name(key.name());
+            let labels = LabelSet::from_key_and_global(&key, &self.global_labels);
             let value = counter.get_inner().load(Ordering::Acquire);
             let entry =
                 counters.entry(name).or_insert_with(HashMap::new).entry(labels).or_insert(0);
@@ -52,7 +53,8 @@ impl Inner {
                 continue;
             }
 
-            let (name, labels) = key_to_parts(&key, Some(&self.global_labels));
+            let name = sanitize_metric_name(key.name());
+            let labels = LabelSet::from_key_and_global(&key, &self.global_labels);
             let value = f64::from_bits(gauge.get_inner().load(Ordering::Acquire));
             let entry =
                 gauges.entry(name).or_insert_with(HashMap::new).entry(labels).or_insert(0.0);
@@ -69,7 +71,8 @@ impl Inner {
                 // Since we store aggregated distributions directly, when we're told that a metric
                 // is not recent enough and should be/was deleted from the registry, we also need to
                 // delete it on our side as well.
-                let (name, labels) = key_to_parts(&key, Some(&self.global_labels));
+                let name = sanitize_metric_name(key.name());
+                let labels = LabelSet::from_key_and_global(&key, &self.global_labels);
                 let mut wg = self.distributions.write().unwrap_or_else(PoisonError::into_inner);
                 let delete_by_name = if let Some(by_name) = wg.get_mut(&name) {
                     by_name.swap_remove(&labels);
@@ -98,7 +101,8 @@ impl Inner {
     fn drain_histograms_to_distributions(&self) {
         let histogram_handles = self.registry.get_histogram_handles();
         for (key, histogram) in histogram_handles {
-            let (name, labels) = key_to_parts(&key, Some(&self.global_labels));
+            let name = sanitize_metric_name(key.name());
+            let labels = LabelSet::from_key_and_global(&key, &self.global_labels);
 
             let mut wg = self.distributions.write().unwrap_or_else(PoisonError::into_inner);
             let entry = wg
@@ -332,12 +336,7 @@ impl PrometheusHandle {
         let snapshot = self.inner.get_recent_metrics();
         let descriptions = self.inner.descriptions.read().unwrap_or_else(PoisonError::into_inner);
 
-        crate::protobuf::render_protobuf(
-            &snapshot,
-            &descriptions,
-            &self.inner.global_labels,
-            self.inner.counter_suffix,
-        )
+        crate::protobuf::render_protobuf(snapshot, &descriptions, self.inner.counter_suffix)
     }
 
     /// Performs upkeeping operations to ensure metrics held by recorder are up-to-date and do not
