use addTimeoutMonitor for server-side query timeout checks

apache · Oct 18, 2023 · 72479a7 · 72479a7
1 parent c5854f7
commit 72479a7
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 2 deletions.
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala
@@ -1876,6 +1876,14 @@ object KyuubiConf {
       .checkValue(_ >= 1000, "must >= 1s if set")
       .createOptional
 
+  val OPERATION_QUERY_TIMEOUT_MONITOR_ENABLED: ConfigEntry[Boolean] =
+    buildConf("kyuubi.operation.query.timeout.monitor.enabled")
+      .doc("Whether to monitor timeout query timeout check on server side.")
+      .version("1.8.0")
+      .internal
+      .booleanConf
+      .createWithDefault(true)
+
   val OPERATION_RESULT_MAX_ROWS: ConfigEntry[Int] =
     buildConf("kyuubi.operation.result.max.rows")
       .doc("Max rows of Spark query results. Rows exceeding the limit would be ignored. " +

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/ExecuteStatement.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/ExecuteStatement.scala
@@ -25,6 +25,7 @@ import org.apache.hive.service.rpc.thrift.TOperationState._
 
 import org.apache.kyuubi.KyuubiSQLException
 import org.apache.kyuubi.config.KyuubiConf
+import org.apache.kyuubi.config.KyuubiConf.OPERATION_QUERY_TIMEOUT_MONITOR_ENABLED
 import org.apache.kyuubi.metrics.{MetricsConstants, MetricsSystem}
 import org.apache.kyuubi.operation.FetchOrientation.FETCH_NEXT
 import org.apache.kyuubi.operation.log.OperationLog
@@ -58,6 +59,10 @@ class ExecuteStatement(
     OperationLog.removeCurrentOperationLog()
   }
 
+  private val isTimeoutMonitorEnabled: Boolean = confOverlay.getOrElse[String](
+    OPERATION_QUERY_TIMEOUT_MONITOR_ENABLED.key,
+    OPERATION_QUERY_TIMEOUT_MONITOR_ENABLED.defaultValStr).toBoolean
+
   private def executeStatement(): Unit = {
     try {
       // We need to avoid executing query in sync mode, because there is no heartbeat mechanism
@@ -84,7 +89,7 @@ class ExecuteStatement(
       var lastStateUpdateTime: Long = 0L
       val stateUpdateInterval =
         session.sessionManager.getConf.get(KyuubiConf.OPERATION_STATUS_UPDATE_INTERVAL)
-      while (!isComplete) {
+      while (!isComplete && !isTerminalState(state)) {
         fetchQueryLog()
         verifyTStatus(statusResp.getStatus)
         if (statusResp.getProgressUpdateResponse != null) {
@@ -159,7 +164,15 @@ class ExecuteStatement(
   override protected def runInternal(): Unit = {
     executeStatement()
     val sessionManager = session.sessionManager
-    val asyncOperation: Runnable = () => waitStatementComplete()
+    val asyncOperation: Runnable = () =>
+      try {
+        if (isTimeoutMonitorEnabled) {
+          addTimeoutMonitor(queryTimeout)
+        }
+        waitStatementComplete()
+      } finally {
+        shutdownTimeoutMonitor()
+      }
     try {
       val opHandle = sessionManager.submitBackgroundOperation(asyncOperation)
       setBackgroundHandle(opHandle)