From 3774658831963aa04123d22be2fb3845f228f8d2 Mon Sep 17 00:00:00 2001
From: simon <zhangshiming@cvte.com>
Date: Wed, 20 Oct 2021 19:36:36 +0800
Subject: [PATCH] [KYUUBI #1258] Implement sessions/execpool/statistic

### _Why are the changes needed?_
This is a subtask of umbrella issue #KPIP-1

### _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

- [ ] [Run test](https://kyuubi.readthedocs.io/en/latest/develop_tools/testing.html#running-tests) locally before make a pull request

Closes #1259 from simon824/h1.

Closes #1258

8b92b2d2 [simon] future.cancel
7aa6416c [simon] merge to statistic
e0c5d1ab [simon] [KYUUBI #1258] Implement /sessions/execpoolsize and /sessions/execpoolactivecount
d9219801 [simon] [SUB-TASK][KPIP-1]

Authored-by: simon <zhangshiming@cvte.com>
Signed-off-by: ulysses-you <ulyssesyou@apache.org>
---
 .../server/api/v1/SessionsResource.scala      |  7 +++++
 .../org/apache/kyuubi/server/api/v1/dto.scala |  2 ++
 .../server/api/v1/SessionsResourceSuite.scala | 31 +++++++++++++++++--
 3 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/SessionsResource.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/SessionsResource.scala
index e4531d976..173329ace 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/SessionsResource.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/SessionsResource.scala
@@ -36,6 +36,13 @@ private[v1] class SessionsResource extends ApiRequestContext {
     SessionOpenCount(backendService.sessionManager.getOpenSessionCount)
   }
 
+  @GET
+  @Path("execpool/statistic")
+  def execPoolStatistic(): ExecPoolStatistic = {
+    ExecPoolStatistic(backendService.sessionManager.getExecPoolSize,
+      backendService.sessionManager.getActiveCount)
+  }
+
   @POST
   @Consumes(Array(MediaType.APPLICATION_JSON))
   def openSession(request: SessionOpenRequest): SessionHandle = {
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/dto.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/dto.scala
index 995e7e1b1..48131f322 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/dto.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/dto.scala
@@ -19,6 +19,8 @@ package org.apache.kyuubi.server.api.v1
 
 case class SessionOpenCount(openSessionCount: Int)
 
+case class ExecPoolStatistic(execPoolSize: Int, execPoolActiveCount: Int)
+
 case class SessionOpenRequest(
   protocolVersion: Int,
   user: String,
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 4f2f0f69e..dc6008616 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
@@ -22,8 +22,7 @@ import javax.ws.rs.core.MediaType
 
 import org.junit.Test
 
-import org.apache.kyuubi.server.RestApiBaseSuite
-import org.apache.kyuubi.server.RestFrontendServiceSuite
+import org.apache.kyuubi.server.{RestApiBaseSuite, RestFrontendService, RestFrontendServiceSuite}
 import org.apache.kyuubi.session.SessionHandle
 
 class SessionsResourceSuite extends RestApiBaseSuite {
@@ -84,4 +83,32 @@ class SessionsResourceSuite extends RestApiBaseSuite {
     }
   }
 
+  @Test
+  def testExecPoolStatistic: Unit = {
+    RestFrontendServiceSuite.withKyuubiRestServer {
+      (restFrontendService: RestFrontendService, _, _) =>
+
+        val sessionManager = restFrontendService.be.sessionManager
+        val future = sessionManager.submitBackgroundOperation(() => {
+          Thread.sleep(3000)
+        })
+
+        // verify the exec pool statistic
+        var response = target("api/v1/sessions/execpool/statistic").request().get()
+        val execPoolStatistic1 = response.readEntity(classOf[ExecPoolStatistic])
+        assert(execPoolStatistic1.execPoolSize == 1 && execPoolStatistic1.execPoolActiveCount == 1)
+
+        future.cancel(true)
+        response = target("api/v1/sessions/execpool/statistic").request().get()
+        val execPoolStatistic2 = response.readEntity(classOf[ExecPoolStatistic])
+        assert(execPoolStatistic2.execPoolSize == 1 && execPoolStatistic2.execPoolActiveCount == 0)
+
+        sessionManager.stop()
+        response = target("api/v1/sessions/execpool/statistic").request().get()
+        val execPoolStatistic3 = response.readEntity(classOf[ExecPoolStatistic])
+        assert(execPoolStatistic3.execPoolSize == 0 && execPoolStatistic3.execPoolActiveCount == 0)
+
+    }
+  }
+
 }