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

rong pushed a commit to branch query-memory-leak
in repository https://gitbox.apache.org/repos/asf/iotdb.git

commit 0cb9ecf6b32756b2e4571516d612c6652742214d
Author: Steve Yurong Su <[email protected]>
AuthorDate: Mon Aug 2 11:53:00 2021 +0800

    fix leak in sessionIdToStatementId & statementIdToQueryId
---
 .../apache/iotdb/db/query/control/SessionManager.java   | 17 +++++++++++------
 .../java/org/apache/iotdb/db/service/TSServiceImpl.java | 13 +++++++++----
 2 files changed, 20 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 799f8a6..c3257fb 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
@@ -26,7 +26,6 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.time.ZoneId;
-import java.util.Collections;
 import java.util.Map;
 import java.util.Set;
 import java.util.TimeZone;
@@ -78,6 +77,7 @@ public class SessionManager {
 
   public long requestSessionId(String username, String zoneId) {
     long sessionId = sessionIdGenerator.incrementAndGet();
+
     currSessionId.set(sessionId);
     sessionIdToUsername.put(sessionId, username);
     sessionIdToZoneId.put(sessionId, ZoneId.of(zoneId));
@@ -88,10 +88,15 @@ public class SessionManager {
   public boolean releaseSessionResource(long sessionId) {
     sessionIdToZoneId.remove(sessionId);
 
-    for (long statementId :
-        sessionIdToStatementId.getOrDefault(sessionId, 
Collections.emptySet())) {
-      for (long queryId : statementIdToQueryId.getOrDefault(statementId, 
Collections.emptySet())) {
-        releaseQueryResourceNoExceptions(queryId);
+    Set<Long> statementIdSet = sessionIdToStatementId.remove(sessionId);
+    if (statementIdSet != null) {
+      for (Long statementId : statementIdSet) {
+        Set<Long> queryIdSet = statementIdToQueryId.remove(statementId);
+        if (queryIdSet != null) {
+          for (Long queryId : queryIdSet) {
+            releaseQueryResourceNoExceptions(queryId);
+          }
+        }
       }
     }
 
@@ -109,7 +114,7 @@ public class SessionManager {
   public void closeStatement(long sessionId, long statementId) {
     Set<Long> queryIdSet = statementIdToQueryId.remove(statementId);
     if (queryIdSet != null) {
-      for (long queryId : queryIdSet) {
+      for (Long queryId : queryIdSet) {
         releaseQueryResourceNoExceptions(queryId);
       }
     }
diff --git 
a/server/src/main/java/org/apache/iotdb/db/service/TSServiceImpl.java 
b/server/src/main/java/org/apache/iotdb/db/service/TSServiceImpl.java
index d1bf262..7073d98 100644
--- a/server/src/main/java/org/apache/iotdb/db/service/TSServiceImpl.java
+++ b/server/src/main/java/org/apache/iotdb/db/service/TSServiceImpl.java
@@ -304,12 +304,17 @@ public class TSServiceImpl implements TSIService.Iface {
     }
 
     try {
-      if (req.isSetStatementId() && req.isSetQueryId()) {
-        sessionManager.closeDataset(req.statementId, req.queryId);
+      if (req.isSetStatementId()) {
+        if (req.isSetQueryId()) {
+          sessionManager.closeDataset(req.statementId, req.queryId);
+        } else {
+          sessionManager.closeStatement(req.sessionId, req.statementId);
+        }
+        return RpcUtils.getStatus(
+            TSStatusCode.CLOSE_OPERATION_ERROR, "statement id not set by 
client.");
       } else {
-        sessionManager.closeStatement(req.sessionId, req.statementId);
+        return RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS);
       }
-      return RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS);
     } catch (Exception e) {
       return onNPEOrUnexpectedException(
           e, "executing closeOperation", TSStatusCode.CLOSE_OPERATION_ERROR);

Reply via email to