save

apache · Dec 26, 2024 · b8f227f · b8f227f
1 parent aa33521
commit b8f227f
Showing 1 changed file with 19 additions and 4 deletions.
diff --git a/kyuubi-metrics/src/main/scala/org/apache/kyuubi/metrics/PrometheusReporterService.scala b/kyuubi-metrics/src/main/scala/org/apache/kyuubi/metrics/PrometheusReporterService.scala
@@ -140,20 +140,35 @@ class PrometheusReporterService(registry: MetricRegistry)
     val metricsSnapshotWriter = new java.io.StringWriter
     val contentType = TextFormat.chooseContentType(null)
     TextFormat.writeFormat(contentType, metricsSnapshotWriter, bridgeRegistry.metricFamilySamples())
-    val labelStr = labelString(labels)
     metricsSnapshotWriter.toString.split("\n").map { line =>
       if (line.startsWith("#")) {
         line
       } else {
         line.split("\\s+", 2) match {
-          case Array(metrics, rest) => s"""$metrics${labelStr} $rest"""
+          case Array(metrics, rest) =>
+            val metricsWithNewLabels = PrometheusReporterService.applyExtraLabels(metrics, labels)
+            s"""$metricsWithNewLabels $rest"""
           case _ => line
         }
       }
     }.mkString("\n")
   }
+}
 
-  private def labelString(labels: Map[String, String]): String = {
-    labels.map { case (k, v) => s"""$k="$v"""" }.toArray.sorted.mkString("{", ",", "}")
+object PrometheusReporterService {
+  def applyExtraLabels(metrics: String, labels: Map[String, String]): String = {
+    if (labels.isEmpty) return metrics
+
+    val labelString = labels.map { case (k, v) => s"""$k="$v"""" }.toArray.sorted.mkString(",")
+    val labelStartIdx = metrics.indexOf("{")
+    val labelEndIdx = metrics.indexOf("}")
+    if (labelStartIdx > 0 && labelEndIdx > 0) {
+      metrics.substring(0, labelEndIdx).stripSuffix(",") +
+        "," +
+        labelString +
+        metrics.substring(labelEndIdx)
+    } else {
+      s"$metrics{${labelString}}"
+    }
   }
 }