Optimize ConfigNode region lookups and reduce metadata copying

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/LoadManager.java

@@ -501,6 +501,15 @@ public void removeRegionGroupRelatedCache(TConsensusGroupId consensusGroupId) {
     eventService.checkAndBroadcastConsensusGroupStatisticsChangeEventIfNecessary();
   }
 
+  /**
+   * Get the latest leader of the specified RegionGroup.
+   *
+   * @return the leader ID, or -1 if the cache is absent or no leader has been selected
+   */
+  public int getRegionLeaderId(TConsensusGroupId regionGroupId) {
+    return loadCache.getRegionLeaderId(regionGroupId);
+  }
+
   /**
    * Get the latest RegionLeaderMap.
    *

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/cache/LoadCache.java

@@ -799,6 +799,16 @@ public void removeRegionGroupCache(TConsensusGroupId consensusGroupId) {
     consensusGroupCacheMap.remove(consensusGroupId);
   }
 
+  /**
+   * Get the latest leader of the specified RegionGroup.
+   *
+   * @return the leader ID, or -1 if the cache is absent or no leader has been selected
+   */
+  public int getRegionLeaderId(TConsensusGroupId regionGroupId) {
+    ConsensusGroupCache cache = consensusGroupCacheMap.get(regionGroupId);
+    return cache == null ? ConsensusGroupCache.UN_READY_LEADER_ID : cache.getLeaderId();
+  }
+
   /**
    * Safely get the latest RegionLeaderMap.
    *

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/partition/PartitionInfo.java

@@ -791,11 +791,9 @@ public List<TRegionReplicaSet> getAllReplicaSets(int dataNodeId) {
    */
   public List<TRegionReplicaSet> getReplicaSets(
       String database, List<TConsensusGroupId> regionGroupIds) {
-    if (databasePartitionTables.containsKey(database)) {
-      return databasePartitionTables.get(database).getReplicaSets(regionGroupIds);
-    } else {
-      return Collections.emptyList();
-    }
+    // A database can be deleted while a Region lookup is in progress.
+    DatabasePartitionTable table = databasePartitionTables.get(database);
+    return table == null ? Collections.emptyList() : table.getReplicaSets(regionGroupIds);
   }
 
   /**

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/env/RegionMaintainHandler.java

@@ -53,6 +53,7 @@
 import org.apache.iotdb.confignode.manager.ConfigManager;
 import org.apache.iotdb.confignode.manager.load.LoadManager;
 import org.apache.iotdb.confignode.manager.load.cache.consensus.ConsensusGroupHeartbeatSample;
+import org.apache.iotdb.confignode.manager.partition.PartitionManager;
 import org.apache.iotdb.confignode.procedure.exception.ProcedureException;
 import org.apache.iotdb.confignode.rpc.thrift.TCreatePipeReq;
 import org.apache.iotdb.consensus.pipe.consensuspipe.ConsensusPipeName;
@@ -72,6 +73,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 
@@ -771,9 +773,17 @@ public List<TDataNodeLocation> findRegionLocations(TConsensusGroupId regionId) {
   }
 
   public Optional<TRegionReplicaSet> getRegionReplicaSet(TConsensusGroupId regionId) {
-    return configManager.getPartitionManager().getAllReplicaSets().stream()
-        .filter(rg -> rg.regionId.equals(regionId))
-        .findAny();
+    if (regionId == null) {
+      return Optional.empty();
+    }
+    PartitionManager partitionManager = configManager.getPartitionManager();
+    String database = partitionManager.getRegionDatabase(regionId);
+    if (database == null) {
+      return Optional.empty();
+    }
+    // Locate by Region ID before copying: unrelated replica sets need not be locked or copied.
+    return partitionManager.getReplicaSets(database, Collections.singletonList(regionId)).stream()
+        .findFirst();
   }
 
   public String getRegionReplicaSetString(TConsensusGroupId regionId) {
@@ -931,11 +941,15 @@ public Optional<TDataNodeLocation> filterDataNodeWithOtherRegionReplica(
       return Optional.empty();
     }
 
-    // Choosing the RUNNING DataNodes to execute firstly
-    // If all DataNodes are not RUNNING, then choose the REMOVING DataNodes secondly
+    Set<Integer> regionDataNodeIds =
+        regionLocations.stream().map(TDataNodeLocation::getDataNodeId).collect(Collectors.toSet());
+    // Filter IDs before reading configurations, so unrelated DataNodes are not deep-copied.
     List<TDataNodeLocation> aliveDataNodes =
-        configManager.getNodeManager().filterDataNodeThroughStatus(allowingStatus).stream()
-            .map(TDataNodeConfiguration::getLocation)
+        configManager.getLoadManager().filterDataNodeThroughStatus(allowingStatus).stream()
+            .filter(regionDataNodeIds::contains)
+            .map(
+                nodeId ->
+                    configManager.getNodeManager().getRegisteredDataNode(nodeId).getLocation())
             .collect(Collectors.toList());
     final int leaderId = getRegionLeaderId(regionId);
     Collections.shuffle(aliveDataNodes);
@@ -962,11 +976,11 @@ public Optional<TDataNodeLocation> filterDataNodeWithOtherRegionReplica(
    */
   private int getRegionLeaderId(TConsensusGroupId regionId) {
     final LoadManager loadManager = configManager.getLoadManager();
-    Integer leaderId = loadManager.getRegionLeaderMap().get(regionId);
-    if (leaderId == null || leaderId == -1) {
+    int leaderId = loadManager.getRegionLeaderId(regionId);
+    if (leaderId == -1) {
       loadManager.waitForRegionGroupReady(Collections.singletonList(regionId));
-      leaderId = loadManager.getRegionLeaderMap().get(regionId);
+      leaderId = loadManager.getRegionLeaderId(regionId);
     }
-    return leaderId == null ? -1 : leaderId;
+    return leaderId;
   }
 }

iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/load/cache/ConsensusGroupCacheTest.java

@@ -18,14 +18,42 @@
  */
 package org.apache.iotdb.confignode.manager.load.cache;
 
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId;
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType;
 import org.apache.iotdb.confignode.manager.load.cache.consensus.ConsensusGroupCache;
 import org.apache.iotdb.confignode.manager.load.cache.consensus.ConsensusGroupHeartbeatSample;
 
 import org.junit.Assert;
 import org.junit.Test;
 
+import java.util.Collections;
+
 public class ConsensusGroupCacheTest {
 
+  @Test
+  public void regionLeaderLookupTest() {
+    LoadCache loadCache = new LoadCache();
+    TConsensusGroupId dataRegion = new TConsensusGroupId(TConsensusGroupType.DataRegion, 1);
+    TConsensusGroupId schemaRegion = new TConsensusGroupId(TConsensusGroupType.SchemaRegion, 1);
+    Assert.assertEquals(-1, loadCache.getRegionLeaderId(dataRegion));
+    Assert.assertEquals(-1, loadCache.getRegionLeaderId(schemaRegion));
+
+    loadCache.createRegionGroupHeartbeatCache("root.test", dataRegion, Collections.singleton(2));
+    loadCache.createRegionGroupHeartbeatCache("root.test", schemaRegion, Collections.singleton(3));
+    Assert.assertEquals(-1, loadCache.getRegionLeaderId(dataRegion));
+    Assert.assertEquals(-1, loadCache.getRegionLeaderId(schemaRegion));
+
+    loadCache.cacheConsensusSample(dataRegion, new ConsensusGroupHeartbeatSample(1L, 2));
+    loadCache.cacheConsensusSample(schemaRegion, new ConsensusGroupHeartbeatSample(1L, 3));
+    loadCache.updateConsensusGroupStatistics();
+    Assert.assertEquals(2, loadCache.getRegionLeaderId(dataRegion));
+    Assert.assertEquals(3, loadCache.getRegionLeaderId(schemaRegion));
+
+    loadCache.removeRegionGroupCache(dataRegion);
+    Assert.assertEquals(-1, loadCache.getRegionLeaderId(dataRegion));
+    Assert.assertEquals(3, loadCache.getRegionLeaderId(schemaRegion));
+  }
+
   @Test
   public void periodicUpdateTest() {
     ConsensusGroupCache consensusGroupCache = new ConsensusGroupCache();

iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/procedure/env/RegionMaintainHandlerTest.java

@@ -91,6 +91,7 @@ public static Collection<Object[]> regionTypes() {
 
   private TConsensusGroupId regionId;
   private LoadManager loadManager;
+  private NodeManager nodeManager;
   private RegionMaintainHandler handler;
   private String originalConsensusProtocol;
 
@@ -104,25 +105,26 @@ public void setUp() {
 
     regionId = new TConsensusGroupId(regionType, 1);
     ConfigManager configManager = mock(ConfigManager.class);
-    NodeManager nodeManager = mock(NodeManager.class);
+    nodeManager = mock(NodeManager.class);
     PartitionManager partitionManager = mock(PartitionManager.class);
     loadManager = mock(LoadManager.class);
     when(configManager.getNodeManager()).thenReturn(nodeManager);
     when(configManager.getPartitionManager()).thenReturn(partitionManager);
     when(configManager.getLoadManager()).thenReturn(loadManager);
-    when(partitionManager.getAllReplicaSets())
+    when(partitionManager.getRegionDatabase(regionId)).thenReturn("root.test");
+    when(partitionManager.getReplicaSets("root.test", Collections.singletonList(regionId)))
         .thenReturn(
             Collections.singletonList(
                 new TRegionReplicaSet(
                     regionId, Arrays.asList(original, coordinator, otherReplica))));
-    when(nodeManager.filterDataNodeThroughStatus(NodeStatus.Running, NodeStatus.ReadOnly))
-        .thenReturn(
-            Arrays.asList(
-                new TDataNodeConfiguration().setLocation(original),
-                new TDataNodeConfiguration().setLocation(coordinator),
-                new TDataNodeConfiguration().setLocation(otherReplica),
-                new TDataNodeConfiguration().setLocation(nonReplica)));
-    when(loadManager.getRegionLeaderMap()).thenReturn(Collections.emptyMap());
+    when(loadManager.filterDataNodeThroughStatus(NodeStatus.Running, NodeStatus.ReadOnly))
+        .thenReturn(Arrays.asList(1, 2, 3, 4));
+    for (TDataNodeLocation location :
+        Arrays.asList(original, coordinator, otherReplica, nonReplica)) {
+      when(nodeManager.getRegisteredDataNode(location.getDataNodeId()))
+          .thenReturn(new TDataNodeConfiguration().setLocation(location));
+    }
+    when(loadManager.getRegionLeaderId(regionId)).thenReturn(-1);
     handler = new RegionMaintainHandler(configManager);
   }
 
@@ -144,7 +146,7 @@ public void testFilterReplicaWithoutLeaderCache() {
 
   @Test
   public void testFilterReplicaWithUnknownLeader() {
-    when(loadManager.getRegionLeaderMap()).thenReturn(Collections.singletonMap(regionId, -1));
+    when(loadManager.getRegionLeaderId(regionId)).thenReturn(-1);
     assertEquals(
         Optional.of(otherReplica),
         handler.filterDataNodeWithOtherRegionReplica(
@@ -154,8 +156,7 @@ public void testFilterReplicaWithUnknownLeader() {
 
   @Test
   public void testFilterReplicaPrefersKnownLeader() {
-    when(loadManager.getRegionLeaderMap())
-        .thenReturn(Collections.singletonMap(regionId, otherReplica.getDataNodeId()));
+    when(loadManager.getRegionLeaderId(regionId)).thenReturn(otherReplica.getDataNodeId());
     assertEquals(
         Optional.of(otherReplica),
         handler.filterDataNodeWithOtherRegionReplica(regionId, original));
@@ -178,27 +179,23 @@ public void testTransferLeaderWithoutLeaderCache() throws Exception {
 
   @Test
   public void testTransferLeaderWithUnknownLeader() throws Exception {
-    when(loadManager.getRegionLeaderMap()).thenReturn(Collections.singletonMap(regionId, -1));
+    when(loadManager.getRegionLeaderId(regionId)).thenReturn(-1);
     assertLeaderTransfer();
   }
 
   @Test
   public void testTransferLeaderWaitsForLeaderElection() throws Exception {
     // The filter has a known leader, while the transfer lookup observes a stale cache first.
-    when(loadManager.getRegionLeaderMap())
-        .thenReturn(Collections.singletonMap(regionId, otherReplica.getDataNodeId()))
-        .thenReturn(Collections.emptyMap())
-        .thenReturn(Collections.singletonMap(regionId, original.getDataNodeId()));
+    when(loadManager.getRegionLeaderId(regionId))
+        .thenReturn(otherReplica.getDataNodeId(), -1, original.getDataNodeId());
     assertLeaderTransfer();
     verify(loadManager).waitForRegionGroupReady(Collections.singletonList(regionId));
   }
 
   @Test
   public void testTransferLeaderWhenCacheEntryDisappears() throws Exception {
     // The cache can change between replica selection and the transfer's leader lookup.
-    when(loadManager.getRegionLeaderMap())
-        .thenReturn(Collections.singletonMap(regionId, -1))
-        .thenReturn(Collections.emptyMap());
+    when(loadManager.getRegionLeaderId(regionId)).thenReturn(otherReplica.getDataNodeId(), -1);
     assertLeaderTransfer();
   }
 
@@ -239,6 +236,52 @@ private void assertLeaderTransfer() throws Exception {
     }
   }
 
+  @Test
+  public void testFilterReplicaDoesNotReadUnrelatedConfigurations() {
+    assertEquals(
+        Optional.of(otherReplica),
+        handler.filterDataNodeWithOtherRegionReplica(
+            regionId, Arrays.asList(original, coordinator)));
+    verify(loadManager, never()).getRegionLeaderMap();
+    verify(nodeManager, never()).getRegisteredDataNode(nonReplica.getDataNodeId());
+    verify(nodeManager, never())
+        .filterDataNodeThroughStatus(NodeStatus.Running, NodeStatus.ReadOnly);
+  }
+
+  @Test
+  public void testFilterReplicaSkipsUnregisteredNode() {
+    when(nodeManager.getRegisteredDataNode(otherReplica.getDataNodeId()))
+        .thenReturn(new TDataNodeConfiguration());
+    assertFalse(
+        handler
+            .filterDataNodeWithOtherRegionReplica(regionId, Arrays.asList(original, coordinator))
+            .isPresent());
+  }
+
+  @Test
+  public void testFilterReplicaRejectsChangedEndpoint() {
+    TDataNodeLocation changedLocation = otherReplica.deepCopy();
+    changedLocation.setInternalEndPoint(new TEndPoint("127.0.0.2", 10730));
+    when(nodeManager.getRegisteredDataNode(otherReplica.getDataNodeId()))
+        .thenReturn(new TDataNodeConfiguration().setLocation(changedLocation));
+    assertFalse(
+        handler
+            .filterDataNodeWithOtherRegionReplica(regionId, Arrays.asList(original, coordinator))
+            .isPresent());
+  }
+
+  @Test
+  public void testFilterReplicaPreservesRequestedStatuses() {
+    when(loadManager.filterDataNodeThroughStatus(NodeStatus.Removing))
+        .thenReturn(Collections.singletonList(otherReplica.getDataNodeId()));
+    assertEquals(
+        Optional.of(otherReplica),
+        handler.filterDataNodeWithOtherRegionReplica(regionId, original, NodeStatus.Removing));
+    verify(nodeManager, never()).getRegisteredDataNode(coordinator.getDataNodeId());
+    verify(loadManager, never()).getRegionLeaderMap();
+    verify(nodeManager, never()).getRegisteredDataNode(nonReplica.getDataNodeId());
+  }
+
   private RegionMaintainHandler stubStatusAndSubmit(
       final NodeStatus status, final TSStatus submitResult) {
     RegionMaintainHandler retryHandler = spy(handler);