save

kyuubi-metrics/src/main/scala/org/apache/kyuubi/metrics/PrometheusReporterService.scala

@@ -17,17 +17,20 @@
 
 package org.apache.kyuubi.metrics
 
+import javax.servlet.http.{HttpServlet, HttpServletRequest, HttpServletResponse}
+
 import com.codahale.metrics.MetricRegistry
 import io.prometheus.client.CollectorRegistry
 import io.prometheus.client.dropwizard.DropwizardExports
-import io.prometheus.client.exporter.MetricsServlet
+import io.prometheus.client.exporter.common.TextFormat
 import org.eclipse.jetty.server.{HttpConfiguration, HttpConnectionFactory, Server, ServerConnector}
 import org.eclipse.jetty.servlet.{ServletContextHandler, ServletHolder}
 
 import org.apache.kyuubi.KyuubiException
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.config.KyuubiConf.FRONTEND_JETTY_SEND_VERSION_ENABLED
 import org.apache.kyuubi.service.AbstractService
+import org.apache.kyuubi.util.JavaUtils
 
 class PrometheusReporterService(registry: MetricRegistry)
   extends AbstractService("PrometheusReporterService") {
@@ -38,11 +41,13 @@ class PrometheusReporterService(registry: MetricRegistry)
   private[metrics] var httpServer: Server = _
   private[metrics] var httpServerConnector: ServerConnector = _
   @volatile protected var isStarted = false
+  private var instance: String = _
 
   override def initialize(conf: KyuubiConf): Unit = {
     val port = conf.get(MetricsConf.METRICS_PROMETHEUS_PORT)
     val contextPath = conf.get(MetricsConf.METRICS_PROMETHEUS_PATH)
     val jettyVersionEnabled = conf.get(FRONTEND_JETTY_SEND_VERSION_ENABLED)
+    instance = s"${JavaUtils.findLocalInetAddress.getCanonicalHostName}:$port"
 
     val httpConf = new HttpConfiguration()
     httpConf.setSendServerVersion(jettyVersionEnabled)
@@ -56,8 +61,7 @@ class PrometheusReporterService(registry: MetricRegistry)
     httpServer.setHandler(context)
 
     new DropwizardExports(registry).register(bridgeRegistry)
-    val metricsServlet = new MetricsServlet(bridgeRegistry)
-    context.addServlet(new ServletHolder(metricsServlet), contextPath)
+    context.addServlet(new ServletHolder(createPrometheusServlet()), contextPath)
 
     super.initialize(conf)
   }
@@ -100,4 +104,42 @@ class PrometheusReporterService(registry: MetricRegistry)
       }
     }
   }
+
+  private def createPrometheusServlet(): HttpServlet = {
+    new HttpServlet {
+      override def doGet(request: HttpServletRequest, response: HttpServletResponse): Unit = {
+        try {
+          response.setContentType("text/plain;charset=utf-8")
+          response.setStatus(HttpServletResponse.SC_OK)
+          response.getWriter.print(getMetricsSnapshot())
+        } catch {
+          case e: IllegalArgumentException =>
+            response.sendError(HttpServletResponse.SC_BAD_REQUEST, e.getMessage)
+          case e: Exception =>
+            warn(s"GET ${request.getRequestURI} failed: $e", e)
+            throw e
+        }
+      }
+
+      // ensure TRACE is not supported
+      override protected def doTrace(req: HttpServletRequest, res: HttpServletResponse): Unit = {
+        res.sendError(HttpServletResponse.SC_METHOD_NOT_ALLOWED)
+      }
+    }
+  }
+
+  private def getMetricsSnapshot(): String = {
+    val stringWriter = new java.io.StringWriter
+    val contentType = TextFormat.chooseContentType(null)
+    TextFormat.writeFormat(contentType, stringWriter, bridgeRegistry.metricFamilySamples())
+    val metricsSnapshot = stringWriter.toString
+    metricsSnapshot.split("\n").map { line =>
+      if (line.startsWith("#")) {
+        line
+      } else {
+        val (name, value) = line.split("\\s+", 2)
+        s"$name{instance=\"$instance\"} $value"
+      }
+    }.mkString("\n")
+  }
 }