Fix two flaky multi-DataNode/multi-region ITs from #18046 (remove-DataNode selection + migrate precondition)

integration-test/src/test/java/org/apache/iotdb/confignode/it/regionmigration/pass/commit/IoTDBMigrateMultiRegionForIoTV1IT.java

@@ -68,7 +68,14 @@ public void multiRegionMigrateTest() throws Exception {
         .setDataRegionConsensusProtocolClass(ConsensusFactory.IOT_CONSENSUS)
         .setSchemaRegionConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS)
         .setDataReplicationFactor(1)
-        .setSchemaReplicationFactor(1);
+        .setSchemaReplicationFactor(1)
+        // Create 6 data region groups (> 5 DataNodes) so that, with replication factor 1, at least
+        // one DataNode is guaranteed by pigeonhole to host >= 2 regions - the precondition of
+        // selectDataNodeHostingMultipleRegions below. Under the default AUTO policy only ~2-3
+        // regions were created and a balanced spread could leave every DataNode with a single
+        // region, making this test flaky ("Cannot find a DataNode hosting at least two regions").
+        .setDataRegionGroupExtensionPolicy("CUSTOM")
+        .setDefaultDataRegionGroupNumPerDatabase(6);
 
     EnvFactory.getEnv().initClusterEnvironment(1, 5);
 

integration-test/src/test/java/org/apache/iotdb/confignode/it/removedatanode/IoTDBRemoveDataNodeNormalIT.java

@@ -357,7 +357,7 @@ allDataNodeId, removeDataNodeNum, getAllRegionMap(statement))
         removeSuccess = true;
       } catch (ConditionTimeoutException e) {
         if (expectRemoveSuccess) {
-          LOGGER.error("Remove DataNodes timeout in 2 minutes");
+          LOGGER.error("Remove DataNodes timeout in 5 minutes");
           Assert.fail();
         }
       }

integration-test/src/test/java/org/apache/iotdb/confignode/it/removedatanode/IoTDBRemoveDataNodeUtils.java

@@ -37,6 +37,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -81,8 +82,9 @@ public static Set<Integer> selectRemoveDataNodes(
    * generated {@code RemoveDataNodesProcedure} try to migrate two replicas of the same group at
    * once, which the ConfigNode rejects ("Only one replica of the same consensus group is allowed to
    * be migrated at the same time."). Randomly picking DataNodes (see {@link
-   * #selectRemoveDataNodes}) therefore makes multi-DataNode-remove tests flaky; this method keeps
-   * the selection valid and deterministic-enough for such tests.
+   * #selectRemoveDataNodes}) therefore makes multi-DataNode-remove tests flaky; this method instead
+   * searches exhaustively (with backtracking) for a conflict-free selection, so it fails only when
+   * no such selection exists rather than when an unlucky pick order dead-ends.
    *
    * @param allDataNodeId all registered DataNode ids
    * @param removeDataNodeNum how many DataNodes to remove
@@ -106,20 +108,20 @@ public static Set<Integer> selectRemoveDataNodesWithoutRegionConflict(
                         .computeIfAbsent(dataNodeId, id -> new HashSet<>())
                         .add(regionId)));
 
+    // Shuffle so that, when several conflict-free selections exist, a random one is returned (keeps
+    // the coverage of this test varied across runs).
     List<Integer> shuffledDataNodeIds = new ArrayList<>(allDataNodeId);
     Collections.shuffle(shuffledDataNodeIds);
 
-    Set<Integer> selected = new HashSet<>();
-    Set<Integer> coveredRegions = new HashSet<>();
-    for (Integer dataNodeId : shuffledDataNodeIds) {
-      Set<Integer> regions = dataNodeToRegions.getOrDefault(dataNodeId, Collections.emptySet());
-      if (Collections.disjoint(coveredRegions, regions)) {
-        selected.add(dataNodeId);
-        coveredRegions.addAll(regions);
-        if (selected.size() == removeDataNodeNum) {
-          return selected;
-        }
-      }
+    // Search exhaustively (with backtracking) for a set of DataNodes whose hosted region groups are
+    // pairwise disjoint. A single-pass greedy that unconditionally commits to the first shuffled
+    // DataNode can dead-end and wrongly throw even though a valid conflict-free selection exists -
+    // e.g. when that first DataNode happens to share a consensus group with every other DataNode -
+    // which made this test flaky.
+    Set<Integer> selected = new LinkedHashSet<>();
+    if (searchConflictFreeDataNodes(
+        shuffledDataNodeIds, 0, removeDataNodeNum, new HashSet<>(), selected, dataNodeToRegions)) {
+      return selected;
     }
     throw new IllegalStateException(
         String.format(
@@ -128,6 +130,45 @@ public static Set<Integer> selectRemoveDataNodesWithoutRegionConflict(
             removeDataNodeNum, allDataNodeId, regionMap));
   }
 
+  /**
+   * Depth-first search with backtracking for {@code need} DataNodes whose hosted region groups are
+   * pairwise disjoint. On success returns {@code true} and leaves the chosen ids in {@code
+   * selected}; on failure returns {@code false} with {@code selected} restored to its prior state.
+   *
+   * @param dataNodeIds candidate DataNode ids (iterated from {@code start} onwards)
+   * @param start index to start picking from, so each combination is visited at most once
+   * @param need how many DataNodes must be selected in total
+   * @param coveredRegions region groups already covered by the currently-selected DataNodes
+   * @param selected the DataNodes chosen so far (mutated during the search)
+   * @param dataNodeToRegions dataNodeId -&gt; the region groups it hosts a replica of
+   */
+  private static boolean searchConflictFreeDataNodes(
+      List<Integer> dataNodeIds,
+      int start,
+      int need,
+      Set<Integer> coveredRegions,
+      Set<Integer> selected,
+      Map<Integer, Set<Integer>> dataNodeToRegions) {
+    if (selected.size() == need) {
+      return true;
+    }
+    for (int i = start; i < dataNodeIds.size(); i++) {
+      Integer dataNodeId = dataNodeIds.get(i);
+      Set<Integer> regions = dataNodeToRegions.getOrDefault(dataNodeId, Collections.emptySet());
+      if (Collections.disjoint(coveredRegions, regions)) {
+        selected.add(dataNodeId);
+        Set<Integer> nextCovered = new HashSet<>(coveredRegions);
+        nextCovered.addAll(regions);
+        if (searchConflictFreeDataNodes(
+            dataNodeIds, i + 1, need, nextCovered, selected, dataNodeToRegions)) {
+          return true;
+        }
+        selected.remove(dataNodeId);
+      }
+    }
+    return false;
+  }
+
   public static void restartDataNodes(List<DataNodeWrapper> dataNodeWrappers) {
     dataNodeWrappers.parallelStream()
         .forEach(