Add units to prometheus metric lines (#535)

fayalalebrun · web-flow · commit ed64fb60152c · 2024-11-25T18:43:24.000-05:00
diff --git a/metrics-exporter-prometheus/src/exporter/builder.rs b/metrics-exporter-prometheus/src/exporter/builder.rs
@@ -48,6 +48,7 @@ pub struct PrometheusBuilder {
     upkeep_timeout: Duration,
     recency_mask: MetricKindMask,
     global_labels: Option<IndexMap<String, String>>,
+    enable_unit_suffix: bool,
 }
 
 impl PrometheusBuilder {
@@ -80,6 +81,7 @@ impl PrometheusBuilder {
             upkeep_timeout,
             recency_mask: MetricKindMask::NONE,
             global_labels: None,
+            enable_unit_suffix: false,
         }
     }
 
@@ -279,6 +281,19 @@ impl PrometheusBuilder {
         Ok(self)
     }
 
+    /// Sets whether a unit suffix is appended to metric names.
+    ///
+    /// When this is enabled and the [`Unit`][metrics::Unit] of metric is
+    /// given, then the exported metric name will be appended to according to
+    /// the [Prometheus Best Practices](https://prometheus.io/docs/practices/naming/).
+    ///
+    /// Defaults to false.
+    #[must_use]
+    pub fn set_enable_unit_suffix(mut self, enabled: bool) -> Self {
+        self.enable_unit_suffix = enabled;
+        self
+    }
+
     /// Sets the bucket for a specific pattern.
     ///
     /// The match pattern can be a full match (equality), prefix match, or suffix match.  The matchers are applied in
@@ -512,6 +527,7 @@ impl PrometheusBuilder {
             ),
             descriptions: RwLock::new(HashMap::new()),
             global_labels: self.global_labels.unwrap_or_default(),
+            enable_unit_suffix: self.enable_unit_suffix,
         };
 
         PrometheusRecorder::from(inner)
diff --git a/metrics-exporter-prometheus/src/formatting.rs b/metrics-exporter-prometheus/src/formatting.rs
@@ -1,7 +1,7 @@
 //! Helpers for rendering metrics in the Prometheus exposition format.
 
 use indexmap::IndexMap;
-use metrics::Key;
+use metrics::{Key, Unit};
 
 /// Breaks a key into the name and label components, with optional default labels.
 ///
