Cherry-pick current

apache · Sep 26, 2024 · 1ce65df · 1ce65df
1 parent 74d15eb
commit 1ce65df
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 21 deletions.
diff --git a/...b-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java b/...b-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/MPPQueryContext.java
@@ -33,10 +33,10 @@
 import org.apache.tsfile.read.filter.basic.Filter;
 
 import java.time.ZoneId;
-import java.util.HashMap;
+import java.util.HashSet;
 import java.util.LinkedList;
 import java.util.List;
-import java.util.Map;
+import java.util.Set;
 
 /**
  * This class is used to record the context of a query including QueryId, query statement, session
@@ -71,7 +71,7 @@ public class MPPQueryContext {
 
   private Filter globalTimeFilter;
 
-  private Map<SchemaLockType, Integer> acquiredLockNumMap = new HashMap<>();
+  private final Set<SchemaLockType> acquiredLocks = new HashSet<>();
 
   private boolean isExplainAnalyze = false;
 
@@ -200,16 +200,12 @@ public String getSql() {
     return sql;
   }
 
-  public Map<SchemaLockType, Integer> getAcquiredLockNumMap() {
-    return acquiredLockNumMap;
+  public Set<SchemaLockType> getAcquiredLocks() {
+    return acquiredLocks;
   }
 
-  public void addAcquiredLockNum(SchemaLockType lockType) {
-    if (acquiredLockNumMap.containsKey(lockType)) {
-      acquiredLockNumMap.put(lockType, acquiredLockNumMap.get(lockType) + 1);
-    } else {
-      acquiredLockNumMap.put(lockType, 1);
-    }
+  public boolean addAcquiredLock(final SchemaLockType lockType) {
+    return acquiredLocks.add(lockType);
   }
 
   public void generateGlobalTimeFilter(Analysis analysis) {

diff --git a/...ain/java/org/apache/iotdb/db/queryengine/plan/analyze/lock/DataNodeSchemaLockManager.java b/...ain/java/org/apache/iotdb/db/queryengine/plan/analyze/lock/DataNodeSchemaLockManager.java
@@ -21,7 +21,6 @@
 
 import org.apache.iotdb.db.queryengine.common.MPPQueryContext;
 
-import java.util.Map;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 
 public class DataNodeSchemaLockManager {
@@ -45,18 +44,16 @@ private DataNodeSchemaLockManager() {
   }
 
   public void takeReadLock(final MPPQueryContext context, final SchemaLockType lockType) {
-    locks[lockType.ordinal()].readLock().lock();
-    context.addAcquiredLockNum(lockType);
+    if (context.addAcquiredLock(lockType)) {
+      locks[lockType.ordinal()].readLock().lock();
+    }
   }
 
   public void releaseReadLock(final MPPQueryContext queryContext) {
-    if (queryContext != null && !queryContext.getAcquiredLockNumMap().isEmpty()) {
-      final Map<SchemaLockType, Integer> lockMap = queryContext.getAcquiredLockNumMap();
-      for (final Map.Entry<SchemaLockType, Integer> entry : lockMap.entrySet()) {
-        for (int i = 0; i < entry.getValue(); i++) {
-          locks[entry.getKey().ordinal()].readLock().unlock();
-        }
-      }
+    if (queryContext != null && !queryContext.getAcquiredLocks().isEmpty()) {
+      queryContext
+          .getAcquiredLocks()
+          .forEach(lockType -> locks[lockType.ordinal()].readLock().unlock());
     }
   }