From efe7a2ffae0c630c54460a796619493d1999c903 Mon Sep 17 00:00:00 2001
From: wforget <643348094@qq.com>
Date: Tue, 5 Nov 2024 10:55:17 +0800
Subject: [PATCH] [KYUUBI-6790] Fix engine cannot exit when gracefully stopped

---
 .../scala/org/apache/kyuubi/ha/client/ServiceDiscovery.scala   | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/ServiceDiscovery.scala b/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/ServiceDiscovery.scala
index 2968c4f961c..2b05a8a7b8c 100644
--- a/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/ServiceDiscovery.scala
+++ b/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/ServiceDiscovery.scala
@@ -66,10 +66,11 @@ abstract class ServiceDiscovery(
 
   // stop the server genteelly
   def stopGracefully(isLost: Boolean = false): Unit = {
-    val activeSessionCount = fe.be.sessionManager.getActiveUserSessionCount
+    var activeSessionCount = fe.be.sessionManager.getActiveUserSessionCount
     while (activeSessionCount > 0) {
       info(s"$activeSessionCount connection(s) are active, delay shutdown")
       Thread.sleep(TimeUnit.SECONDS.toMillis(10))
+      activeSessionCount = fe.be.sessionManager.getActiveUserSessionCount
     }
     isServerLost.set(isLost)
     gracefulShutdownLatch.countDown()