fix: ensure counter metrics include _total suffix in HELP and TYPE lines (#597)

michelhe · web-flow · commit 2d99f82481df · 2025-07-31T08:33:36.000-04:00
diff --git a/metrics-exporter-prometheus/src/exporter/builder.rs b/metrics-exporter-prometheus/src/exporter/builder.rs
@@ -631,40 +631,101 @@ mod tests {
     }
 
     #[test]
-    fn test_render_with_recommended_naming() {
-        // test 1 - no unit or description
+    fn test_render_recommended_naming_no_unit_or_description() {
         let recorder = PrometheusBuilder::new().with_recommended_naming(true).build_recorder();
 
         let key = Key::from_name("basic_counter");
-        let counter1 = recorder.register_counter(&key, &METADATA);
-        counter1.increment(42);
+        let counter = recorder.register_counter(&key, &METADATA);
+        counter.increment(42);
 
         let handle = recorder.handle();
         let rendered = handle.render();
-        let expected_counter = "# TYPE basic_counter counter\nbasic_counter_total 42\n\n";
+        let expected = "# TYPE basic_counter_total counter\nbasic_counter_total 42\n\n";
 
-        assert_eq!(rendered, expected_counter);
+        assert_eq!(rendered, expected);
+    }
 
-        // test 2 - with unit and description
+    #[test]
+    fn test_render_recommended_naming_with_unit_and_description() {
         // Note: we need to create a new recorder, as the render order is not deterministic
         let recorder = PrometheusBuilder::new().with_recommended_naming(true).build_recorder();
 
         let key_name = KeyName::from_const_str("counter_with_unit");
         let key = Key::from_name(key_name.clone());
         recorder.describe_counter(key_name, Some(Unit::Bytes), "A counter with a unit".into());
-        let counter2 = recorder.register_counter(&key, &METADATA);
-        counter2.increment(42);
+        let counter = recorder.register_counter(&key, &METADATA);
+        counter.increment(42);
 
         let handle = recorder.handle();
         let rendered = handle.render();
-        let expected_counter = concat!(
-            "# HELP counter_with_unit_bytes A counter with a unit\n",
-            "# TYPE counter_with_unit_bytes counter\n",
+        let expected: &'static str = concat!(
+            "# HELP counter_with_unit_bytes_total A counter with a unit\n",
+            "# TYPE counter_with_unit_bytes_total counter\n",
             "counter_with_unit_bytes_total 42\n",
             "\n",
         );
+        assert_eq!(rendered, expected);
+    }
 
-        assert_eq!(rendered, expected_counter);
+    #[test]
+    fn test_render_recommended_naming_manual_total_suffix_with_unit() {
+        let recorder = PrometheusBuilder::new().with_recommended_naming(true).build_recorder();
+        let key_name = KeyName::from_const_str("foo_total");
+        let key = Key::from_name(key_name.clone());
+        recorder.describe_counter(key_name, Some(Unit::Bytes), "Some help".into());
+        let counter = recorder.register_counter(&key, &METADATA);
+        counter.increment(42);
+
+        let handle = recorder.handle();
+        let rendered = handle.render();
+        let expected = concat!(
+            "# HELP foo_bytes_total Some help\n",
+            "# TYPE foo_bytes_total counter\n",
+            "foo_bytes_total 42\n",
+            "\n",
+        );
+        assert_eq!(rendered, expected);
+    }
+
+    #[test]
+    fn test_render_recommended_naming_manual_counter_suffixes() {
+        let recorder = PrometheusBuilder::new().with_recommended_naming(true).build_recorder();
+        let key_name = KeyName::from_const_str("foo_bytes_total");
+        let key = Key::from_name(key_name.clone());
+        recorder.describe_counter(key_name, Some(Unit::Bytes), "Some help".into());
+        let counter = recorder.register_counter(&key, &METADATA);
+        counter.increment(42);
+
+        let handle = recorder.handle();
+        let rendered = handle.render();
+        let expected = concat!(
+            "# HELP foo_bytes_total Some help\n",
+            "# TYPE foo_bytes_total counter\n",
+            "foo_bytes_total 42\n",
+            "\n",
+        );
+        assert_eq!(rendered, expected);
+    }
+
+    #[test]
+    fn test_render_recommended_naming_gauge_with_unit_in_name() {
+        let recorder = PrometheusBuilder::new().with_recommended_naming(true).build_recorder();
+
+        let key_name = KeyName::from_const_str("gauge_with_unit_bytes");
+        let key = Key::from_name(key_name.clone());
+        recorder.describe_gauge(key_name, Some(Unit::Bytes), "A gauge with a unit".into());
+        let gauge = recorder.register_gauge(&key, &METADATA);
+        gauge.set(42.0);
+
+        let handle = recorder.handle();
+        let rendered = handle.render();
+        let expected = concat!(
+            "# HELP gauge_with_unit_bytes A gauge with a unit\n",
+            "# TYPE gauge_with_unit_bytes gauge\n",
+            "gauge_with_unit_bytes 42\n",
+            "\n",
+        );
+        assert_eq!(rendered, expected);
     }
 
     #[test]
diff --git a/metrics-exporter-prometheus/src/formatting.rs b/metrics-exporter-prometheus/src/formatting.rs
@@ -29,13 +29,15 @@ pub fn key_to_parts(
 /// Writes a help (description) line in the Prometheus [exposition format].
 ///
 /// [exposition format]: https://github.com/prometheus/docs/blob/main/content/docs/instrumenting/exposition_formats.md#text-format-details
-pub fn write_help_line(buffer: &mut String, name: &str, unit: Option<Unit>, desc: &str) {
+pub fn write_help_line(
+    buffer: &mut String,
+    name: &str,
+    unit: Option<Unit>,
+    suffix: Option<&'static str>,
+    desc: &str,
+) {
     buffer.push_str("# HELP ");
-    buffer.push_str(name);
-    if let Some(unit) = unit {
-        buffer.push('_');
-        buffer.push_str(unit.as_str());
-    }
+    add_metric_name(buffer, name, unit, suffix);
     buffer.push(' ');
     let desc = sanitize_description(desc);
     buffer.push_str(&desc);
@@ -45,13 +47,15 @@ pub fn write_help_line(buffer: &mut String, name: &str, unit: Option<Unit>, desc
 /// Writes a metric type line in the Prometheus [exposition format].
 ///
 /// [exposition format]: https://github.com/prometheus/docs/blob/main/content/docs/instrumenting/exposition_formats.md#text-format-details
-pub fn write_type_line(buffer: &mut String, name: &str, unit: Option<Unit>, metric_type: &str) {
+pub fn write_type_line(
+    buffer: &mut String,
+    name: &str,
+    unit: Option<Unit>,
+    suffix: Option<&'static str>,
+    metric_type: &str,
+) {
     buffer.push_str("# TYPE ");
-    buffer.push_str(name);
-    if let Some(unit) = unit {
-        buffer.push('_');
-        buffer.push_str(unit.as_str());
-    }
+    add_metric_name(buffer, name, unit, suffix);
     buffer.push(' ');
     buffer.push_str(metric_type);
     buffer.push('\n');
@@ -77,18 +81,7 @@ pub fn write_metric_line<T, T2>(
     T: std::fmt::Display,
     T2: std::fmt::Display,
 {
-    buffer.push_str(name);
-
-    match unit {
-        Some(Unit::Count) | None => {}
-        Some(Unit::Percent) => add_unit(buffer, "ratio"),
-        Some(unit) => add_unit(buffer, unit.as_str()),
-    }
-
-    if let Some(suffix) = suffix {
-        buffer.push('_');
-        buffer.push_str(suffix);
-    }
+    add_metric_name(buffer, name, unit, suffix);
 
     if !labels.is_empty() || additional_label.is_some() {
         buffer.push('{');
@@ -121,24 +114,58 @@ pub fn write_metric_line<T, T2>(
     buffer.push('\n');
 }
 
-fn add_unit(buffer: &mut String, unit: &str) {
-    const SUM_SUFFIX: &str = "_sum";
-    const COUNT_SUFFIX: &str = "_count";
-    const BUCKET_SUFFIX: &str = "_bucket";
-
-    if buffer.ends_with(SUM_SUFFIX) {
-        let suffix_pos = buffer.len() - SUM_SUFFIX.len();
-        buffer.insert(suffix_pos, '_');
-        buffer.insert_str(suffix_pos + 1, unit);
-    } else if buffer.ends_with(COUNT_SUFFIX) {
-        let suffix_pos = buffer.len() - COUNT_SUFFIX.len();
-        buffer.insert(suffix_pos, '_');
-        buffer.insert_str(suffix_pos + 1, unit);
-    } else if buffer.ends_with(BUCKET_SUFFIX) {
-        let suffix_pos = buffer.len() - BUCKET_SUFFIX.len();
-        buffer.insert(suffix_pos, '_');
-        buffer.insert_str(suffix_pos + 1, unit);
-    } else {
+fn add_metric_name(
+    buffer: &mut String,
+    name: &str,
+    unit: Option<Unit>,
+    suffix: Option<&'static str>,
+) {
+    buffer.push_str(name);
+    if let Some(unit) = unit {
+        add_unit_if_missing(buffer, unit);
+    }
+    if let Some(suffix) = suffix {
+        add_suffix_if_missing(buffer, suffix);
+    }
+}
+
+/// Adds a suffix to the metric name if it is not already in the name.
+fn add_suffix_if_missing(buffer: &mut String, suffix: &str) {
+    if !buffer.ends_with(suffix) {
+        buffer.push('_');
+        buffer.push_str(suffix);
+    }
+}
+
+/// Adds a unit to the metric name if it is not already in the name.
+/// If the metric ends with a known suffix, we try to insert the unit before the suffix.
+/// Otherwise, we append the unit to the end of the metric name.
+fn add_unit_if_missing(buffer: &mut String, unit: Unit) {
+    const KNOWN_SUFFIXES: [&str; 4] = ["_sum", "_count", "_bucket", "_total"];
+
+    let unit = match unit {
+        Unit::Count => {
+            // For count, we don't suffix the unit.
+            return;
+        }
+        Unit::Percent => "ratio",
+        unit => unit.as_str(),
+    };
+
+    let mut handled = false;
+    for suffix in KNOWN_SUFFIXES {
+        if buffer.ends_with(suffix) {
+            let suffix_pos = buffer.len() - suffix.len();
+            // Check if name before suffix already has the unit
+            if !&buffer[..suffix_pos].ends_with(unit) {
+                buffer.insert(suffix_pos, '_');
+                buffer.insert_str(suffix_pos + 1, unit);
+            }
+            handled = true;
+            break;
+        }
+    }
+    if !handled && !buffer.ends_with(unit) {
         buffer.push('_');
         buffer.push_str(unit);
     }
diff --git a/metrics-exporter-prometheus/src/recorder.rs b/metrics-exporter-prometheus/src/recorder.rs
@@ -120,11 +120,11 @@ impl Inner {
         for (name, mut by_labels) in counters.drain() {
             let unit = descriptions.get(name.as_str()).and_then(|(desc, unit)| {
                 let unit = unit.filter(|_| self.enable_unit_suffix);
-                write_help_line(&mut output, name.as_str(), unit, desc);
+                write_help_line(&mut output, name.as_str(), unit, self.counter_suffix, desc);
                 unit
             });
 
-            write_type_line(&mut output, name.as_str(), unit, "counter");
+            write_type_line(&mut output, name.as_str(), unit, self.counter_suffix, "counter");
             for (labels, value) in by_labels.drain() {
                 write_metric_line::<&str, u64>(
                     &mut output,
@@ -142,11 +142,11 @@ impl Inner {
         for (name, mut by_labels) in gauges.drain() {
             let unit = descriptions.get(name.as_str()).and_then(|(desc, unit)| {
                 let unit = unit.filter(|_| self.enable_unit_suffix);
-                write_help_line(&mut output, name.as_str(), unit, desc);
+                write_help_line(&mut output, name.as_str(), unit, None, desc);
                 unit
             });
 
-            write_type_line(&mut output, name.as_str(), unit, "gauge");
+            write_type_line(&mut output, name.as_str(), unit, None, "gauge");
             for (labels, value) in by_labels.drain() {
                 write_metric_line::<&str, f64>(
                     &mut output,
@@ -164,12 +164,12 @@ impl Inner {
         for (name, mut by_labels) in distributions.drain() {
             let unit = descriptions.get(name.as_str()).and_then(|(desc, unit)| {
                 let unit = unit.filter(|_| self.enable_unit_suffix);
-                write_help_line(&mut output, name.as_str(), unit, desc);
+                write_help_line(&mut output, name.as_str(), unit, None, desc);
                 unit
             });
 
             let distribution_type = self.distribution_builder.get_distribution_type(name.as_str());
-            write_type_line(&mut output, name.as_str(), unit, distribution_type);
+            write_type_line(&mut output, name.as_str(), unit, None, distribution_type);
             for (labels, distribution) in by_labels.drain(..) {
                 let (sum, count) = match distribution {
                     Distribution::Summary(summary, quantiles, sum) => {
