From 4b3417f21f58fc3f29bc67f8e5c8fc2420f4d49f Mon Sep 17 00:00:00 2001
From: liupeiyue <liupeiyue@yy.com>
Date: Mon, 16 Dec 2024 14:46:31 +0800
Subject: [PATCH] [KYUUBI #6843] FIX 'query-timeout-thread' thread leak

---
 .../apache/kyuubi/engine/flink/operation/FlinkOperation.scala    | 1 +
 .../apache/kyuubi/engine/flink/operation/PlanOnlyStatement.scala | 1 -
 .../scala/org/apache/kyuubi/operation/ExecuteStatement.scala     | 1 -
 .../main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala | 1 +
 4 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/FlinkOperation.scala b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/FlinkOperation.scala
index df067a888c6..95457ecaa7e 100644
--- a/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/FlinkOperation.scala
+++ b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/FlinkOperation.scala
@@ -163,6 +163,7 @@ abstract class FlinkOperation(session: Session) extends AbstractOperation(sessio
           val ke = KyuubiSQLException(s"Error operating $opType: $errMsg", e)
           setOperationException(ke)
           setState(OperationState.ERROR)
+          shutdownTimeoutMonitor()
           throw ke
         }
       }
diff --git a/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/PlanOnlyStatement.scala b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/PlanOnlyStatement.scala
index 66e54428d76..276aafb2314 100644
--- a/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/PlanOnlyStatement.scala
+++ b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/PlanOnlyStatement.scala
@@ -71,7 +71,6 @@ class PlanOnlyStatement(
           explainOperation(statement)
       }
     } catch {
-      shutdownTimeoutMonitor()
       onError()
     }
   }
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
index 27f0b1818be..46ca2e8f6c0 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/ExecuteStatement.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/ExecuteStatement.scala
@@ -173,7 +173,6 @@ class ExecuteStatement(
       val opHandle = sessionManager.submitBackgroundOperation(asyncOperation)
       setBackgroundHandle(opHandle)
     } catch {
-      shutdownTimeoutMonitor()
       onError("submitting query in background, query rejected")
     }
 
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala
index a543bddb6c0..9de1dfb2b96 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala
@@ -95,6 +95,7 @@ abstract class KyuubiOperation(session: Session) extends AbstractOperation(sessi
           }
           setOperationException(ke)
           setState(OperationState.ERROR)
+          shutdownTimeoutMonitor()
           throw ke
         }
       }