This is an automated email from the ASF dual-hosted git repository.

xingtanzjr pushed a commit to branch test_ratis_0623
in repository https://gitbox.apache.org/repos/asf/iotdb.git


The following commit(s) were added to refs/heads/test_ratis_0623 by this push:
     new 7b6e5c83a8 spotless
7b6e5c83a8 is described below

commit 7b6e5c83a82ca4877199b2eeee5168ee10065535
Author: Jinrui.Zhang <[email protected]>
AuthorDate: Thu Jun 23 20:26:27 2022 +0800

    spotless
---
 .../iotdb/db/query/control/SessionManager.java     | 30 +++++++++++++++++-----
 .../thrift/impl/DataNodeTSIServiceImpl.java        |  9 ++++---
 2 files changed, 29 insertions(+), 10 deletions(-)

diff --git 
a/server/src/main/java/org/apache/iotdb/db/query/control/SessionManager.java 
b/server/src/main/java/org/apache/iotdb/db/query/control/SessionManager.java
index 68018310d8..f6c51a267b 100644
--- a/server/src/main/java/org/apache/iotdb/db/query/control/SessionManager.java
+++ b/server/src/main/java/org/apache/iotdb/db/query/control/SessionManager.java
@@ -157,7 +157,8 @@ public class SessionManager {
       long queryId,
       long statementId,
       boolean haveStatementId,
-      boolean haveSetQueryId) {
+      boolean haveSetQueryId,
+      Consumer<Long> releaseByQueryId) {
     if (!checkLogin(sessionId)) {
       return RpcUtils.getStatus(
           TSStatusCode.NOT_LOGIN_ERROR,
@@ -174,9 +175,9 @@ public class SessionManager {
     try {
       if (haveStatementId) {
         if (haveSetQueryId) {
-          this.closeDataset(statementId, queryId);
+          this.closeDataset(statementId, queryId, releaseByQueryId);
         } else {
-          this.closeStatement(sessionId, statementId);
+          this.closeStatement(sessionId, statementId, releaseByQueryId);
         }
         return RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS);
       } else {
@@ -189,6 +190,21 @@ public class SessionManager {
     }
   }
 
+  public TSStatus closeOperation(
+      long sessionId,
+      long queryId,
+      long statementId,
+      boolean haveStatementId,
+      boolean haveSetQueryId) {
+    return closeOperation(
+        sessionId,
+        queryId,
+        statementId,
+        haveStatementId,
+        haveSetQueryId,
+        this::releaseQueryResourceNoExceptions);
+  }
+
   /**
    * Check whether current user has logged in.
    *
@@ -261,11 +277,11 @@ public class SessionManager {
     return statementId;
   }
 
-  public void closeStatement(long sessionId, long statementId) {
+  public void closeStatement(long sessionId, long statementId, Consumer<Long> 
releaseByQueryId) {
     Set<Long> queryIdSet = statementIdToQueryId.remove(statementId);
     if (queryIdSet != null) {
       for (Long queryId : queryIdSet) {
-        releaseQueryResourceNoExceptions(queryId);
+        releaseByQueryId.accept(queryId);
       }
     }
 
@@ -377,8 +393,8 @@ public class SessionManager {
     queryIdToDataSet.remove(queryId);
   }
 
-  public void closeDataset(Long statementId, Long queryId) {
-    releaseQueryResourceNoExceptions(queryId);
+  public void closeDataset(Long statementId, Long queryId, Consumer<Long> 
releaseByQueryId) {
+    releaseByQueryId.accept(queryId);
     if (statementIdToQueryId.containsKey(statementId)) {
       statementIdToQueryId.get(statementId).remove(queryId);
     }
diff --git 
a/server/src/main/java/org/apache/iotdb/db/service/thrift/impl/DataNodeTSIServiceImpl.java
 
b/server/src/main/java/org/apache/iotdb/db/service/thrift/impl/DataNodeTSIServiceImpl.java
index 15ffeae432..b035aa05ee 100644
--- 
a/server/src/main/java/org/apache/iotdb/db/service/thrift/impl/DataNodeTSIServiceImpl.java
+++ 
b/server/src/main/java/org/apache/iotdb/db/service/thrift/impl/DataNodeTSIServiceImpl.java
@@ -218,15 +218,18 @@ public class DataNodeTSIServiceImpl implements 
TSIEventHandler {
   @Override
   public TSStatus cancelOperation(TSCancelOperationReq req) {
     // TODO implement
-    cleanupQueryExecution(req.queryId);
     return RpcUtils.getStatus(TSStatusCode.QUERY_NOT_ALLOWED, "Cancellation is 
not implemented");
   }
 
   @Override
   public TSStatus closeOperation(TSCloseOperationReq req) {
-    cleanupQueryExecution(req.queryId);
     return SESSION_MANAGER.closeOperation(
-        req.sessionId, req.queryId, req.statementId, req.isSetStatementId(), 
req.isSetQueryId());
+        req.sessionId,
+        req.queryId,
+        req.statementId,
+        req.isSetStatementId(),
+        req.isSetQueryId(),
+        this::cleanupQueryExecution);
   }
 
   @Override

Reply via email to