[To dev/1.3][remove datanode] cherry pick related commits (#16711)

HxpSerein · web-flow · commit b1ae6dc5cfb2 · 2025-11-07T19:15:43.000+08:00
* [remove datanode] GCR load balancing implement for removing datanode (#15282) (cherry picked from commit 7650b47) * [remove datanode] Fix IoTDBRemoveDataNodeNormalIT #15429 (cherry picked from commit 9537806) * [remove datanode] Accelerate GCR load balancing implement (#15535) (cherry picked from commit 51bad1e) * [remove datanode] Fix ArrayIndexOutOfBoundsException in computeInitialDbLoad (#15718) (cherry picked from commit 346ee72)
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/balancer/region/GreedyCopySetRegionGroupAllocator.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/balancer/region/GreedyCopySetRegionGroupAllocator.java
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/balancer/region/GreedyRegionGroupAllocator.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/balancer/region/GreedyRegionGroupAllocator.java
@@ -57,6 +57,19 @@ public TRegionReplicaSet generateOptimalRegionReplicasDistribution(
         weightList.stream().limit(replicationFactor).collect(Collectors.toList()));
   }
 
+  @Override
+  public Map<TConsensusGroupId, TDataNodeConfiguration> removeNodeReplicaSelect(
+      Map<Integer, TDataNodeConfiguration> availableDataNodeMap,
+      Map<Integer, Double> freeDiskSpaceMap,
+      List<TRegionReplicaSet> allocatedRegionGroups,
+      Map<TConsensusGroupId, String> regionDatabaseMap,
+      Map<String, List<TRegionReplicaSet>> databaseAllocatedRegionGroupMap,
+      Map<TConsensusGroupId, TRegionReplicaSet> remainReplicasMap) {
+    // TODO: Implement this method
+    throw new UnsupportedOperationException(
+        "The removeNodeReplicaSelect method of GreedyRegionGroupAllocator is yet to be implemented.");
+  }
+
   private List<TDataNodeLocation> buildWeightList(
       Map<Integer, TDataNodeConfiguration> availableDataNodeMap,
       Map<Integer, Double> freeDiskSpaceMap,
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/balancer/region/IRegionGroupAllocator.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/balancer/region/IRegionGroupAllocator.java
@@ -47,4 +47,25 @@ TRegionReplicaSet generateOptimalRegionReplicasDistribution(
       List<TRegionReplicaSet> databaseAllocatedRegionGroups,
       int replicationFactor,
       TConsensusGroupId consensusGroupId);
+
+  /**
+   * Select the optimal DataNode to place the new replica on along with the remaining replica set.
+   *
+   * @param availableDataNodeMap DataNodes that can be used for allocation
+   * @param freeDiskSpaceMap The free disk space of the DataNodes
+   * @param allocatedRegionGroups Allocated RegionGroups
+   * @param regionDatabaseMap A mapping from each region identifier to its corresponding database
+   *     name
+   * @param databaseAllocatedRegionGroupMap Allocated RegionGroups within the same Database with the
+   *     replica set
+   * @param remainReplicasMap the remaining replica set excluding the removed DataNodes
+   * @return The optimal DataNode to place the new replica on along with the remaining replicas
+   */
+  Map<TConsensusGroupId, TDataNodeConfiguration> removeNodeReplicaSelect(
+      Map<Integer, TDataNodeConfiguration> availableDataNodeMap,
+      Map<Integer, Double> freeDiskSpaceMap,
+      List<TRegionReplicaSet> allocatedRegionGroups,
+      Map<TConsensusGroupId, String> regionDatabaseMap,
+      Map<String, List<TRegionReplicaSet>> databaseAllocatedRegionGroupMap,
+      Map<TConsensusGroupId, TRegionReplicaSet> remainReplicasMap);
 }
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/balancer/region/PartiteGraphReplicationRegionGroupAllocator.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/balancer/region/PartiteGraphReplicationRegionGroupAllocator.java
@@ -113,6 +113,19 @@ public TRegionReplicaSet generateOptimalRegionReplicasDistribution(
     return result;
   }
 
+  @Override
+  public Map<TConsensusGroupId, TDataNodeConfiguration> removeNodeReplicaSelect(
+      Map<Integer, TDataNodeConfiguration> availableDataNodeMap,
+      Map<Integer, Double> freeDiskSpaceMap,
+      List<TRegionReplicaSet> allocatedRegionGroups,
+      Map<TConsensusGroupId, String> regionDatabaseMap,
+      Map<String, List<TRegionReplicaSet>> databaseAllocatedRegionGroupMap,
+      Map<TConsensusGroupId, TRegionReplicaSet> remainReplicasMap) {
+    // TODO: Implement this method
+    throw new UnsupportedOperationException(
+        "The removeNodeReplicaSelect method of PartiteGraphPlacementRegionGroupAllocator is yet to be implemented.");
+  }
+
   private void prepare(
       int replicationFactor,
       Map<Integer, TDataNodeConfiguration> availableDataNodeMap,
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/env/RemoveDataNodeHandler.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/env/RemoveDataNodeHandler.java
@@ -20,6 +20,7 @@
 package org.apache.iotdb.confignode.procedure.env;
 
 import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId;
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType;
 import org.apache.iotdb.common.rpc.thrift.TDataNodeConfiguration;
 import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation;
 import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet;
@@ -37,6 +38,8 @@
 import org.apache.iotdb.confignode.consensus.request.write.datanode.RemoveDataNodePlan;
 import org.apache.iotdb.confignode.consensus.response.datanode.DataNodeToStatusResp;
 import org.apache.iotdb.confignode.manager.ConfigManager;
+import org.apache.iotdb.confignode.manager.load.balancer.region.GreedyCopySetRegionGroupAllocator;
+import org.apache.iotdb.confignode.manager.load.balancer.region.IRegionGroupAllocator;
 import org.apache.iotdb.confignode.manager.load.cache.node.NodeHeartbeatSample;
 import org.apache.iotdb.confignode.manager.load.cache.region.RegionHeartbeatSample;
 import org.apache.iotdb.confignode.manager.partition.PartitionMetrics;
@@ -51,10 +54,13 @@
 
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.TreeMap;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 import static org.apache.iotdb.confignode.conf.ConfigNodeConstant.REMOVE_DATANODE_PROCESS;
@@ -70,8 +76,22 @@ public class RemoveDataNodeHandler {
 
   private final ConfigManager configManager;
 
+  private final IRegionGroupAllocator regionGroupAllocator;
+
   public RemoveDataNodeHandler(ConfigManager configManager) {
     this.configManager = configManager;
+
+    switch (ConfigNodeDescriptor.getInstance().getConf().getRegionGroupAllocatePolicy()) {
+      case GREEDY:
+        this.regionGroupAllocator = new GreedyCopySetRegionGroupAllocator();
+        break;
+      case PGR:
+        this.regionGroupAllocator = new GreedyCopySetRegionGroupAllocator();
+        break;
+      case GCR:
+      default:
+        this.regionGroupAllocator = new GreedyCopySetRegionGroupAllocator();
+    }
   }
 
   /**
@@ -193,6 +213,172 @@ public List<RegionMigrationPlan> getRegionMigrationPlans(
     return regionMigrationPlans;
   }
 
+  /**
+   * Retrieves all region migration plans for the specified removed DataNodes and selects the
+   * destination.
+   *
+   * @param removedDataNodes the list of DataNodes from which to obtain migration plans
+   * @return a list of region migration plans associated with the removed DataNodes
+   */
+  public List<RegionMigrationPlan> selectedRegionMigrationPlans(
+      List<TDataNodeLocation> removedDataNodes) {
+
+    Set<Integer> removedDataNodesSet = new HashSet<>();
+    for (TDataNodeLocation removedDataNode : removedDataNodes) {
+      removedDataNodesSet.add(removedDataNode.dataNodeId);
+    }
+
+    final List<TDataNodeConfiguration> availableDataNodes =
+        configManager
+            .getNodeManager()
+            .filterDataNodeThroughStatus(NodeStatus.Running, NodeStatus.Unknown)
+            .stream()
+            .filter(node -> !removedDataNodesSet.contains(node.getLocation().getDataNodeId()))
+            .collect(Collectors.toList());
+
+    List<RegionMigrationPlan> regionMigrationPlans = new ArrayList<>();
+
+    regionMigrationPlans.addAll(
+        selectMigrationPlans(availableDataNodes, TConsensusGroupType.DataRegion, removedDataNodes));
+
+    regionMigrationPlans.addAll(
+        selectMigrationPlans(
+            availableDataNodes, TConsensusGroupType.SchemaRegion, removedDataNodes));
+
+    return regionMigrationPlans;
+  }
+
+  public List<RegionMigrationPlan> selectMigrationPlans(
+      List<TDataNodeConfiguration> availableDataNodes,
+      TConsensusGroupType consensusGroupType,
+      List<TDataNodeLocation> removedDataNodes) {
+
+    // Retrieve all allocated replica sets for the given consensus group type
+    List<TRegionReplicaSet> allocatedReplicaSets =
+        configManager.getPartitionManager().getAllReplicaSets(consensusGroupType);
+
+    // Step 1: Identify affected replica sets and record the removed DataNode for each replica set
+    Map<TConsensusGroupId, TDataNodeLocation> removedNodeMap = new HashMap<>();
+    Set<TRegionReplicaSet> affectedReplicaSets =
+        identifyAffectedReplicaSets(allocatedReplicaSets, removedDataNodes, removedNodeMap);
+
+    // Step 2: Update affected replica sets by removing the removed DataNode
+    updateReplicaSets(allocatedReplicaSets, affectedReplicaSets, removedNodeMap);
+
+    // Build a mapping of available DataNodes and their free disk space (computed only once)
+    Map<Integer, TDataNodeConfiguration> availableDataNodeMap =
+        buildAvailableDataNodeMap(availableDataNodes);
+    Map<Integer, Double> freeDiskSpaceMap = buildFreeDiskSpaceMap(availableDataNodes);
+
+    // Step 3: For each affected replica set, select a new destination DataNode and create a
+    // migration plan
+    List<RegionMigrationPlan> migrationPlans = new ArrayList<>();
+
+    Map<TConsensusGroupId, TRegionReplicaSet> remainReplicasMap = new HashMap<>();
+    Map<TConsensusGroupId, String> regionDatabaseMap = new HashMap<>();
+    Map<String, List<TRegionReplicaSet>> databaseAllocatedRegionGroupMap = new HashMap<>();
+
+    for (TRegionReplicaSet replicaSet : affectedReplicaSets) {
+      remainReplicasMap.put(replicaSet.getRegionId(), replicaSet);
+      String database =
+          configManager.getPartitionManager().getRegionDatabase(replicaSet.getRegionId());
+      List<TRegionReplicaSet> databaseAllocatedReplicaSets =
+          configManager.getPartitionManager().getAllReplicaSets(database, consensusGroupType);
+      regionDatabaseMap.put(replicaSet.getRegionId(), database);
+      databaseAllocatedRegionGroupMap.put(database, databaseAllocatedReplicaSets);
+    }
+
+    Map<TConsensusGroupId, TDataNodeConfiguration> result =
+        regionGroupAllocator.removeNodeReplicaSelect(
+            availableDataNodeMap,
+            freeDiskSpaceMap,
+            allocatedReplicaSets,
+            regionDatabaseMap,
+            databaseAllocatedRegionGroupMap,
+            remainReplicasMap);
+
+    for (TConsensusGroupId regionId : result.keySet()) {
+
+      TDataNodeConfiguration selectedNode = result.get(regionId);
+      LOGGER.info(
+          "Selected DataNode {} for Region {}",
+          selectedNode.getLocation().getDataNodeId(),
+          regionId);
+
+      // Create the migration plan
+      RegionMigrationPlan plan = RegionMigrationPlan.create(regionId, removedNodeMap.get(regionId));
+      plan.setToDataNode(selectedNode.getLocation());
+      migrationPlans.add(plan);
+    }
+    return migrationPlans;
+  }
+
+  /**
+   * Identifies affected replica sets from allocatedReplicaSets that contain any DataNode in
+   * removedDataNodes, and records the removed DataNode for each replica set.
+   */
+  private Set<TRegionReplicaSet> identifyAffectedReplicaSets(
+      List<TRegionReplicaSet> allocatedReplicaSets,
+      List<TDataNodeLocation> removedDataNodes,
+      Map<TConsensusGroupId, TDataNodeLocation> removedNodeMap) {
+
+    Set<TRegionReplicaSet> affectedReplicaSets = new HashSet<>();
+    // Create a copy of allocatedReplicaSets to avoid concurrent modifications
+    List<TRegionReplicaSet> allocatedCopy = new ArrayList<>(allocatedReplicaSets);
+
+    for (TDataNodeLocation removedNode : removedDataNodes) {
+      allocatedCopy.stream()
+          .filter(replicaSet -> replicaSet.getDataNodeLocations().contains(removedNode))
+          .forEach(
+              replicaSet -> {
+                removedNodeMap.put(replicaSet.getRegionId(), removedNode);
+                affectedReplicaSets.add(replicaSet);
+              });
+    }
+    return affectedReplicaSets;
+  }
+
+  /**
+   * Updates each affected replica set by removing the removed DataNode from its list. The
+   * allocatedReplicaSets list is updated accordingly.
+   */
+  private void updateReplicaSets(
+      List<TRegionReplicaSet> allocatedReplicaSets,
+      Set<TRegionReplicaSet> affectedReplicaSets,
+      Map<TConsensusGroupId, TDataNodeLocation> removedNodeMap) {
+    for (TRegionReplicaSet replicaSet : affectedReplicaSets) {
+      // Remove the replica set, update its node list, then re-add it
+      allocatedReplicaSets.remove(replicaSet);
+      replicaSet.getDataNodeLocations().remove(removedNodeMap.get(replicaSet.getRegionId()));
+      allocatedReplicaSets.add(replicaSet);
+    }
+  }
+
+  /**
+   * Constructs a mapping from DataNodeId to TDataNodeConfiguration from the available DataNodes.
+   */
+  private Map<Integer, TDataNodeConfiguration> buildAvailableDataNodeMap(
+      List<TDataNodeConfiguration> availableDataNodes) {
+    return availableDataNodes.stream()
+        .collect(
+            Collectors.toMap(
+                dataNode -> dataNode.getLocation().getDataNodeId(), Function.identity()));
+  }
+
+  /** Constructs a mapping of free disk space for each DataNode. */
+  private Map<Integer, Double> buildFreeDiskSpaceMap(
+      List<TDataNodeConfiguration> availableDataNodes) {
+    Map<Integer, Double> freeDiskSpaceMap = new HashMap<>(availableDataNodes.size());
+    availableDataNodes.forEach(
+        dataNode ->
+            freeDiskSpaceMap.put(
+                dataNode.getLocation().getDataNodeId(),
+                configManager
+                    .getLoadManager()
+                    .getFreeDiskSpace(dataNode.getLocation().getDataNodeId())));
+    return freeDiskSpaceMap;
+  }
+
   /**
    * Broadcasts DataNodes' status change, preventing disabled DataNodes from accepting read or write
    * requests.
diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/node/RemoveDataNodesProcedure.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/node/RemoveDataNodesProcedure.java
@@ -121,7 +121,8 @@ protected Flow executeFromState(ConfigNodeProcedureEnv env, RemoveDataNodeState
           removedDataNodes.forEach(
               dataNode -> removedNodeStatusMap.put(dataNode.getDataNodeId(), NodeStatus.Removing));
           removeDataNodeHandler.changeDataNodeStatus(removedDataNodes, removedNodeStatusMap);
-          regionMigrationPlans = removeDataNodeHandler.getRegionMigrationPlans(removedDataNodes);
+          regionMigrationPlans =
+              removeDataNodeHandler.selectedRegionMigrationPlans(removedDataNodes);
           LOG.info(
               "{}, DataNode regions to be removed is {}",
               REMOVE_DATANODE_PROCESS,
@@ -165,8 +166,7 @@ private void submitChildRegionMigrate(ConfigNodeProcedureEnv env) {
         regionMigrationPlan -> {
           TConsensusGroupId regionId = regionMigrationPlan.getRegionId();
           TDataNodeLocation removedDataNode = regionMigrationPlan.getFromDataNode();
-          TDataNodeLocation destDataNode =
-              env.getRegionMaintainHandler().findDestDataNode(regionId);
+          TDataNodeLocation destDataNode = regionMigrationPlan.getToDataNode();
           // TODO: need to improve the coordinator selection method here, maybe through load
           // balancing and other means.
           final TDataNodeLocation coordinatorForAddPeer =
diff --git a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/load/balancer/region/GreedyCopySetRemoveNodeReplicaSelectTest.java b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/load/balancer/region/GreedyCopySetRemoveNodeReplicaSelectTest.java
