Add method for sanitizing Unit names

Signed-off-by: Fabian Stäber <fabian@fstab.de>

prometheus-metrics-model/src/main/java/io/prometheus/metrics/model/snapshots/PrometheusNaming.java

@@ -20,6 +20,11 @@ public class PrometheusNaming {
      */
     private static final Pattern LABEL_NAME_PATTERN = Pattern.compile("^[a-zA-Z_.][a-zA-Z0-9_.]*$");
 
+    /**
+     * Legal characters for unit names, including dot.
+     */
+    private static final Pattern UNIT_NAME_PATTERN = Pattern.compile("^[a-zA-Z0-9_.:]+$");
+
     /**
      * According to OpenMetrics {@code _count} and {@code _sum} (and {@code _gcount}, {@code _gsum}) should also be
      * reserved metric name suffixes. However, popular instrumentation libraries have Gauges with names
@@ -83,6 +88,32 @@ public static boolean isValidLabelName(String name) {
                 !(name.startsWith("__") || name.startsWith("._") || name.startsWith("..") || name.startsWith("_."));
     }
 
+    /**
+     * Units may not have illegal characters, and they may not end with a reserved suffix like 'total'.
+     */
+    public static boolean isValidUnitName(String name) {
+        return validateUnitName(name) == null;
+    }
+
+    /**
+     * Same as {@link #isValidUnitName(String)} but returns an error message.
+     */
+    public static String validateUnitName(String name) {
+        if (name.isEmpty()) {
+            return "The unit name must not be empty.";
+        }
+        for (String reservedSuffix : RESERVED_METRIC_NAME_SUFFIXES) {
+            String suffixName = reservedSuffix.substring(1);
+            if (name.endsWith(suffixName)) {
+                return suffixName + " is a reserved suffix in Prometheus";
+            }
+        }
+        if (!UNIT_NAME_PATTERN.matcher(name).matches()) {
+            return "The unit name contains unsupported characters";
+        }
+        return null;
+    }
+
     /**
      * Get the metric or label name that is used in Prometheus exposition format.
      *
@@ -149,6 +180,42 @@ public static String sanitizeLabelName(String labelName) {
         return sanitizedName;
     }
 
+    /**
+     * Convert an arbitrary string to a name where {@link #isValidUnitName(String) isValidUnitName(name)} is true.
+     *
+     * @throws IllegalArgumentException if the {@code unitName} cannot be converted, for example if you call {@code sanitizeUnitName("total")} or {@code sanitizeUnitName("")}.
+     * @throws NullPointerException if {@code unitName} is null.
+     */
+    public static String sanitizeUnitName(String unitName) {
+        if (unitName.isEmpty()) {
+            throw new IllegalArgumentException("Cannot convert an empty string to a valid unit name.");
+        }
+        String sanitizedName = replaceIllegalCharsInUnitName(unitName);
+        boolean modified = true;
+        while (modified) {
+            modified = false;
+            while (sanitizedName.startsWith("_") || sanitizedName.startsWith(".")) {
+                sanitizedName = sanitizedName.substring(1);
+                modified = true;
+            }
+            while (sanitizedName.endsWith(".") || sanitizedName.endsWith("_")) {
+                sanitizedName = sanitizedName.substring(0, sanitizedName.length()-1);
+                modified = true;
+            }
+            for (String reservedSuffix : RESERVED_METRIC_NAME_SUFFIXES) {
+                String suffixName = reservedSuffix.substring(1);
+                if (sanitizedName.endsWith(suffixName)) {
+                    sanitizedName = sanitizedName.substring(0, sanitizedName.length() - suffixName.length());
+                    modified = true;
+                }
+            }
+        }
+        if (sanitizedName.isEmpty()) {
+            throw new IllegalArgumentException("Cannot convert '" + unitName + "' into a valid unit name.");
+        }
+        return sanitizedName;
+    }
+
     /**
      * Returns a string that matches {@link #METRIC_NAME_PATTERN}.
      */
@@ -189,4 +256,25 @@ private static String replaceIllegalCharsInLabelName(String name) {
         }
         return new String(sanitized);
     }
