[kyuubi] branch master updated: [KYUUBI #4829] Support to expose operation metrics

[fe...@apache.org]

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

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


The following commit(s) were added to refs/heads/master by this push:
     new f20f61e68 [KYUUBI #4829] Support to expose operation metrics
f20f61e68 is described below

commit f20f61e68603660f9bc806d14d47c515ca94541a
Author: fwang12 <fw...@ebay.com>
AuthorDate: Tue May 16 13:18:46 2023 +0800

    [KYUUBI #4829] Support to expose operation metrics
    
    ### _Why are the changes needed?_
    
    We need operation metrics to know the operation details.
    
    ### _How was this patch tested?_
    - [x] Add some test cases that check the changes thoroughly including negative and positive cases if possible
    
    - [ ] Add screenshots for manual tests if appropriate
    
    - [x] [Run test](https://kyuubi.readthedocs.io/en/master/develop_tools/testing.html#running-tests) locally before make a pull request
    
    Closes #4829 from turboFei/op_metrics.
    
    Closes #4829
    
    8749405b2 [fwang12] comment
    188fbf5ae [fwang12] fix ut
    658f1f1c5 [fwang12] expose operation metrics
    
    Authored-by: fwang12 <fw...@ebay.com>
    Signed-off-by: fwang12 <fw...@ebay.com>
---
 .../apache/kyuubi/events/KyuubiOperationEvent.scala  |  7 +++++--
 .../apache/kyuubi/operation/KyuubiOperation.scala    | 15 +++++++++++++++
 .../apache/kyuubi/server/BackendServiceMetric.scala  | 12 +++++++++++-
 .../operation/KyuubiOperationPerUserSuite.scala      | 20 ++++++++++++++++++++
 4 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/events/KyuubiOperationEvent.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/events/KyuubiOperationEvent.scala
index 7147cb424..2a103213e 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/events/KyuubiOperationEvent.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/events/KyuubiOperationEvent.scala
@@ -43,6 +43,7 @@ import org.apache.kyuubi.session.KyuubiSession
  * @param sessionUser the authenticated client user
  * @param sessionType the type of the parent session
  * @param kyuubiInstance the parent session connection url
+ * @param metrics the operation metrics
  */
 case class KyuubiOperationEvent private (
     statementId: String,
@@ -58,7 +59,8 @@ case class KyuubiOperationEvent private (
     sessionId: String,
     sessionUser: String,
     sessionType: String,
-    kyuubiInstance: String) extends KyuubiEvent {
+    kyuubiInstance: String,
+    metrics: Map[String, String]) extends KyuubiEvent {
 
   // operation events are partitioned by the date when the corresponding operations are
   // created.
@@ -88,6 +90,7 @@ object KyuubiOperationEvent {
       session.handle.identifier.toString,
       session.user,
       session.sessionType.toString,
-      session.connectionUrl)
+      session.connectionUrl,
+      operation.metrics)
   }
 }
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala
index e0475394e..783581a06 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala
@@ -53,6 +53,21 @@ abstract class KyuubiOperation(session: Session) extends AbstractOperation(sessi
 
   def remoteOpHandle(): TOperationHandle = _remoteOpHandle
 
+  @volatile protected var _fetchLogCount = 0L
+  @volatile protected var _fetchResultsCount = 0L
+
+  protected[kyuubi] def increaseFetchLogCount(count: Int): Unit = {
+    _fetchLogCount += count
+  }
+
+  protected[kyuubi] def increaseFetchResultsCount(count: Int): Unit = {
+    _fetchResultsCount += count
+  }
+
+  def metrics: Map[String, String] = Map(
+    "fetchLogCount" -> _fetchLogCount.toString,
+    "fetchResultsCount" -> _fetchResultsCount.toString)
+
   protected def verifyTStatus(tStatus: TStatus): Unit = {
     ThriftUtils.verifyTStatus(tStatus)
   }
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/BackendServiceMetric.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/BackendServiceMetric.scala
index 68bf11d7f..3a92a5ad0 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/BackendServiceMetric.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/BackendServiceMetric.scala
@@ -20,7 +20,7 @@ package org.apache.kyuubi.server
 import org.apache.hive.service.rpc.thrift._
 
 import org.apache.kyuubi.metrics.{MetricsConstants, MetricsSystem}
-import org.apache.kyuubi.operation.{OperationHandle, OperationStatus}
+import org.apache.kyuubi.operation.{KyuubiOperation, OperationHandle, OperationStatus}
 import org.apache.kyuubi.operation.FetchOrientation.FetchOrientation
 import org.apache.kyuubi.service.BackendService
 import org.apache.kyuubi.session.SessionHandle
@@ -207,6 +207,16 @@ trait BackendServiceMetric extends BackendService {
         else MetricsConstants.BS_FETCH_RESULT_ROWS_RATE,
         rowsSize))
 
+      val operation = sessionManager.operationManager
+        .getOperation(operationHandle)
+        .asInstanceOf[KyuubiOperation]
+
+      if (fetchLog) {
+        operation.increaseFetchLogCount(rowsSize)
+      } else {
+        operation.increaseFetchResultsCount(rowsSize)
+      }
+
       rowSet
     }
   }
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationPerUserSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationPerUserSuite.scala
index 87a2dc7d2..f86d75247 100644
--- a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationPerUserSuite.scala
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationPerUserSuite.scala
@@ -339,4 +339,24 @@ class KyuubiOperationPerUserSuite
       }
     }
   }
+
+  test("support to expose kyuubi operation metrics") {
+    withSessionConf()(Map.empty)(Map.empty) {
+      withJdbcStatement() { statement =>
+        val uuid = UUID.randomUUID().toString
+        val query = s"select '$uuid'"
+        val res = statement.executeQuery(query)
+        assert(res.next())
+        assert(!res.next())
+
+        val operationMetrics =
+          server.backendService.sessionManager.operationManager.allOperations()
+            .map(_.asInstanceOf[KyuubiOperation])
+            .filter(_.statement == query)
+            .head.metrics
+        assert(operationMetrics.get("fetchResultsCount") == Some("1"))
+        assert(operationMetrics.get("fetchLogCount") == Some("0"))
+      }
+    }
+  }
 }