check gauge exists

apache · Jan 15, 2025 · 32dce6f · 32dce6f
1 parent fff1841
commit 32dce6f
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 2 deletions.
diff --git a/kyuubi-metrics/src/main/scala/org/apache/kyuubi/metrics/MetricsSystem.scala b/kyuubi-metrics/src/main/scala/org/apache/kyuubi/metrics/MetricsSystem.scala
@@ -20,6 +20,8 @@ package org.apache.kyuubi.metrics
 import java.lang.management.ManagementFactory
 import java.util.concurrent.TimeUnit
 
+import scala.collection.JavaConverters._
+
 import com.codahale.metrics.{Gauge, MetricRegistry, Snapshot}
 import com.codahale.metrics.jvm._
 
@@ -58,8 +60,8 @@ class MetricsSystem extends CompositeService("MetricsSystem") {
     meter.mark(value)
   }
 
-  def getGauge[T](name: String): Option[Gauge[T]] = {
-    Option(registry.gauge(name))
+  def getGauge(name: String): Option[Gauge[_]] = {
+    registry.getGauges((metricsName, _) => { metricsName == name }).asScala.map(_._2).headOption
   }
 
   def registerGauge[T](name: String, value: => T, default: T): Unit = {

diff --git a/kyuubi-metrics/src/test/scala/org/apache/kyuubi/metrics/MetricsSystemSuite.scala b/kyuubi-metrics/src/test/scala/org/apache/kyuubi/metrics/MetricsSystemSuite.scala
@@ -94,4 +94,26 @@ class MetricsSystemSuite extends KyuubiFunSuite {
     checkJsonFileMetrics(reportFile, "20181117")
     metricsSystem.stop()
   }
+
+  test("metrics - get gauge") {
+    val testContextPath = "/prometheus-metrics"
+
+    val conf = KyuubiConf()
+      .set(MetricsConf.METRICS_ENABLED, true)
+      .set(MetricsConf.METRICS_REPORTERS, Set(ReporterType.PROMETHEUS.toString))
+      .set(MetricsConf.METRICS_PROMETHEUS_PORT, 0) // random port
+      .set(MetricsConf.METRICS_PROMETHEUS_PATH, testContextPath)
+    val metricsSystem = new MetricsSystem()
+    metricsSystem.initialize(conf)
+    metricsSystem.start()
+
+    assert(metricsSystem.getGauge(MetricsConstants.THRIFT_SSL_CERT_EXPIRATION).isEmpty)
+    metricsSystem.registerGauge(
+      MetricsConstants.THRIFT_SSL_CERT_EXPIRATION,
+      () => System.currentTimeMillis(),
+      0)
+    assert(metricsSystem.getGauge(MetricsConstants.THRIFT_SSL_CERT_EXPIRATION).isDefined)
+
+    metricsSystem.stop()
+  }
 }