+
+    /**
+     * Returns a string that matches {@link #UNIT_NAME_PATTERN}.
+     */
+    private static String replaceIllegalCharsInUnitName(String name) {
+        int length = name.length();
+        char[] sanitized = new char[length];
+        for (int i = 0; i < length; i++) {
+            char ch = name.charAt(i);
+            if (ch == ':' ||
+                    ch == '.' ||
+                    (ch >= 'a' && ch <= 'z') ||
+                    (ch >= 'A' && ch <= 'Z') ||
+                    (ch >= '0' && ch <= '9')) {
+                sanitized[i] = ch;
+            } else {
+                sanitized[i] = '_';
+            }
+        }
+        return new String(sanitized);
+    }
 }

prometheus-metrics-model/src/main/java/io/prometheus/metrics/model/snapshots/Unit.java

@@ -28,10 +28,12 @@ public Unit(String name) {
         if (name == null) {
             throw new NullPointerException("Unit name cannot be null.");
         }
-        if (name.trim().isEmpty()) {
-            throw new IllegalArgumentException("Unit name cannot be empty.");
+        name = name.trim();
+        String error = PrometheusNaming.validateUnitName(name);
+        if (error != null) {
+            throw new IllegalArgumentException(name + ": Illegal unit name: " + error);
         }
-        this.name = name.trim();
+        this.name = name;
     }
 
     @Override

prometheus-metrics-model/src/test/java/io/prometheus/metrics/model/snapshots/PrometheusNamingTest.java

@@ -3,9 +3,7 @@
 import org.junit.Assert;
 import org.junit.Test;
 
-import static io.prometheus.metrics.model.snapshots.PrometheusNaming.prometheusName;
-import static io.prometheus.metrics.model.snapshots.PrometheusNaming.sanitizeLabelName;
-import static io.prometheus.metrics.model.snapshots.PrometheusNaming.sanitizeMetricName;
+import static io.prometheus.metrics.model.snapshots.PrometheusNaming.*;
 
 public class PrometheusNamingTest {
 
@@ -19,6 +17,8 @@ public void testSanitizeMetricName() {
         Assert.assertEquals("jvm", sanitizeMetricName("jvm_info"));
         Assert.assertEquals("jvm", sanitizeMetricName("jvm.info"));
         Assert.assertEquals("a.b", sanitizeMetricName("a.b"));
+        Assert.assertEquals("total", sanitizeMetricName("_total"));
+        Assert.assertEquals("total", sanitizeMetricName("total"));
     }
 
     @Test
@@ -31,4 +31,46 @@ public void testSanitizeLabelName() {
         Assert.assertEquals("abc.def", sanitizeLabelName("abc.def"));
         Assert.assertEquals("abc.def2", sanitizeLabelName("abc.def2"));
     }
+
+    @Test
+    public void testValidateUnitName() {
+        Assert.assertNotNull(validateUnitName("secondstotal"));
+        Assert.assertNotNull(validateUnitName("total"));
+        Assert.assertNotNull(validateUnitName("seconds_total"));
+        Assert.assertNotNull(validateUnitName("_total"));
+        Assert.assertNotNull(validateUnitName(""));
+
+        Assert.assertNull(validateUnitName("seconds"));
+        Assert.assertNull(validateUnitName("2"));
+    }
+
+    @Test
+    public void testSanitizeUnitName() {
+        Assert.assertEquals("seconds", sanitizeUnitName("seconds"));
+        Assert.assertEquals("seconds", sanitizeUnitName("seconds_total"));
+        Assert.assertEquals("seconds", sanitizeUnitName("seconds_total_total"));
+        Assert.assertEquals("m_s", sanitizeUnitName("m/s"));
+        Assert.assertEquals("seconds", sanitizeUnitName("secondstotal"));
+        Assert.assertEquals("2", sanitizeUnitName("2"));
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testInvalidUnitName1() {
+        sanitizeUnitName("total");
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testInvalidUnitName2() {
+        sanitizeUnitName("_total");
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testInvalidUnitName3() {
+        sanitizeUnitName("%");
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testEmptyUnitName() {
+        sanitizeUnitName("");
+    }
 }