[incubator-kyuubi] branch branch-1.6 updated: [KYUUBI #3156] Expose REST frontend connection metrics

[fe...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

feiwang pushed a commit to branch branch-1.6
in repository https://gitbox.apache.org/repos/asf/incubator-kyuubi.git


The following commit(s) were added to refs/heads/branch-1.6 by this push:
     new d15ca5182 [KYUUBI #3156] Expose REST frontend connection metrics
d15ca5182 is described below

commit d15ca5182e22959c75663e2d8f6132fdd15f1bee
Author: Tianlin Liao <ti...@ebay.com>
AuthorDate: Thu Aug 18 17:08:14 2022 +0800

    [KYUUBI #3156] Expose REST frontend connection metrics
    
    ### _Why are the changes needed?_
    
    ### _How was this patch tested?_
    - [ ] Add some test cases that check the changes thoroughly including negative and positive cases if possible
    
    - [ ] Add screenshots for manual tests if appropriate
    
    - [ ] [Run test](https://kyuubi.apache.org/docs/latest/develop_tools/testing.html#running-tests) locally before make a pull request
    
    Closes #3233 from lightning-L/kyuubi-3156.
    
    Closes #3156
    
    7e297de7 [Tianlin Liao] [KYUUBI #3156] Expose REST frontend connection metrics
    
    Authored-by: Tianlin Liao <ti...@ebay.com>
    Signed-off-by: Fei Wang <fw...@ebay.com>
---
 .../apache/kyuubi/metrics/MetricsConstants.scala   |  5 +++++
 .../KyuubiHttpAuthenticationFactory.scala          | 15 ++++++++++++++
 .../server/api/v1/SessionsResourceSuite.scala      | 23 ++++++++++++++++++++++
 .../kyuubi/server/rest/client/BatchCliSuite.scala  | 11 +++++++++++
 .../server/rest/client/BatchRestApiSuite.scala     | 16 +++++++++++++++
 5 files changed, 70 insertions(+)

diff --git a/kyuubi-metrics/src/main/scala/org/apache/kyuubi/metrics/MetricsConstants.scala b/kyuubi-metrics/src/main/scala/org/apache/kyuubi/metrics/MetricsConstants.scala
index 342f6e38f..022992901 100644
--- a/kyuubi-metrics/src/main/scala/org/apache/kyuubi/metrics/MetricsConstants.scala
+++ b/kyuubi-metrics/src/main/scala/org/apache/kyuubi/metrics/MetricsConstants.scala
@@ -33,6 +33,7 @@ object MetricsConstants {
   final private val CONN = KYUUBI + "connection."
   final private val THRIFT_HTTP_CONN = KYUUBI + "thrift.http.connection."
   final private val THRIFT_BINARY_CONN = KYUUBI + "thrift.binary.connection."
+  final private val REST_CONN = KYUUBI + "rest.connection."
 
   final val CONN_OPEN: String = CONN + "opened"
   final val CONN_FAIL: String = CONN + "failed"
@@ -46,6 +47,10 @@ object MetricsConstants {
   final val THRIFT_BINARY_CONN_FAIL: String = THRIFT_BINARY_CONN + "failed"
   final val THRIFT_BINARY_CONN_TOTAL: String = THRIFT_BINARY_CONN + "total"
 
+  final val REST_CONN_OPEN: String = REST_CONN + "opened"
+  final val REST_CONN_FAIL: String = REST_CONN + "failed"
+  final val REST_CONN_TOTAL: String = REST_CONN + "total"
+
   final private val ENGINE = KYUUBI + "engine."
   final val ENGINE_FAIL: String = ENGINE + "failed"
   final val ENGINE_TIMEOUT: String = ENGINE + "timeout"
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/http/authentication/KyuubiHttpAuthenticationFactory.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/http/authentication/KyuubiHttpAuthenticationFactory.scala
index 16dc607b8..4556f1c82 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/http/authentication/KyuubiHttpAuthenticationFactory.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/http/authentication/KyuubiHttpAuthenticationFactory.scala
@@ -26,6 +26,8 @@ import org.eclipse.jetty.server.handler.HandlerWrapper
 
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.config.KyuubiConf.{AUTHENTICATION_METHOD, ENGINE_SECURITY_ENABLED}
+import org.apache.kyuubi.metrics.MetricsConstants.{REST_CONN_FAIL, REST_CONN_OPEN, REST_CONN_TOTAL}
+import org.apache.kyuubi.metrics.MetricsSystem
 import org.apache.kyuubi.service.authentication.{AuthTypes, InternalSecurityAccessor}
 import org.apache.kyuubi.service.authentication.AuthTypes.KERBEROS
 
@@ -51,6 +53,10 @@ class KyuubiHttpAuthenticationFactory(conf: KyuubiConf) {
             baseRequest: Request,
             request: HttpServletRequest,
             response: HttpServletResponse): Unit = {
+          MetricsSystem.tracing { ms =>
+            ms.incCount(REST_CONN_TOTAL)
+            ms.incCount(REST_CONN_OPEN)
+          }
           try {
             if (kerberosEnabled) {
               ugi.doAs(new PrivilegedAction[Unit] {
@@ -62,6 +68,15 @@ class KyuubiHttpAuthenticationFactory(conf: KyuubiConf) {
               handler.handle(target, baseRequest, request, response)
             }
           } finally {
+            val statusCode = response.getStatus
+            if (statusCode < 200 || statusCode >= 300) {
+              MetricsSystem.tracing { ms =>
+                ms.incCount(REST_CONN_FAIL)
+              }
+            }
+            MetricsSystem.tracing { ms =>
+              ms.decCount(REST_CONN_OPEN)
+            }
             AuthenticationFilter.HTTP_CLIENT_USER_NAME.remove()
             AuthenticationFilter.HTTP_CLIENT_IP_ADDRESS.remove()
           }
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/SessionsResourceSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/SessionsResourceSuite.scala
index daf92b1cd..62aeb97cb 100644
--- a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/SessionsResourceSuite.scala
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/SessionsResourceSuite.scala
@@ -23,13 +23,23 @@ import javax.ws.rs.core.{MediaType, Response}
 
 import scala.collection.JavaConverters._
 
+import org.scalatest.time.SpanSugar.convertIntToGrainOfTime
+
 import org.apache.kyuubi.{KyuubiFunSuite, RestFrontendTestHelper}
 import org.apache.kyuubi.client.api.v1.dto._
 import org.apache.kyuubi.events.KyuubiSessionEvent
+import org.apache.kyuubi.metrics.{MetricsConstants, MetricsSystem}
 import org.apache.kyuubi.operation.OperationHandle
 
 class SessionsResourceSuite extends KyuubiFunSuite with RestFrontendTestHelper {
 
+  override protected def beforeEach(): Unit = {
+    super.beforeEach()
+    eventually(timeout(10.seconds), interval(200.milliseconds)) {
+      assert(MetricsSystem.counterValue(MetricsConstants.REST_CONN_OPEN).getOrElse(0L) === 0)
+    }
+  }
+
   test("open/close and count session") {
     val requestObj = new SessionOpenRequest(
       1,
@@ -127,6 +137,11 @@ class SessionsResourceSuite extends KyuubiFunSuite with RestFrontendTestHelper {
   }
 
   test("get infoType") {
+    val totalConnections =
+      MetricsSystem.counterValue(MetricsConstants.REST_CONN_TOTAL).getOrElse(0L)
+    val failedConnections =
+      MetricsSystem.counterValue(MetricsConstants.REST_CONN_FAIL).getOrElse(0L)
+
     val requestObj = new SessionOpenRequest(
       1,
       "admin",
@@ -160,6 +175,14 @@ class SessionsResourceSuite extends KyuubiFunSuite with RestFrontendTestHelper {
     response = webTarget.path(s"api/v1/sessions/$sessionHandle/info/str")
       .request().get()
     assert(404 == response.getStatus)
+
+    eventually(timeout(3.seconds), interval(200.milliseconds)) {
+      assert(MetricsSystem.counterValue(
+        MetricsConstants.REST_CONN_TOTAL).getOrElse(0L) - totalConnections === 6)
+      assert(MetricsSystem.counterValue(MetricsConstants.REST_CONN_OPEN).getOrElse(0L) === 0)
+      assert(MetricsSystem.counterValue(
+        MetricsConstants.REST_CONN_FAIL).getOrElse(0L) - failedConnections === 4)
+    }
   }
 
   test("submit operation and get operation handle") {
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/rest/client/BatchCliSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/rest/client/BatchCliSuite.scala
index f71363b8c..d45f6f215 100644
--- a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/rest/client/BatchCliSuite.scala
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/rest/client/BatchCliSuite.scala
@@ -23,9 +23,11 @@ import java.nio.file.{Files, Paths}
 
 import org.apache.hadoop.security.UserGroupInformation
 import org.apache.hive.service.rpc.thrift.TProtocolVersion
+import org.scalatest.time.SpanSugar.convertIntToGrainOfTime
 
 import org.apache.kyuubi.{BatchTestHelper, RestClientTestHelper, Utils}
 import org.apache.kyuubi.ctl.TestPrematureExit
+import org.apache.kyuubi.metrics.{MetricsConstants, MetricsSystem}
 import org.apache.kyuubi.session.KyuubiSessionManager
 
 class BatchCliSuite extends RestClientTestHelper with TestPrematureExit with BatchTestHelper {
@@ -74,6 +76,9 @@ class BatchCliSuite extends RestClientTestHelper with TestPrematureExit with Bat
   }
 
   test("basic batch rest client") {
+    val totalConnections =
+      MetricsSystem.counterValue(MetricsConstants.REST_CONN_TOTAL).getOrElse(0L)
+
     val createArgs = Array(
       "create",
       "batch",
@@ -121,6 +126,12 @@ class BatchCliSuite extends RestClientTestHelper with TestPrematureExit with Bat
       "--password",
       ldapUserPasswd)
     result = testPrematureExitForControlCli(deleteArgs, "\"success\":true")
+
+    eventually(timeout(3.seconds), interval(200.milliseconds)) {
+      assert(MetricsSystem.counterValue(
+        MetricsConstants.REST_CONN_TOTAL).getOrElse(0L) - totalConnections === 5)
+      assert(MetricsSystem.counterValue(MetricsConstants.REST_CONN_OPEN).getOrElse(0L) === 0)
+    }
   }
 
   test("spnego batch rest client") {
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/rest/client/BatchRestApiSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/rest/client/BatchRestApiSuite.scala
index 7bf3777b8..d144d61be 100644
--- a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/rest/client/BatchRestApiSuite.scala
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/rest/client/BatchRestApiSuite.scala
@@ -19,10 +19,13 @@ package org.apache.kyuubi.server.rest.client
 
 import java.util.Base64
 
+import org.scalatest.time.SpanSugar.convertIntToGrainOfTime
+
 import org.apache.kyuubi.{BatchTestHelper, RestClientTestHelper}
 import org.apache.kyuubi.client.{BatchRestApi, KyuubiRestClient}
 import org.apache.kyuubi.client.api.v1.dto.Batch
 import org.apache.kyuubi.client.exception.KyuubiRestException
+import org.apache.kyuubi.metrics.{MetricsConstants, MetricsSystem}
 
 class BatchRestApiSuite extends RestClientTestHelper with BatchTestHelper {
 
@@ -59,6 +62,11 @@ class BatchRestApiSuite extends RestClientTestHelper with BatchTestHelper {
   }
 
   test("basic batch rest client with invalid user") {
+    val totalConnections =
+      MetricsSystem.counterValue(MetricsConstants.REST_CONN_TOTAL).getOrElse(0L)
+    val failedConnections =
+      MetricsSystem.counterValue(MetricsConstants.REST_CONN_FAIL).getOrElse(0L)
+
     val basicKyuubiRestClient: KyuubiRestClient =
       KyuubiRestClient.builder(baseUri.toString)
         .authHeaderMethod(KyuubiRestClient.AuthHeaderMethod.BASIC)
@@ -75,6 +83,14 @@ class BatchRestApiSuite extends RestClientTestHelper with BatchTestHelper {
     assert(e.getCause.toString.contains(s"Error validating LDAP user: uid=${customUser}"))
 
     basicKyuubiRestClient.close()
+
+    eventually(timeout(3.seconds), interval(200.milliseconds)) {
+      assert(MetricsSystem.counterValue(
+        MetricsConstants.REST_CONN_TOTAL).getOrElse(0L) - totalConnections === 1)
+      assert(MetricsSystem.counterValue(MetricsConstants.REST_CONN_OPEN).getOrElse(0L) === 0)
+      assert(MetricsSystem.counterValue(
+        MetricsConstants.REST_CONN_FAIL).getOrElse(0L) - failedConnections === 1)
+    }
   }
 
   test("spnego batch rest client") {