@@ -64,6 +64,7 @@ pub fn write_metric_line<T, T2>(
     labels: &[String],
     additional_label: Option<(&'static str, T)>,
     value: T2,
+    unit: Option<Unit>,
 ) where
     T: std::fmt::Display,
     T2: std::fmt::Display,
@@ -74,6 +75,18 @@ pub fn write_metric_line<T, T2>(
         buffer.push_str(suffix);
     }
 
+    match unit {
+        Some(Unit::Count) | None => {}
+        Some(Unit::Percent) => {
+            buffer.push('_');
+            buffer.push_str("ratio");
+        }
+        Some(unit) => {
+            buffer.push('_');
+            buffer.push_str(unit.as_str());
+        }
+    }
+
     if !labels.is_empty() || additional_label.is_some() {
         buffer.push('{');
 
diff --git a/metrics-exporter-prometheus/src/recorder.rs b/metrics-exporter-prometheus/src/recorder.rs
@@ -21,8 +21,9 @@ pub(crate) struct Inner {
     pub recency: Recency<Key>,
     pub distributions: RwLock<HashMap<String, IndexMap<Vec<String>, Distribution>>>,
     pub distribution_builder: DistributionBuilder,
-    pub descriptions: RwLock<HashMap<String, SharedString>>,
+    pub descriptions: RwLock<HashMap<String, (SharedString, Option<Unit>)>>,
     pub global_labels: IndexMap<String, String>,
+    pub enable_unit_suffix: bool,
 }
 
 impl Inner {
@@ -116,33 +117,52 @@ impl Inner {
         let descriptions = self.descriptions.read().unwrap_or_else(PoisonError::into_inner);
 
         for (name, mut by_labels) in counters.drain() {
-            if let Some(desc) = descriptions.get(name.as_str()) {
+            let unit = descriptions.get(name.as_str()).and_then(|(desc, unit)| {
                 write_help_line(&mut output, name.as_str(), desc);
-            }
+                *unit
+            });
 
             write_type_line(&mut output, name.as_str(), "counter");
             for (labels, value) in by_labels.drain() {
-                write_metric_line::<&str, u64>(&mut output, &name, None, &labels, None, value);
+                write_metric_line::<&str, u64>(
+                    &mut output,
+                    &name,
+                    None,
+                    &labels,
+                    None,
+                    value,
+                    unit.filter(|_| self.enable_unit_suffix),
+                );
             }
             output.push('\n');
         }
 
         for (name, mut by_labels) in gauges.drain() {
-            if let Some(desc) = descriptions.get(name.as_str()) {
+            let unit = descriptions.get(name.as_str()).and_then(|(desc, unit)| {
                 write_help_line(&mut output, name.as_str(), desc);
-            }
+                *unit
+            });
 
             write_type_line(&mut output, name.as_str(), "gauge");
             for (labels, value) in by_labels.drain() {
-                write_metric_line::<&str, f64>(&mut output, &name, None, &labels, None, value);
+                write_metric_line::<&str, f64>(
+                    &mut output,
+                    &name,
+                    None,
+                    &labels,
+                    None,
+                    value,
+                    unit.filter(|_| self.enable_unit_suffix),
+                );
             }
             output.push('\n');
         }
 
         for (name, mut by_labels) in distributions.drain() {
-            if let Some(desc) = descriptions.get(name.as_str()) {
+            let unit = descriptions.get(name.as_str()).and_then(|(desc, unit)| {
                 write_help_line(&mut output, name.as_str(), desc);
-            }
+                *unit
+            });
 
             let distribution_type = self.distribution_builder.get_distribution_type(name.as_str());
             write_type_line(&mut output, name.as_str(), distribution_type);
@@ -159,6 +179,7 @@ impl Inner {
                                 &labels,
                                 Some(("quantile", quantile.value())),
                                 value,
+                                unit.filter(|_| self.enable_unit_suffix),
                             );
                         }
 
@@ -173,6 +194,7 @@ impl Inner {
                                 &labels,
                                 Some(("le", le)),
                                 count,
+                                unit.filter(|_| self.enable_unit_suffix),
                             );
                         }
                         write_metric_line(
@@ -182,20 +204,30 @@ impl Inner {
                             &labels,
                             Some(("le", "+Inf")),
                             histogram.count(),
+                            unit.filter(|_| self.enable_unit_suffix),
                         );
 
                         (histogram.sum(), histogram.count())
                     }
                 };
 
-                write_metric_line::<&str, f64>(&mut output, &name, Some("sum"), &labels, None, sum);
+                write_metric_line::<&str, f64>(
+                    &mut output,
+                    &name,
+                    Some("sum"),
+                    &labels,
+                    None,
+                    sum,
+                    unit,
+                );
                 write_metric_line::<&str, u64>(
                     &mut output,
                     &name,
                     Some("count"),
                     &labels,
                     None,
                     count,
+                    unit,
                 );
             }
 
@@ -226,11 +258,16 @@ impl PrometheusRecorder {
         PrometheusHandle { inner: self.inner.clone() }
     }
 
-    fn add_description_if_missing(&self, key_name: &KeyName, description: SharedString) {
+    fn add_description_if_missing(
+        &self,
+        key_name: &KeyName,
+        description: SharedString,
+        unit: Option<Unit>,
+    ) {
         let sanitized = sanitize_metric_name(key_name.as_str());
         let mut descriptions =
             self.inner.descriptions.write().unwrap_or_else(PoisonError::into_inner);
-        descriptions.entry(sanitized).or_insert(description);
+        descriptions.entry(sanitized).or_insert((description, unit));
     }
 }
 
@@ -241,21 +278,16 @@ impl From<Inner> for PrometheusRecorder {
 }
 
 impl Recorder for PrometheusRecorder {
-    fn describe_counter(&self, key_name: KeyName, _unit: Option<Unit>, description: SharedString) {
-        self.add_description_if_missing(&key_name, description);
+    fn describe_counter(&self, key_name: KeyName, unit: Option<Unit>, description: SharedString) {
+        self.add_description_if_missing(&key_name, description, unit);
     }
 
-    fn describe_gauge(&self, key_name: KeyName, _unit: Option<Unit>, description: SharedString) {
-        self.add_description_if_missing(&key_name, description);
+    fn describe_gauge(&self, key_name: KeyName, unit: Option<Unit>, description: SharedString) {
+        self.add_description_if_missing(&key_name, description, unit);
     }
 
-    fn describe_histogram(
-        &self,
-        key_name: KeyName,
-        _unit: Option<Unit>,
-        description: SharedString,
-    ) {
-        self.add_description_if_missing(&key_name, description);
+    fn describe_histogram(&self, key_name: KeyName, unit: Option<Unit>, description: SharedString) {
+        self.add_description_if_missing(&key_name, description, unit);
     }
 
     fn register_counter(&self, key: &Key, _metadata: &Metadata<'_>) -> Counter {
