Move stmtHandleChangeLock to the head when cancel KyuubiStatement

Author: ruanwenjun

kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/KyuubiStatement.java

@@ -45,7 +45,7 @@ public class KyuubiStatement implements SQLStatement, KyuubiLoggable {
   private volatile TOperationHandle stmtHandle = null;
   // This lock must be acquired before modifying or judge stmt
   // to ensure there are no concurrent accesses or race conditions.
-  private final Lock stmtHandleChangeLock = new ReentrantLock();
+  private final Lock stmtHandleAccessLock = new ReentrantLock();
   private final TSessionHandle sessHandle;
   Map<String, String> sessConf = new HashMap<>();
   private int fetchSize = DEFAULT_FETCH_SIZE;
@@ -130,11 +130,11 @@ public KyuubiStatement(
   @Override
   public void cancel() throws SQLException {
     try {
+      stmtHandleAccessLock.lock();
       checkConnection("cancel");
       if (isCancelled) {
         return;
       }
-      stmtHandleChangeLock.lock();
       if (stmtHandle != null) {
         TCancelOperationReq cancelReq = new TCancelOperationReq(stmtHandle);
         TCancelOperationResp cancelResp = client.CancelOperation(cancelReq);
@@ -146,7 +146,7 @@ public void cancel() throws SQLException {
     } catch (Exception e) {
       throw new KyuubiSQLException(e.toString(), "08S01", e);
     } finally {
-      stmtHandleChangeLock.unlock();
+      stmtHandleAccessLock.unlock();
     }
   }
 
@@ -185,7 +185,7 @@ void closeClientOperation() throws SQLException {
   @Override
   public void close() throws SQLException {
     try {
-      stmtHandleChangeLock.lock();
+      stmtHandleAccessLock.lock();
       if (isClosed) {
         return;
       }
@@ -194,7 +194,7 @@ public void close() throws SQLException {
       closeResultSet();
       isClosed = true;
     } finally {
-      stmtHandleChangeLock.unlock();
+      stmtHandleAccessLock.unlock();
     }
   }
 
@@ -324,7 +324,7 @@ private void runAsyncOnServer(String sql) throws SQLException {
   }
 
   private void runAsyncOnServer(String sql, Map<String, String> confOneTime) throws SQLException {
-    stmtHandleChangeLock.lock();
+    stmtHandleAccessLock.lock();
     try {
       checkConnection("execute");
 
@@ -360,7 +360,7 @@ private void runAsyncOnServer(String sql, Map<String, String> confOneTime) throw
         throw new KyuubiSQLException(ex.toString(), "08S01", ex);
       }
     } finally {
-      stmtHandleChangeLock.unlock();
+      stmtHandleAccessLock.unlock();
     }
   }
 

kyuubi-hive-jdbc/src/test/java/org/apache/kyuubi/jdbc/hive/KyuubiStatementTest.java

@@ -21,9 +21,8 @@
 import static org.junit.Assert.assertThrows;
 
 import java.sql.SQLException;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-import org.junit.Assert;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.atomic.AtomicReference;
 import org.junit.Test;
 
 public class KyuubiStatementTest {
@@ -60,27 +59,47 @@ public void testaddBatch() throws SQLException {
   }
 
   @Test
-  public void testThrowKyuubiSQLExceptionWhenExecuteSqlOnClosedStmt() throws SQLException {
-    KyuubiStatement stmt = new KyuubiStatement(null, null, null);
-    try {
-      ExecutorService executorService = Executors.newFixedThreadPool(2);
-      executorService.submit(
-          () -> {
-            try {
-              stmt.close();
-            } catch (SQLException e) {
-              throw new RuntimeException(e);
-            }
-          });
-      executorService.submit(
-          () -> {
-            Assert.assertEquals(
-                "Can't exectue after statement has been closed",
-                assertThrows(KyuubiSQLException.class, () -> stmt.execute("SELECT 1"))
-                    .getMessage());
-          });
-    } finally {
-      stmt.close();
+  public void testThrowKyuubiSQLExceptionWhenExecuteSqlOnClosedStmt()
+      throws SQLException, InterruptedException {
+    try (KyuubiStatement stmt = new KyuubiStatement(null, null, null)) {
+      AtomicReference<Throwable> assertionFailure = new AtomicReference<>();
+      CountDownLatch latch = new CountDownLatch(1);
+
+      Thread thread1 =
+          new Thread(
+              () -> {
+                try {
+                  latch.countDown(); // 通知 thread2 可以开始了
+                  stmt.close();
+                } catch (SQLException e) {
+                  assertionFailure.set(e);
+                }
+              });
+
+      Thread thread2 =
+          new Thread(
+              () -> {
+                try {
+                  latch.await();
+                  KyuubiSQLException ex =
+                      assertThrows(KyuubiSQLException.class, () -> stmt.execute("SELECT 1"));
+                  assertEquals("Can't execute after statement has been closed", ex.getMessage());
+                } catch (Throwable t) {
+                  // 捕获所有异常，包括断言失败
+                  assertionFailure.set(t);
+                }
+              });
+
+      thread1.start();
+      thread2.start();
+
+      thread1.join();
+      thread2.join();
+
+      // 主线程判断是否有断言或异常
+      if (assertionFailure.get() != null) {
+        throw new AssertionError("Assertion failed in thread", assertionFailure.get());
+      }
     }
   }
 }