From 63d1e84b55bcd499c41ff86c9abf5036fbbc7962 Mon Sep 17 00:00:00 2001
From: Alexandre <4717855+alexcsf@users.noreply.github.com>
Date: Mon, 5 Sep 2022 17:08:08 +0100
Subject: [PATCH] Add OpenMetrics units to metrics.

---
 .../kotlin/org/jitsi/metrics/CounterMetric.kt     |  4 +++-
 .../kotlin/org/jitsi/metrics/DoubleGaugeMetric.kt |  4 +++-
 .../kotlin/org/jitsi/metrics/LongGaugeMetric.kt   |  4 +++-
 .../src/main/kotlin/org/jitsi/metrics/Metric.kt   | 15 +++++++++++++++
 4 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/jicoco-metrics/src/main/kotlin/org/jitsi/metrics/CounterMetric.kt b/jicoco-metrics/src/main/kotlin/org/jitsi/metrics/CounterMetric.kt
index 7e37f76f..6a7e68aa 100644
--- a/jicoco-metrics/src/main/kotlin/org/jitsi/metrics/CounterMetric.kt
+++ b/jicoco-metrics/src/main/kotlin/org/jitsi/metrics/CounterMetric.kt
@@ -37,7 +37,9 @@ class CounterMetric @JvmOverloads constructor(
     internal val initialValue: Long = 0L
 ) : Metric<Long>() {
     private val counter =
-        Counter.build(name, help).namespace(namespace).create().apply { inc(initialValue.toDouble()) }
+        Counter.build(name, help).namespace(namespace)
+            .apply { setUnit() }.create()
+            .apply { inc(initialValue.toDouble()) }
 
     override fun get() = counter.get().toLong()
 
diff --git a/jicoco-metrics/src/main/kotlin/org/jitsi/metrics/DoubleGaugeMetric.kt b/jicoco-metrics/src/main/kotlin/org/jitsi/metrics/DoubleGaugeMetric.kt
index 51f4db51..3a88c899 100644
--- a/jicoco-metrics/src/main/kotlin/org/jitsi/metrics/DoubleGaugeMetric.kt
+++ b/jicoco-metrics/src/main/kotlin/org/jitsi/metrics/DoubleGaugeMetric.kt
@@ -34,7 +34,9 @@ class DoubleGaugeMetric @JvmOverloads constructor(
     /** an optional initial value for this metric */
     internal val initialValue: Double = 0.0
 ) : Metric<Double>() {
-    private val gauge = Gauge.build(name, help).namespace(namespace).create().apply { set(initialValue) }
+    private val gauge = Gauge.build(name, help).namespace(namespace)
+        .apply { setUnit() }.create()
+        .apply { set(initialValue) }
 
     override fun get() = gauge.get()
 
diff --git a/jicoco-metrics/src/main/kotlin/org/jitsi/metrics/LongGaugeMetric.kt b/jicoco-metrics/src/main/kotlin/org/jitsi/metrics/LongGaugeMetric.kt
index 8ad39d15..d8813a93 100644
--- a/jicoco-metrics/src/main/kotlin/org/jitsi/metrics/LongGaugeMetric.kt
+++ b/jicoco-metrics/src/main/kotlin/org/jitsi/metrics/LongGaugeMetric.kt
@@ -34,7 +34,9 @@ class LongGaugeMetric @JvmOverloads constructor(
     /** an optional initial value for this metric */
     internal val initialValue: Long = 0L
 ) : Metric<Long>() {
-    private val gauge = Gauge.build(name, help).namespace(namespace).create().apply { set(initialValue.toDouble()) }
+    private val gauge = Gauge.build(name, help).namespace(namespace)
+        .apply { setUnit() }.create()
+        .apply { set(initialValue.toDouble()) }
 
     override fun get() = gauge.get().toLong()
 
diff --git a/jicoco-metrics/src/main/kotlin/org/jitsi/metrics/Metric.kt b/jicoco-metrics/src/main/kotlin/org/jitsi/metrics/Metric.kt
index 7cdd3b14..833be418 100644
--- a/jicoco-metrics/src/main/kotlin/org/jitsi/metrics/Metric.kt
+++ b/jicoco-metrics/src/main/kotlin/org/jitsi/metrics/Metric.kt
@@ -16,6 +16,7 @@
 package org.jitsi.metrics
 
 import io.prometheus.client.CollectorRegistry
+import io.prometheus.client.SimpleCollector
 
 /**
  * `Metric` provides methods common to all Prometheus metric type wrappers.
@@ -44,4 +45,18 @@ sealed class Metric<T> {
      * Registers this metric with the given [CollectorRegistry] and returns it.
      */
     internal abstract fun register(registry: CollectorRegistry): Metric<T>
+
+    /**
+     * Sets the OpenMetrics format unit of this metric from its name (suffix delimited by an underscore).
+     *
+     * See: [OpenMetrics](https://github.com/OpenObservability/OpenMetrics/blob/main/specification/OpenMetrics.md#unit)
+     */
+    internal fun SimpleCollector.Builder<*, *>.setUnit() {
+        val suffix = name.substringAfterLast("_", "")
+        if (UNITS.contains(suffix)) {
+            unit(suffix)
+        }
+    }
 }
+
+private val UNITS = setOf("milliseconds", "seconds", "bits", "kilobits", "megabits", "bytes", "kilobytes", "megabytes")