From 32dce6fb1d3d3231ddac85683f3579ea6001ddc7 Mon Sep 17 00:00:00 2001 From: "Wang, Fei" Date: Wed, 15 Jan 2025 09:29:36 +0800 Subject: [PATCH] check gauge exists --- .../apache/kyuubi/metrics/MetricsSystem.scala | 6 +++-- .../kyuubi/metrics/MetricsSystemSuite.scala | 22 +++++++++++++++++++ 2 files changed, 26 insertions(+), 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 index 3db6daba4b2..12732bbd88c 100644 --- 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 index bac20181ca5..d895e89640b 100644 --- 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() + } }