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,10 @@
 import static org.junit.Assert.assertThrows;
 
 import java.sql.SQLException;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
-import org.junit.Assert;
+import java.util.concurrent.Future;
 import org.junit.Test;
 
 public class KyuubiStatementTest {
@@ -60,25 +61,31 @@ public void testaddBatch() throws SQLException {
   }
 
   @Test
-  public void testThrowKyuubiSQLExceptionWhenExecuteSqlOnClosedStmt() throws SQLException {
+  public void testThrowKyuubiSQLExceptionWhenExecuteSqlOnClosedStmt()
+      throws SQLException, ExecutionException, InterruptedException {
     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());
-          });
+      Future<?> closeFuture =
+          executorService.submit(
+              () -> {
+                try {
+                  stmt.close();
+                } catch (SQLException e) {
+                  throw new RuntimeException(e);
+                }
+              });
+      Future<?> submitFuture =
+          executorService.submit(
+              () -> {
+                assertEquals(
+                    "Can't execute after statement has been closed",
+                    assertThrows(KyuubiSQLException.class, () -> stmt.execute("SELECT 1"))
+                        .getMessage());
+              });
+      closeFuture.get();
+      submitFuture.get();
+      executorService.shutdown();
     } finally {
       stmt.close();
     }