[SPARK-32026][CORE][TEST] Add PrometheusServletSuite

core/src/main/scala/org/apache/spark/metrics/sink/PrometheusServlet.scala

@@ -43,6 +43,7 @@ private[spark] class PrometheusServlet(
   extends Sink {
 
   val SERVLET_KEY_PATH = "path"
+  val DRIVER_SUBSTRING = ".driver."
 
   val servletPath = property.getProperty(SERVLET_KEY_PATH)
 
@@ -56,71 +57,84 @@ private[spark] class PrometheusServlet(
   def getMetricsSnapshot(request: HttpServletRequest): String = {
     import scala.collection.JavaConverters._
 
-    val guagesLabel = """{type="gauges"}"""
-    val countersLabel = """{type="counters"}"""
-    val metersLabel = countersLabel
-    val histogramslabels = """{type="histograms"}"""
-    val timersLabels = """{type="timers"}"""
+    def gaugesLabel(applicationId: String) =
+      s"""{application_id="${applicationId}", type="gauges"}"""
+    def countersLabel(applicationId: String) =
+      s"""{application_id="${applicationId}", type="counters"}"""
+    def metersLabel(applicationId: String) = countersLabel(applicationId)
+    def histogramslabels(applicationId: String) =
+      s"""{application_id="${applicationId}", type="histograms"}"""
+    def timersLabels(applicationId: String) =
+      s"""{application_id="${applicationId}", type="timers"}"""
 
     val sb = new StringBuilder()
     registry.getGauges.asScala.foreach { case (k, v) =>
       if (!v.getValue.isInstanceOf[String]) {
-        sb.append(s"${normalizeKey(k)}Number$guagesLabel ${v.getValue}\n")
-        sb.append(s"${normalizeKey(k)}Value$guagesLabel ${v.getValue}\n")
+        val (applicationId, prefix) = normalizeKey(k)
+        val label = gaugesLabel(applicationId)
+        sb.append(s"${prefix}Number$label ${v.getValue}\n")
+        sb.append(s"${prefix}Value$label ${v.getValue}\n")
       }
     }
     registry.getCounters.asScala.foreach { case (k, v) =>
-      sb.append(s"${normalizeKey(k)}Count$countersLabel ${v.getCount}\n")
+      val (applicationId, prefix) = normalizeKey(k)
+      val label = countersLabel(applicationId)
+      sb.append(s"${prefix}Count$label ${v.getCount}\n")
     }
     registry.getHistograms.asScala.foreach { case (k, h) =>
       val snapshot = h.getSnapshot
-      val prefix = normalizeKey(k)
-      sb.append(s"${prefix}Count$histogramslabels ${h.getCount}\n")
-      sb.append(s"${prefix}Max$histogramslabels ${snapshot.getMax}\n")
-      sb.append(s"${prefix}Mean$histogramslabels ${snapshot.getMean}\n")
-      sb.append(s"${prefix}Min$histogramslabels ${snapshot.getMin}\n")
-      sb.append(s"${prefix}50thPercentile$histogramslabels ${snapshot.getMedian}\n")
-      sb.append(s"${prefix}75thPercentile$histogramslabels ${snapshot.get75thPercentile}\n")
-      sb.append(s"${prefix}95thPercentile$histogramslabels ${snapshot.get95thPercentile}\n")
-      sb.append(s"${prefix}98thPercentile$histogramslabels ${snapshot.get98thPercentile}\n")
-      sb.append(s"${prefix}99thPercentile$histogramslabels ${snapshot.get99thPercentile}\n")
-      sb.append(s"${prefix}999thPercentile$histogramslabels ${snapshot.get999thPercentile}\n")
-      sb.append(s"${prefix}StdDev$histogramslabels ${snapshot.getStdDev}\n")
+      val (applicationId, prefix) = normalizeKey(k)
+      val label = histogramslabels(applicationId)
+      sb.append(s"${prefix}Count$label ${h.getCount}\n")
+      sb.append(s"${prefix}Max$label ${snapshot.getMax}\n")
+      sb.append(s"${prefix}Mean$label ${snapshot.getMean}\n")
+      sb.append(s"${prefix}Min$label ${snapshot.getMin}\n")
+      sb.append(s"${prefix}50thPercentile$label ${snapshot.getMedian}\n")
+      sb.append(s"${prefix}75thPercentile$label ${snapshot.get75thPercentile}\n")
+      sb.append(s"${prefix}95thPercentile$label ${snapshot.get95thPercentile}\n")
+      sb.append(s"${prefix}98thPercentile$label ${snapshot.get98thPercentile}\n")
+      sb.append(s"${prefix}99thPercentile$label ${snapshot.get99thPercentile}\n")
+      sb.append(s"${prefix}999thPercentile$label ${snapshot.get999thPercentile}\n")
+      sb.append(s"${prefix}StdDev$label ${snapshot.getStdDev}\n")
     }
     registry.getMeters.entrySet.iterator.asScala.foreach { kv =>
-      val prefix = normalizeKey(kv.getKey)
+      val (applicationId, prefix) = normalizeKey(kv.getKey)
+      val label = metersLabel(applicationId)
       val meter = kv.getValue
-      sb.append(s"${prefix}Count$metersLabel ${meter.getCount}\n")
-      sb.append(s"${prefix}MeanRate$metersLabel ${meter.getMeanRate}\n")
-      sb.append(s"${prefix}OneMinuteRate$metersLabel ${meter.getOneMinuteRate}\n")
-      sb.append(s"${prefix}FiveMinuteRate$metersLabel ${meter.getFiveMinuteRate}\n")
-      sb.append(s"${prefix}FifteenMinuteRate$metersLabel ${meter.getFifteenMinuteRate}\n")
+      sb.append(s"${prefix}Count$label ${meter.getCount}\n")
+      sb.append(s"${prefix}MeanRate$label ${meter.getMeanRate}\n")
+      sb.append(s"${prefix}OneMinuteRate$label ${meter.getOneMinuteRate}\n")
+      sb.append(s"${prefix}FiveMinuteRate$label ${meter.getFiveMinuteRate}\n")
+      sb.append(s"${prefix}FifteenMinuteRate$label ${meter.getFifteenMinuteRate}\n")
     }
     registry.getTimers.entrySet.iterator.asScala.foreach { kv =>
-      val prefix = normalizeKey(kv.getKey)
+      val (applicationId, prefix) = normalizeKey(kv.getKey)
+      val label = timersLabels(applicationId)
       val timer = kv.getValue
       val snapshot = timer.getSnapshot
-      sb.append(s"${prefix}Count$timersLabels ${timer.getCount}\n")
-      sb.append(s"${prefix}Max$timersLabels ${snapshot.getMax}\n")
-      sb.append(s"${prefix}Mean$timersLabels ${snapshot.getMax}\n")
-      sb.append(s"${prefix}Min$timersLabels ${snapshot.getMin}\n")
-      sb.append(s"${prefix}50thPercentile$timersLabels ${snapshot.getMedian}\n")
-      sb.append(s"${prefix}75thPercentile$timersLabels ${snapshot.get75thPercentile}\n")
-      sb.append(s"${prefix}95thPercentile$timersLabels ${snapshot.get95thPercentile}\n")
-      sb.append(s"${prefix}98thPercentile$timersLabels ${snapshot.get98thPercentile}\n")
-      sb.append(s"${prefix}99thPercentile$timersLabels ${snapshot.get99thPercentile}\n")
-      sb.append(s"${prefix}999thPercentile$timersLabels ${snapshot.get999thPercentile}\n")
-      sb.append(s"${prefix}StdDev$timersLabels ${snapshot.getStdDev}\n")
-      sb.append(s"${prefix}FifteenMinuteRate$timersLabels ${timer.getFifteenMinuteRate}\n")
-      sb.append(s"${prefix}FiveMinuteRate$timersLabels ${timer.getFiveMinuteRate}\n")
-      sb.append(s"${prefix}OneMinuteRate$timersLabels ${timer.getOneMinuteRate}\n")
-      sb.append(s"${prefix}MeanRate$timersLabels ${timer.getMeanRate}\n")
+      sb.append(s"${prefix}Count$label ${timer.getCount}\n")
+      sb.append(s"${prefix}Max$label ${snapshot.getMax}\n")
+      sb.append(s"${prefix}Mean$label ${snapshot.getMax}\n")
+      sb.append(s"${prefix}Min$label ${snapshot.getMin}\n")
+      sb.append(s"${prefix}50thPercentile$label ${snapshot.getMedian}\n")
+      sb.append(s"${prefix}75thPercentile$label ${snapshot.get75thPercentile}\n")
+      sb.append(s"${prefix}95thPercentile$label ${snapshot.get95thPercentile}\n")
+      sb.append(s"${prefix}98thPercentile$label ${snapshot.get98thPercentile}\n")
+      sb.append(s"${prefix}99thPercentile$label ${snapshot.get99thPercentile}\n")
+      sb.append(s"${prefix}999thPercentile$label ${snapshot.get999thPercentile}\n")
+      sb.append(s"${prefix}StdDev$label ${snapshot.getStdDev}\n")
+      sb.append(s"${prefix}FifteenMinuteRate$label ${timer.getFifteenMinuteRate}\n")
+      sb.append(s"${prefix}FiveMinuteRate$label ${timer.getFiveMinuteRate}\n")
+      sb.append(s"${prefix}OneMinuteRate$label ${timer.getOneMinuteRate}\n")
+      sb.append(s"${prefix}MeanRate$label ${timer.getMeanRate}\n")
     }
     sb.toString()
   }
 
-  private def normalizeKey(key: String): String = {
-    s"metrics_${key.replaceAll("[^a-zA-Z0-9]", "_")}_"
+  private def normalizeKey(key: String): (String, String) = {
+    val applicationId = key.substring(0, key.indexOf(DRIVER_SUBSTRING))
+    val suffix = key.substring(key.indexOf(DRIVER_SUBSTRING), key.length)
+    (applicationId, s"metrics${suffix.replaceAll("[^a-zA-Z0-9]", "_")}_")
   }
 
   override def start(): Unit = { }

core/src/test/scala/org/apache/spark/metrics/sink/PrometheusServletSuite.scala

@@ -0,0 +1,64 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.metrics.sink
+
+import java.util.Properties
+
+import scala.collection.JavaConverters._
+import com.codahale.metrics.{Gauge, MetricRegistry}
+import org.apache.spark.{SecurityManager, SparkConf, SparkFunSuite}
+import org.scalatest.PrivateMethodTester
+
+class PrometheusServletSuite extends SparkFunSuite with PrivateMethodTester {
+
+  test("PrometheusServletSuite default test") {
+    val props = new Properties
+    props.put("host", "127.0.0.1")
+    props.put("port", "12340")
+    val registry = new MetricRegistry
+    val securityMgr = new SecurityManager(new SparkConf(false))
+
+    val sink = new PrometheusServlet(props, registry, securityMgr)
+    val gauge = new Gauge[Double] {
+      override def getValue: Double = 1.23
+    }
+    sink.registry.register("gauge", gauge)
+    sink.registry.register("anothergauge", gauge)
+    sink.registry.register("streaminggauge", gauge)
+
+    val metricKeys = sink.registry.getGauges().keySet.asScala
+
+    assert(metricKeys.equals(Set("gauge", "anothergauge", "streaminggauge")),
+      "Should contain all metrics registered")
+  }
+
+  test("PrometheusServletSuite private function test") {
+    val key = "local-1592132938718.driver.LiveListenerBus.listenerProcessingTime.org.apache.spark.HeartbeatReceiver"
+    val props = new Properties
+    props.put("host", "127.0.0.1")
+    props.put("port", "12340")
+    val registry = new MetricRegistry
+    val securityMgr = new SecurityManager(new SparkConf(false))
+
+    val sink = new PrometheusServlet(props, registry, securityMgr)
+    val (applicationId, suffix) = sink invokePrivate PrivateMethod[(String, String)]('normalizeKey)(key)
+    assert(applicationId == "local-1592132938718")
+    assert(suffix == "metrics_driver_LiveListenerBus_listenerProcessingTime_org_apache_spark_HeartbeatReceiver_")
+  }
+
+}