From 6757c00354f3ac8c7573b682946de1664ecd731f Mon Sep 17 00:00:00 2001
From: Shaoyuan Chen <chensy20@mails.tsinghua.edu.cn>
Date: Wed, 6 Apr 2022 00:43:44 +0800
Subject: [PATCH 1/3] fix prometheus metric name and label key sanitizer

---
 metrics-exporter-prometheus/src/formatting.rs | 28 +++++++++++++++++--
 1 file changed, 25 insertions(+), 3 deletions(-)

diff --git a/metrics-exporter-prometheus/src/formatting.rs b/metrics-exporter-prometheus/src/formatting.rs
index b0fd4b78..a78998f0 100644
--- a/metrics-exporter-prometheus/src/formatting.rs
+++ b/metrics-exporter-prometheus/src/formatting.rs
@@ -128,9 +128,31 @@ 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 ret = key
+        .chars()
+        .enumerate()
+        .map(|(idx, ch)| {
+            if idx == 0 {
+                if invalid_label_key_start_character(ch) {
+                    '_'
+                } else {
+                    ch
+                }
+            } else {
+                if invalid_label_key_character(ch) {
+                    '_'
+                } else {
+                    ch
+                }
+            }
+        })
+        .collect::<String>();
+
+    if key.starts_with("__") && !key.starts_with("___") {
+        ret = "_".to_string() + &ret;
+    }
+
+    ret
 }
 
 /// Sanitizes a label value to be valid under the Prometheus [data model].

From e30495a2959a64a596d95fa605edb6dfad9c44b2 Mon Sep 17 00:00:00 2001
From: Shaoyuan Chen <chensy20@mails.tsinghua.edu.cn>
Date: Wed, 6 Apr 2022 00:49:33 +0800
Subject: [PATCH 2/3] add some test cases

---
 metrics-exporter-prometheus/src/formatting.rs | 40 +++++++------------
 1 file changed, 15 insertions(+), 25 deletions(-)

diff --git a/metrics-exporter-prometheus/src/formatting.rs b/metrics-exporter-prometheus/src/formatting.rs
index a78998f0..42929eb6 100644
--- a/metrics-exporter-prometheus/src/formatting.rs
+++ b/metrics-exporter-prometheus/src/formatting.rs
@@ -128,31 +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_].
-    let mut ret = key
-        .chars()
-        .enumerate()
-        .map(|(idx, ch)| {
-            if idx == 0 {
-                if invalid_label_key_start_character(ch) {
-                    '_'
-                } else {
-                    ch
-                }
-            } else {
-                if invalid_label_key_character(ch) {
-                    '_'
-                } else {
-                    ch
-                }
-            }
-        })
-        .collect::<String>();
-
-    if key.starts_with("__") && !key.starts_with("___") {
-        ret = "_".to_string() + &ret;
+    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;
     }
-
-    ret
+    out
 }
 
 /// Sanitizes a label value to be valid under the Prometheus [data model].
@@ -263,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 {
@@ -279,7 +267,9 @@ mod tests {
             (":", "_"),
             ("foo_bar", "foo_bar"),
             ("1foobar", "_foobar"),
-            ("__foobar", "___foobar"),
+            ("__foobar", "__foobar"),
+            ("foo1bar2", "foo1bar2"),
+            ("123", "_23"),
         ];
 
         for (input, expected) in cases {

From 5f8f694d13ef3bcef86a1e86a722da1ce81c8cbd Mon Sep 17 00:00:00 2001
From: Toby Lawrence <tobz@users.noreply.github.com>
Date: Fri, 8 Apr 2022 16:07:50 -0400
Subject: [PATCH 3/3] Update formatting.rs

---
 metrics-exporter-prometheus/src/formatting.rs | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/metrics-exporter-prometheus/src/formatting.rs b/metrics-exporter-prometheus/src/formatting.rs
index 42929eb6..0236b1dc 100644
--- a/metrics-exporter-prometheus/src/formatting.rs
+++ b/metrics-exporter-prometheus/src/formatting.rs
@@ -341,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");