[iotdb] 01/02: fix format and fix memory leak in SessionManager

[ro...@apache.org]

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

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

commit 559f1a6de4a315f18330e1ed6e911d4ad74f6c67
Author: Steve Yurong Su <ro...@apache.org>
AuthorDate: Wed Aug 4 12:02:07 2021 +0800

    fix format and fix memory leak in SessionManager
---
 client-cpp/src/test/main.cpp                            |  6 ++++--
 .../apache/iotdb/db/query/control/SessionManager.java   | 17 +++++++++++------
 .../java/org/apache/iotdb/db/service/TSServiceImpl.java | 13 +++++++++----
 3 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/client-cpp/src/test/main.cpp b/client-cpp/src/test/main.cpp
index 2256526..9476343 100644
--- a/client-cpp/src/test/main.cpp
+++ b/client-cpp/src/test/main.cpp
@@ -18,6 +18,7 @@
  */
 
 #define CATCH_CONFIG_MAIN
+
 #include <catch.hpp>
 #include "Session.h"
 
@@ -27,14 +28,15 @@ struct SessionListener : Catch::TestEventListenerBase {
 
     using TestEventListenerBase::TestEventListenerBase;
 
-    void testCaseStarting( Catch::TestCaseInfo const& testInfo ) override {
+    void testCaseStarting(Catch::TestCaseInfo const &testInfo) override {
         // Perform some setup before a test case is run
         session->open(false);
     }
 
-    void testCaseEnded( Catch::TestCaseStats const& testCaseStats ) override {
+    void testCaseEnded(Catch::TestCaseStats const &testCaseStats) override {
         // Tear-down after a test case is run
         session->close();
     }
 };
+
 CATCH_REGISTER_LISTENER( SessionListener )
\ No newline at end of file
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 dc892e1..c479953 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);
+          }
+        }
       }
     }
 
@@ -123,7 +128,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 4c45e09..6c627e7 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
@@ -308,12 +308,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.SUCCESS_STATUS);
       } else {
-        sessionManager.closeStatement(req.sessionId, req.statementId);
+        return RpcUtils.getStatus(
+            TSStatusCode.CLOSE_OPERATION_ERROR, "statement id not set by client.");
       }
-      return RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS);
     } catch (Exception e) {
       return onNPEOrUnexpectedException(
           e, "executing closeOperation", TSStatusCode.CLOSE_OPERATION_ERROR);