[remove datanode] Fix ArrayIndexOutOfBoundsException in computeInitialDbLoad (apache#15718)

Author: HxpSerein

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/balancer/region/GreedyCopySetRegionGroupAllocator.java

@@ -155,7 +155,7 @@ public Map<TConsensusGroupId, TDataNodeConfiguration> removeNodeReplicaSelect(
       // 1. prepare: compute regionCounter, databaseRegionCounter, and combinationCounter
 
       prepare(availableDataNodeMap, allocatedRegionGroups, Collections.emptyList());
-      computeInitialDbLoad(databaseAllocatedRegionGroupMap);
+      computeInitialDbLoad(availableDataNodeMap, databaseAllocatedRegionGroupMap);
 
       // 2. Build allowed candidate set for each region that needs to be migrated.
       // For each region in remainReplicasMap, the candidate destination nodes are all nodes in
@@ -432,9 +432,11 @@ private int[] getCurrentMetrics(
   /**
    * Compute the initial load for each database on each data node.
    *
+   * @param availableDataNodeMap currently available DataNodes, ensure size() >= replicationFactor
    * @param databaseAllocatedRegionGroupMap Mapping of each database to its list of replica sets.
    */
   private void computeInitialDbLoad(
+      Map<Integer, TDataNodeConfiguration> availableDataNodeMap,
       Map<String, List<TRegionReplicaSet>> databaseAllocatedRegionGroupMap) {
     initialDbLoad = new HashMap<>();
 
@@ -446,7 +448,9 @@ private void computeInitialDbLoad(
       for (TRegionReplicaSet replicaSet : replicaSets) {
         for (TDataNodeLocation location : replicaSet.getDataNodeLocations()) {
           int nodeId = location.getDataNodeId();
-          load[nodeId]++;
+          if (availableDataNodeMap.containsKey(nodeId)) {
+            load[nodeId]++;
+          }
         }
       }
       initialDbLoad.put(database, load);

iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/load/balancer/region/GreedyCopySetRemoveNodeReplicaSelectTest.java

@@ -81,15 +81,20 @@ public void testSelectDestNode() {
         DATA_REGION_PER_DATA_NODE * TEST_DATA_NODE_NUM / DATA_REPLICATION_FACTOR;
 
     List<TRegionReplicaSet> allocateResult = new ArrayList<>();
+    List<TRegionReplicaSet> databaseAllocateResult = new ArrayList<>();
     for (int index = 0; index < dataRegionGroupNum; index++) {
-      allocateResult.add(
+      TRegionReplicaSet replicaSet =
           GCR_ALLOCATOR.generateOptimalRegionReplicasDistribution(
               AVAILABLE_DATA_NODE_MAP,
               FREE_SPACE_MAP,
               allocateResult,
               allocateResult,
               DATA_REPLICATION_FACTOR,
-              new TConsensusGroupId(TConsensusGroupType.DataRegion, index)));
+              new TConsensusGroupId(TConsensusGroupType.DataRegion, index));
+      TRegionReplicaSet replicaSetCopy = new TRegionReplicaSet(replicaSet);
+
+      allocateResult.add(replicaSet);
+      databaseAllocateResult.add(replicaSetCopy);
     }
 
     List<TRegionReplicaSet> migratedReplicas =
@@ -158,7 +163,7 @@ public void testSelectDestNode() {
     }
     Map<TConsensusGroupId, TRegionReplicaSet> remainReplicasMap = new HashMap<>();
     Map<String, List<TRegionReplicaSet>> databaseAllocatedRegionGroupMap = new HashMap<>();
-    databaseAllocatedRegionGroupMap.put("database", allocateResult);
+    databaseAllocatedRegionGroupMap.put("database", databaseAllocateResult);
 
     for (TRegionReplicaSet remainReplicaSet : remainReplicas) {
       remainReplicasMap.put(remainReplicaSet.getRegionId(), remainReplicaSet);