[KYUUBI apache#6843] FIX 'query-timeout-thread' thread leak

ASiegeLion · Dec 9, 2024 · 78a9fde · 78a9fde
1 parent dc3ac89
commit 78a9fde
Show file tree

Hide file tree

Showing 6 changed files with 12 additions and 2 deletions.
diff --git a/...ql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/PlanOnlyStatement.scala b/...ql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/PlanOnlyStatement.scala
@@ -67,7 +67,10 @@ class PlanOnlyStatement(
             statement)
           resultSet =
             ResultSetUtil.fromResultFetcher(resultFetcher, resultMaxRows, resultFetchTimeout);
-        case _ => explainOperation(statement)
+          shutdownTimeoutMonitor
+        case _ =>
+          explainOperation(statement)
+          shutdownTimeoutMonitor
       }
     } catch {
       onError()

diff --git a/...rk-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecutePython.scala b/...rk-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecutePython.scala
@@ -127,6 +127,7 @@ class ExecutePython(
           val ke =
             KyuubiSQLException("Error submitting python in background", rejected)
           setOperationException(ke)
+          shutdownTimeoutMonitor
           throw ke
       }
     } else {

diff --git a/...ark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecuteScala.scala b/...ark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecuteScala.scala
@@ -154,6 +154,7 @@ class ExecuteScala(
           val ke =
             KyuubiSQLException("Error submitting scala in background", rejected)
           setOperationException(ke)
+          shutdownTimeoutMonitor
           throw ke
       }
     } else {

diff --git a/...sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecuteStatement.scala b/...sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecuteStatement.scala
@@ -124,6 +124,7 @@ class ExecuteStatement(
           val ke =
             KyuubiSQLException("Error submitting query in background, query rejected", rejected)
           setOperationException(ke)
+          shutdownTimeoutMonitor
           throw ke
       }
     } else {

diff --git a/...ino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/ExecuteStatement.scala b/...ino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/ExecuteStatement.scala
@@ -76,6 +76,7 @@ class ExecuteStatement(
           val ke =
             KyuubiSQLException("Error submitting query in background, query rejected", rejected)
           setOperationException(ke)
+          shutdownTimeoutMonitor
           throw ke
       }
     } else {

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
@@ -172,7 +172,10 @@ class ExecuteStatement(
     try {
       val opHandle = sessionManager.submitBackgroundOperation(asyncOperation)
       setBackgroundHandle(opHandle)
-    } catch onError("submitting query in background, query rejected")
+    } catch {
+      shutdownTimeoutMonitor
+      onError("submitting query in background, query rejected")
+    }
 
     if (!shouldRunAsync) getBackgroundHandle.get()
   }