From 533039977d0685b488d14aab78c4881e0ed24db9 Mon Sep 17 00:00:00 2001
From: Shaoyuan CHEN <chensy20@mails.tsinghua.edu.cn>
Date: Sat, 9 Apr 2022 04:34:47 +0800
Subject: [PATCH] fix prometheus metric name and label key sanitizer (#296)

Co-authored-by: Toby Lawrence <tobz@users.noreply.github.com>
---
 metrics-exporter-prometheus/src/formatting.rs | 28 +++++++++++++++----
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/metrics-exporter-prometheus/src/formatting.rs b/metrics-exporter-prometheus/src/formatting.rs
index b0fd4b78..0236b1dc 100644
--- a/metrics-exporter-prometheus/src/formatting.rs
+++ b/metrics-exporter-prometheus/src/formatting.rs
@@ -128,9 +128,17 @@ pub fn sanitize_metric_name(name: &str) -> String {
 /// [data model]: https://prometheus.io/docs/concepts/data_model/#metric-names-and-labels
 pub fn sanitize_label_key(key: &str) -> String {
     // The first character must be [a-zA-Z_], and all subsequent characters must be [a-zA-Z0-9_].
-    key.replacen(invalid_label_key_start_character, "_", 1)
-        .replace(invalid_label_key_character, "_")
-        .replacen("__", "___", 1)
+    let mut out = String::with_capacity(key.len());
+    let mut is_invalid: fn(char) -> bool = invalid_label_key_start_character;
+    for c in key.chars() {
+        if is_invalid(c) {
+            out.push('_');
+        } else {
+            out.push(c);
+        }
+        is_invalid = invalid_label_key_character;
+    }
+    out
 }
 
 /// Sanitizes a label value to be valid under the Prometheus [data model].
@@ -241,6 +249,8 @@ mod tests {
             ("foo_bar", "foo_bar"),
             ("foo1_bar", "foo1_bar"),
             ("1foobar", "_foobar"),
+            ("foo1:bar2", "foo1:bar2"),
+            ("123", "_23"),
         ];
 
         for (input, expected) in cases {
@@ -257,7 +267,9 @@ mod tests {
             (":", "_"),
             ("foo_bar", "foo_bar"),
             ("1foobar", "_foobar"),
-            ("__foobar", "___foobar"),
+            ("__foobar", "__foobar"),
+            ("foo1bar2", "foo1bar2"),
+            ("123", "_23"),
         ];
 
         for (input, expected) in cases {
@@ -329,13 +341,17 @@ mod tests {
 
             // Label keys cannot begin with two underscores, as that format is reserved for internal
             // use.
-            if as_chars.len() == 2 {
+            //
+            // TODO: More closely examine how official Prometheus client libraries handle label key sanitization
+            // and follow whatever they do, so it's not actually clear if transforming `__foo` to `___foo` would
+            // be valid, given that it still technically starts with two underscores.
+            /*if as_chars.len() == 2 {
                 assert!(!(as_chars[0] == '_' && as_chars[1] == '_'));
             } else if as_chars.len() == 3 {
                 if as_chars[0] == '_' && as_chars[1] == '_' {
                     assert_eq!(as_chars[2], '_');
                 }
-            }
+            }*/
 
             assert!(!as_chars.iter().any(|c| invalid_label_key_character(*c)),
                 "invalid character in label key");