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

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

commit bfafdfe75ab1af331daae41de6a1a5019c34efe7
Author: OneSizeFitQuorum <[email protected]>
AuthorDate: Wed Apr 23 11:43:11 2025 +0800

    finish
    
    Signed-off-by: OneSizeFitQuorum <[email protected]>
---
 .../plan/analyze/ClusterPartitionFetcher.java      | 19 ++++---
 .../analyze/cache/partition/PartitionCache.java    | 62 ++++++++++++++--------
 .../plan/analyze/cache/PartitionCacheTest.java     |  5 +-
 3 files changed, 54 insertions(+), 32 deletions(-)

diff --git 
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/ClusterPartitionFetcher.java
 
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/ClusterPartitionFetcher.java
index ae92d6d7240..04bf26eba0e 100644
--- 
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/ClusterPartitionFetcher.java
+++ 
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/ClusterPartitionFetcher.java
@@ -60,7 +60,7 @@ import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
-import java.util.LinkedList;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -490,16 +490,17 @@ public class ClusterPartitionFetcher implements 
IPartitionFetcher {
         new HashMap<>();
     for (final Map.Entry<String, Map<TSeriesPartitionSlot, TConsensusGroupId>> 
entry1 :
         schemaPartitionTableResp.getSchemaPartitionTable().entrySet()) {
+      String database = entry1.getKey();
       final Map<TSeriesPartitionSlot, TRegionReplicaSet> result1 =
-          regionReplicaMap.computeIfAbsent(entry1.getKey(), k -> new 
HashMap<>());
+          regionReplicaMap.computeIfAbsent(database, k -> new HashMap<>());
+      List<TRegionReplicaSet> regionReplicaSets =
+          partitionCache.getRegionReplicaSet(new 
ArrayList<>(entry1.getValue().values()));
+      Iterator<TRegionReplicaSet> iterator = regionReplicaSets.iterator();
       for (final Map.Entry<TSeriesPartitionSlot, TConsensusGroupId> entry2 :
           entry1.getValue().entrySet()) {
-        final TSeriesPartitionSlot seriesPartitionSlot = entry2.getKey();
-        final TConsensusGroupId consensusGroupId = entry2.getValue();
-        result1.put(seriesPartitionSlot, 
partitionCache.getRegionReplicaSet(consensusGroupId));
+        result1.put(entry2.getKey(), iterator.next());
       }
     }
-
     return new SchemaPartition(
         regionReplicaMap,
         
IoTDBDescriptor.getInstance().getConfig().getSeriesPartitionExecutorClass(),
@@ -530,10 +531,8 @@ public class ClusterPartitionFetcher implements 
IPartitionFetcher {
             result1.computeIfAbsent(entry2.getKey(), k -> new HashMap<>());
         for (final Map.Entry<TTimePartitionSlot, List<TConsensusGroupId>> 
entry3 :
             entry2.getValue().entrySet()) {
-          final List<TRegionReplicaSet> regionReplicaSets = new LinkedList<>();
-          for (final TConsensusGroupId consensusGroupId : entry3.getValue()) {
-            
regionReplicaSets.add(partitionCache.getRegionReplicaSet(consensusGroupId));
-          }
+          final List<TRegionReplicaSet> regionReplicaSets =
+              partitionCache.getRegionReplicaSet(entry3.getValue());
           result2.put(entry3.getKey(), regionReplicaSets);
         }
       }
diff --git 
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/cache/partition/PartitionCache.java
 
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/cache/partition/PartitionCache.java
index 15828538c7f..1789d587af3 100644
--- 
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/cache/partition/PartitionCache.java
+++ 
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/cache/partition/PartitionCache.java
@@ -72,7 +72,7 @@ import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
-import java.util.LinkedList;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -537,26 +537,27 @@ public class PartitionCache {
   /**
    * get regionReplicaSet from local and configNode
    *
-   * @param consensusGroupId the id of consensus group
-   * @return regionReplicaSet
+   * @param consensusGroupIds the ids of consensus group
+   * @return List<regionReplicaSet>
    * @throws RuntimeException if failed to get regionReplicaSet from configNode
    * @throws StatementAnalyzeException if there are exception when try to get 
latestRegionRouteMap
    */
-  public TRegionReplicaSet getRegionReplicaSet(TConsensusGroupId 
consensusGroupId) {
-    TRegionReplicaSet result;
+  public List<TRegionReplicaSet> getRegionReplicaSet(List<TConsensusGroupId> 
consensusGroupIds) {
+    List<TRegionReplicaSet> result;
     // try to get regionReplicaSet from cache
     regionReplicaSetLock.readLock().lock();
     try {
-      result = groupIdToReplicaSetMap.get(consensusGroupId);
+      result = getRegionReplicaSetInternal(consensusGroupIds);
     } finally {
       regionReplicaSetLock.readLock().unlock();
     }
-    if (result == null) {
+    if (result.isEmpty()) {
       // if not hit then try to get regionReplicaSet from configNode
       regionReplicaSetLock.writeLock().lock();
       try {
-        // verify that there are not hit in cache
-        if (!groupIdToReplicaSetMap.containsKey(consensusGroupId)) {
+        // double check after getting the write lock
+        result = getRegionReplicaSetInternal(consensusGroupIds);
+        if (result.isEmpty()) {
           try (ConfigNodeClient client =
               
configNodeClientManager.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) {
             TRegionRouteMapResp resp = client.getLatestRegionRouteMap();
@@ -568,18 +569,18 @@ public class PartitionCache {
                   resp.getStatus(),
                   resp.getRegionRouteMap());
             }
+            result = getRegionReplicaSetInternal(consensusGroupIds);
             // if configNode don't have then will throw RuntimeException
-            if (!groupIdToReplicaSetMap.containsKey(consensusGroupId)) {
+            if (result.isEmpty()) {
               // failed to get RegionReplicaSet from configNode
               throw new RuntimeException(
-                  "Failed to get replicaSet of consensus group[id= " + 
consensusGroupId + "]");
+                  "Failed to get replicaSet of consensus groups[ids= " + 
consensusGroupIds + "]");
             }
           } catch (ClientManagerException | TException e) {
             throw new StatementAnalyzeException(
                 "An error occurred when executing getRegionReplicaSet():" + 
e.getMessage());
           }
         }
-        result = groupIdToReplicaSetMap.get(consensusGroupId);
       } finally {
         regionReplicaSetLock.writeLock().unlock();
       }
@@ -588,6 +589,20 @@ public class PartitionCache {
     return result;
   }
 
+  private List<TRegionReplicaSet> getRegionReplicaSetInternal(
+      List<TConsensusGroupId> consensusGroupIds) {
+    List<TRegionReplicaSet> result = new ArrayList<>(consensusGroupIds.size());
+    for (TConsensusGroupId groupId : consensusGroupIds) {
+      TRegionReplicaSet replicaSet = groupIdToReplicaSetMap.get(groupId);
+      if (replicaSet != null) {
+        result.add(replicaSet);
+      } else {
+        return Collections.emptyList();
+      }
+    }
+    return result;
+  }
+
   /**
    * update regionReplicaSetMap according to timestamp
    *
@@ -660,6 +675,8 @@ public class PartitionCache {
         final Map<TSeriesPartitionSlot, TConsensusGroupId> map =
             schemaPartitionTable.getSchemaPartitionMap();
         // check cache for each device
+        List<TSeriesPartitionSlot> seriesPartitionSlots = new ArrayList<>();
+        List<TConsensusGroupId> consensusGroupIds = new ArrayList<>();
         for (final IDeviceID device : entry.getValue()) {
           final TSeriesPartitionSlot seriesPartitionSlot =
               partitionExecutor.getSeriesPartitionSlot(device);
@@ -672,9 +689,12 @@ public class PartitionCache {
             cacheMetrics.record(false, 
CacheMetrics.SCHEMA_PARTITION_CACHE_NAME);
             return null;
           }
-          final TConsensusGroupId consensusGroupId = 
map.get(seriesPartitionSlot);
-          final TRegionReplicaSet regionReplicaSet = 
getRegionReplicaSet(consensusGroupId);
-          regionReplicaSetMap.put(seriesPartitionSlot, regionReplicaSet);
+          seriesPartitionSlots.add(seriesPartitionSlot);
+          consensusGroupIds.add(map.get(seriesPartitionSlot));
+        }
+        List<TRegionReplicaSet> replicaSets = 
getRegionReplicaSet(consensusGroupIds);
+        for (int i = 0; i < replicaSets.size(); i++) {
+          regionReplicaSetMap.put(seriesPartitionSlots.get(i), 
replicaSets.get(i));
         }
       }
       logger.debug("[{} Cache] hit", CacheMetrics.SCHEMA_PARTITION_CACHE_NAME);
@@ -710,9 +730,13 @@ public class PartitionCache {
           new HashMap<>();
       Map<TSeriesPartitionSlot, TRegionReplicaSet> regionReplicaSetMap =
           schemaPartitionMap.computeIfAbsent(database, k -> new HashMap<>());
+      List<TRegionReplicaSet> regionReplicaSets =
+          getRegionReplicaSet(
+              new 
ArrayList<>(schemaPartitionTable.getSchemaPartitionMap().values()));
+      Iterator<TRegionReplicaSet> iterator = regionReplicaSets.iterator();
       for (Map.Entry<TSeriesPartitionSlot, TConsensusGroupId> entry :
           schemaPartitionTable.getSchemaPartitionMap().entrySet()) {
-        regionReplicaSetMap.put(entry.getKey(), 
getRegionReplicaSet(entry.getValue()));
+        regionReplicaSetMap.put(entry.getKey(), iterator.next());
       }
       logger.debug("[{} Cache] hit", CacheMetrics.SCHEMA_PARTITION_CACHE_NAME);
       // cache hit
@@ -907,11 +931,7 @@ public class PartitionCache {
           timePartitionSlot);
       return false;
     }
-    List<TRegionReplicaSet> regionReplicaSets = new LinkedList<>();
-    for (TConsensusGroupId consensusGroupId : cacheConsensusGroupId) {
-      regionReplicaSets.add(getRegionReplicaSet(consensusGroupId));
-    }
-    timePartitionSlotListMap.put(timePartitionSlot, regionReplicaSets);
+    timePartitionSlotListMap.put(timePartitionSlot, 
getRegionReplicaSet(cacheConsensusGroupId));
     return true;
   }
 
diff --git 
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/analyze/cache/PartitionCacheTest.java
 
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/analyze/cache/PartitionCacheTest.java
index c198c3971b4..e08a2b610c9 100644
--- 
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/analyze/cache/PartitionCacheTest.java
+++ 
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/analyze/cache/PartitionCacheTest.java
@@ -276,7 +276,10 @@ public class PartitionCacheTest {
 
   private void checkRegionReplicaSet(TConsensusGroupId consensusGroupId) {
     try {
-      assertNotNull(partitionCache.getRegionReplicaSet(consensusGroupId));
+      List<TRegionReplicaSet> regionReplicaSets =
+          
partitionCache.getRegionReplicaSet(Collections.singletonList(consensusGroupId));
+      assertEquals(1, regionReplicaSets.size());
+      assertNotNull(regionReplicaSets.get(0));
     } catch (Exception e) {
       fail(e.getMessage());
     }

Reply via email to