[KYUUBI #5952] Disconnect connections without running operations after engine maxlife time graceful period

externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/SparkSQLEngine.scala

@@ -129,7 +129,7 @@ case class SparkSQLEngine(spark: SparkSession) extends Serverable("SparkSQLEngin
         info(s"Spark engine is de-registering from engine discovery space.")
         frontendServices.flatMap(_.discoveryService).foreach(_.stop())
         while (backendService.sessionManager.getOpenSessionCount > 0) {
-          Thread.sleep(1000 * 60)
+          Thread.sleep(TimeUnit.SECONDS.toMillis(10))
         }
         info(s"Spark engine has no open session now, terminating.")
         stop()
@@ -170,8 +170,10 @@ case class SparkSQLEngine(spark: SparkSession) extends Serverable("SparkSQLEngin
     val maxLifetime = conf.get(ENGINE_SPARK_MAX_LIFETIME)
     val deregistered = new AtomicBoolean(false)
     if (maxLifetime > 0) {
+      val gracefulPeriod = conf.get(ENGINE_SPARK_MAX_LIFETIME_GRACEFUL_PERIOD)
       val checkTask: Runnable = () => {
-        if (!shutdown.get && System.currentTimeMillis() - getStartTime > maxLifetime) {
+        val elapsedTime = System.currentTimeMillis() - getStartTime
+        if (!shutdown.get && elapsedTime > maxLifetime) {
           if (deregistered.compareAndSet(false, true)) {
             info(s"Spark engine has been running for more than $maxLifetime ms," +
               s" deregistering from engine discovery space.")
@@ -182,6 +184,24 @@ case class SparkSQLEngine(spark: SparkSession) extends Serverable("SparkSQLEngin
             info(s"Spark engine has been running for more than $maxLifetime ms" +
               s" and no open session now, terminating.")
             stop()
+          } else if (gracefulPeriod > 0 && elapsedTime > maxLifetime + gracefulPeriod) {
+            backendService.sessionManager.allSessions().foreach { session =>
+              val operationCount =
+                backendService.sessionManager.operationManager.allOperations()
+                  .filter(_.getSession == session)
+                  .size
+              if (operationCount == 0) {
+                warn(s"Closing session ${session.handle.identifier} forcibly that has no" +
+                  s" operation and has been running for more than $gracefulPeriod ms after engine" +
+                  s" max lifetime.")
+                try {
+                  backendService.sessionManager.closeSession(session.handle)
+                } catch {
+                  case e: Throwable =>
+                    error(s"Error closing session ${session.handle.identifier}", e)
+                }
+              }
+            }
           }
         }
       }

externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/EtcdShareLevelSparkEngineSuite.scala

@@ -17,19 +17,13 @@
 
 package org.apache.kyuubi.engine.spark
 
-import org.apache.kyuubi.config.KyuubiConf.{ENGINE_CHECK_INTERVAL, ENGINE_SHARE_LEVEL, ENGINE_SPARK_MAX_INITIAL_WAIT, ENGINE_SPARK_MAX_LIFETIME}
 import org.apache.kyuubi.engine.ShareLevel
 import org.apache.kyuubi.engine.ShareLevel.ShareLevel
 
 trait EtcdShareLevelSparkEngineSuite
   extends ShareLevelSparkEngineTests with WithEtcdCluster {
   override def withKyuubiConf: Map[String, String] = {
-    super.withKyuubiConf ++
-      etcdConf ++ Map(
-        ENGINE_SHARE_LEVEL.key -> shareLevel.toString,
-        ENGINE_SPARK_MAX_LIFETIME.key -> "PT20s",
-        ENGINE_SPARK_MAX_INITIAL_WAIT.key -> "0",
-        ENGINE_CHECK_INTERVAL.key -> "PT5s")
+    super.withKyuubiConf ++ etcdConf
   }
 }
 

externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/ShareLevelSparkEngineTests.scala

@@ -22,6 +22,7 @@ import java.util.UUID
 import org.scalatest.concurrent.PatienceConfiguration.Timeout
 import org.scalatest.time.SpanSugar.convertIntToGrainOfTime
 
+import org.apache.kyuubi.config.KyuubiConf.{ENGINE_CHECK_INTERVAL, ENGINE_SHARE_LEVEL, ENGINE_SPARK_MAX_INITIAL_WAIT, ENGINE_SPARK_MAX_LIFETIME, ENGINE_SPARK_MAX_LIFETIME_GRACEFUL_PERIOD}
 import org.apache.kyuubi.engine.ShareLevel
 import org.apache.kyuubi.engine.ShareLevel.ShareLevel
 import org.apache.kyuubi.operation.HiveJDBCTestHelper
@@ -35,6 +36,13 @@ trait ShareLevelSparkEngineTests
   extends WithDiscoverySparkSQLEngine with HiveJDBCTestHelper {
   def shareLevel: ShareLevel
 
+  override def withKyuubiConf: Map[String, String] = super.withKyuubiConf ++ Map(
+    ENGINE_SHARE_LEVEL.key -> shareLevel.toString,
+    ENGINE_SPARK_MAX_LIFETIME.key -> "PT5s",
+    ENGINE_SPARK_MAX_INITIAL_WAIT.key -> "0",
+    ENGINE_CHECK_INTERVAL.key -> "PT2s",
+    ENGINE_SPARK_MAX_LIFETIME_GRACEFUL_PERIOD.key -> "100")
+
   override protected def jdbcUrl: String = getJdbcUrl
   override val namespace: String = {
     // for test, we always use uuid as namespace
@@ -76,4 +84,23 @@ trait ShareLevelSparkEngineTests
       }
     }
   }
+
+  test("test spark engine max life-time with graceful period") {
+    withDiscoveryClient { discoveryClient =>
+      assert(engine.getServiceState == ServiceState.STARTED)
+      assert(discoveryClient.pathExists(namespace))
+      withJdbcStatement() { _ =>
+        eventually(Timeout(30.seconds)) {
+          shareLevel match {
+            case ShareLevel.CONNECTION =>
+              assert(engine.getServiceState == ServiceState.STOPPED)
+              assert(discoveryClient.pathNonExists(namespace))
+            case _ =>
+              assert(engine.getServiceState == ServiceState.STOPPED)
+              assert(discoveryClient.pathExists(namespace))
+          }
+        }
+      }
+    }
+  }
 }

externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/ZookeeperShareLevelSparkEngineSuite.scala

@@ -17,22 +17,13 @@
 
 package org.apache.kyuubi.engine.spark
 
-import org.apache.kyuubi.config.KyuubiConf.ENGINE_CHECK_INTERVAL
-import org.apache.kyuubi.config.KyuubiConf.ENGINE_SHARE_LEVEL
-import org.apache.kyuubi.config.KyuubiConf.ENGINE_SPARK_MAX_INITIAL_WAIT
-import org.apache.kyuubi.config.KyuubiConf.ENGINE_SPARK_MAX_LIFETIME
 import org.apache.kyuubi.engine.ShareLevel
 import org.apache.kyuubi.engine.ShareLevel.ShareLevel
 
 trait ZookeeperShareLevelSparkEngineSuite
   extends ShareLevelSparkEngineTests with WithEmbeddedZookeeper {
   override def withKyuubiConf: Map[String, String] = {
-    super.withKyuubiConf ++
-      zookeeperConf ++ Map(
-        ENGINE_SHARE_LEVEL.key -> shareLevel.toString,
-        ENGINE_SPARK_MAX_LIFETIME.key -> "PT20s",
-        ENGINE_SPARK_MAX_INITIAL_WAIT.key -> "0",
-        ENGINE_CHECK_INTERVAL.key -> "PT5s")
+    super.withKyuubiConf ++ zookeeperConf
   }
 }
 

kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala

@@ -1330,6 +1330,16 @@ object KyuubiConf {
       .timeConf
       .createWithDefault(0)
 
+  val ENGINE_SPARK_MAX_LIFETIME_GRACEFUL_PERIOD: ConfigEntry[Long] =
+    buildConf("kyuubi.session.engine.spark.max.lifetime.gracefulPeriod")
+      .doc("Graceful period for Spark engine to wait the connections disconnected after reaching" +
+        " the end of life. After the graceful period, all the connections without running" +
+        " operations will be forcibly disconnected. 0 or negative means always waiting the" +
+        " connections disconnected.")
+      .version("1.8.1")
+      .timeConf
+      .createWithDefault(0)
+
   val ENGINE_SPARK_MAX_INITIAL_WAIT: ConfigEntry[Long] =
     buildConf("kyuubi.session.engine.spark.max.initial.wait")
       .doc("Max wait time for the initial connection to Spark engine. The engine will" +

kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/ServiceDiscovery.scala

@@ -17,7 +17,7 @@
 
 package org.apache.kyuubi.ha.client
 
-import java.util.concurrent.CountDownLatch
+import java.util.concurrent.{CountDownLatch, TimeUnit}
 import java.util.concurrent.atomic.AtomicBoolean
 
 import org.apache.kyuubi.Logging
@@ -68,7 +68,7 @@ abstract class ServiceDiscovery(
   def stopGracefully(isLost: Boolean = false): Unit = {
     while (fe.be.sessionManager.getOpenSessionCount > 0) {
       info(s"${fe.be.sessionManager.getOpenSessionCount} connection(s) are active, delay shutdown")
-      Thread.sleep(1000 * 60)
+      Thread.sleep(TimeUnit.SECONDS.toMillis(10))
     }
     isServerLost.set(isLost)
     gracefulShutdownLatch.countDown()

kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/etcd/EtcdDiscoveryClient.scala

@@ -53,6 +53,7 @@ import org.apache.kyuubi.ha.client.DiscoveryPaths
 import org.apache.kyuubi.ha.client.ServiceDiscovery
 import org.apache.kyuubi.ha.client.ServiceNodeInfo
 import org.apache.kyuubi.ha.client.etcd.EtcdDiscoveryClient._
+import org.apache.kyuubi.util.ThreadUtils
 
 class EtcdDiscoveryClient(conf: KyuubiConf) extends DiscoveryClient {
 
@@ -381,7 +382,12 @@ class EtcdDiscoveryClient(conf: KyuubiConf) extends DiscoveryClient {
         .filter(_.getEventType == WatchEvent.EventType.DELETE).foreach(_ => {
           warn(s"This Kyuubi instance ${instance} is now de-registered from" +
             s" ETCD. The server will be shut down after the last client session completes.")
-          serviceDiscovery.stopGracefully()
+          // for jetcd, the watcher event process might block the main thread,
+          // so start a new thread to do the de-register work as a workaround,
+          // see details in https://github.com/etcd-io/jetcd/issues/1089
+          ThreadUtils.runInNewThread("deregister-watcher-thread", isDaemon = false) {
+            serviceDiscovery.stopGracefully()
+          }
         })
     }