Fix ConfigNode NPE when extend/reconstruct region targets a non-DataNode id

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

@@ -36,6 +36,7 @@
 import org.slf4j.LoggerFactory;
 
 import java.sql.Connection;
+import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.ArrayList;
 import java.util.List;
@@ -120,6 +121,64 @@ public void singleRegionTest() throws Exception {
     }
   }
 
+  /**
+   * Regression test for TIMECHODB-0689: when the target id of "extend region" does not belong to
+   * any registered DataNode (e.g. it is a ConfigNode id or simply does not exist), the ConfigNode
+   * used to throw a NullPointerException in {@code checkExtendRegion} and the client only saw a
+   * misleading "Fail to connect to any config node" message. After the fix a clear, actionable
+   * error message must be returned instead.
+   */
+  @Test
+  public void extendRegionToInvalidDataNodeTest() throws Exception {
+    EnvFactory.getEnv()
+        .getConfig()
+        .getCommonConfig()
+        .setDataRegionConsensusProtocolClass(ConsensusFactory.IOT_CONSENSUS)
+        .setSchemaRegionConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS)
+        .setDataReplicationFactor(1)
+        .setSchemaReplicationFactor(1);
+
+    EnvFactory.getEnv().initClusterEnvironment(1, 3);
+
+    try (final Connection connection = makeItCloseQuietly(EnvFactory.getEnv().getConnection());
+        final Statement statement = makeItCloseQuietly(connection.createStatement())) {
+      // prepare data so that at least one region exists
+      statement.execute(INSERTION1);
+      statement.execute(FLUSH_COMMAND);
+
+      Map<Integer, Set<Integer>> regionMap = getAllRegionMap(statement);
+      Set<Integer> allDataNodeId = getAllDataNodes(statement);
+      Assert.assertFalse(regionMap.isEmpty());
+
+      int selectedRegion = regionMap.keySet().iterator().next();
+      // an id that is guaranteed not to belong to any registered DataNode; a ConfigNode id triggers
+      // the exact same code path (getRegisteredDataNode returns an empty configuration whose
+      // location is null)
+      int invalidDataNodeId = 9999;
+      Assert.assertFalse(allDataNodeId.contains(invalidDataNodeId));
+
+      try {
+        statement.execute(String.format(EXPAND_FORMAT, selectedRegion, invalidDataNodeId));
+        Assert.fail("extend region to a non-existent DataNode is expected to fail");
+      } catch (SQLException e) {
+        String message = e.getMessage();
+        LOGGER.info("extend region to invalid DataNode failed as expected: {}", message);
+        Assert.assertNotNull(message);
+        // the ConfigNode must not crash with an NPE any more ...
+        Assert.assertFalse(
+            "ConfigNode should not throw NullPointerException, but got: " + message,
+            message.contains("NullPointerException"));
+        // ... and the submission must be rejected cleanly. "extend region" wraps every region's
+        // result, so the top-level message only reports the aggregate counts; the concrete "does
+        // not
+        // exist in the cluster" reason is carried in the per-region sub-status.
+        Assert.assertTrue(
+            "Expected the extend submission to be rejected but got: " + message,
+            message.contains("failed to submit: 1"));
+      }
+    }
+  }
+
   private void regionGroupExpand(
       Statement statement,
       SyncConfigNodeIServiceClient client,

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

@@ -45,6 +45,7 @@
 
 import java.io.File;
 import java.sql.Connection;
+import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.Collections;
 import java.util.Map;
@@ -220,4 +221,58 @@ public void normal1C3DTest() throws Exception {
       }
     }
   }
+
+  /**
+   * Regression test for TIMECHODB-0689 (reconstruct path): targeting "reconstruct region" at an id
+   * that is not a registered DataNode (a ConfigNode id or a non-existent id) used to trigger a
+   * NullPointerException in {@code checkReconstructRegion}. After the fix a clear, correct error
+   * message must be returned instead of crashing the ConfigNode RPC.
+   */
+  @Test
+  public void reconstructRegionToInvalidDataNodeTest() throws Exception {
+    EnvFactory.getEnv()
+        .getConfig()
+        .getCommonConfig()
+        .setDataRegionConsensusProtocolClass(ConsensusFactory.IOT_CONSENSUS)
+        .setSchemaRegionConsensusProtocolClass(ConsensusFactory.RATIS_CONSENSUS)
+        .setDataReplicationFactor(1)
+        .setSchemaReplicationFactor(1);
+
+    EnvFactory.getEnv().initClusterEnvironment(1, 3);
+
+    try (Connection connection = makeItCloseQuietly(EnvFactory.getEnv().getConnection());
+        Statement statement = makeItCloseQuietly(connection.createStatement())) {
+      // prepare data so that at least one region exists
+      statement.execute(INSERTION1);
+      statement.execute(FLUSH_COMMAND);
+
+      Map<Integer, Set<Integer>> regionMap = getAllRegionMap(statement);
+      Set<Integer> allDataNodeId = getAllDataNodes(statement);
+      Assert.assertFalse(regionMap.isEmpty());
+
+      int selectedRegion = regionMap.keySet().iterator().next();
+      // an id that is guaranteed not to belong to any registered DataNode; a ConfigNode id triggers
+      // the exact same code path (getRegisteredDataNode returns an empty configuration whose
+      // location is null)
+      int invalidDataNodeId = 9999;
+      Assert.assertFalse(allDataNodeId.contains(invalidDataNodeId));
+
+      try {
+        statement.execute(String.format(RECONSTRUCT_FORMAT, selectedRegion, invalidDataNodeId));
+        Assert.fail("reconstruct region on a non-existent DataNode is expected to fail");
+      } catch (SQLException e) {
+        String message = e.getMessage();
+        LOGGER.info("reconstruct region on invalid DataNode failed as expected: {}", message);
+        Assert.assertNotNull(message);
+        // the ConfigNode must not crash with an NPE any more ...
+        Assert.assertFalse(
+            "ConfigNode should not throw NullPointerException, but got: " + message,
+            message.contains("NullPointerException"));
+        // ... and the client should receive a clear, correct error message
+        Assert.assertTrue(
+            "Expected a 'does not exist in the cluster' error but got: " + message,
+            message.contains("does not exist in the cluster"));
+      }
+    }
+  }
 }

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

@@ -56,12 +56,14 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 
+import static org.apache.iotdb.confignode.it.regionmigration.IoTDBRegionOperationReliabilityITFramework.getAllRegionMap;
 import static org.apache.iotdb.confignode.it.regionmigration.IoTDBRegionOperationReliabilityITFramework.getDataRegionMap;
 import static org.apache.iotdb.confignode.it.removedatanode.IoTDBRemoveDataNodeUtils.awaitUntilSuccess;
 import static org.apache.iotdb.confignode.it.removedatanode.IoTDBRemoveDataNodeUtils.generateRemoveString;
 import static org.apache.iotdb.confignode.it.removedatanode.IoTDBRemoveDataNodeUtils.getConnectionWithSQLType;
 import static org.apache.iotdb.confignode.it.removedatanode.IoTDBRemoveDataNodeUtils.restartDataNodes;
 import static org.apache.iotdb.confignode.it.removedatanode.IoTDBRemoveDataNodeUtils.selectRemoveDataNodes;
+import static org.apache.iotdb.confignode.it.removedatanode.IoTDBRemoveDataNodeUtils.selectRemoveDataNodesWithoutRegionConflict;
 import static org.apache.iotdb.util.MagicUtils.makeItCloseQuietly;
 
 @Category({ClusterIT.class})
@@ -286,8 +288,17 @@ public void testRemoveDataNode(
         allDataNodeId.add(result.getInt(ColumnHeaderConstant.NODE_ID));
       }
 
-      // Select data nodes to remove
-      final Set<Integer> removeDataNodes = selectRemoveDataNodes(allDataNodeId, removeDataNodeNum);
+      // Select data nodes to remove. When removing more than one DataNode we must avoid picking two
+      // DataNodes that host replicas of the same consensus group, otherwise the
+      // RemoveDataNodesProcedure would try to migrate two replicas of one group at once, which the
+      // ConfigNode rejects ("Only one replica of the same consensus group is allowed to be migrated
+      // at the same time."). Selecting purely at random therefore makes this test flaky, so use a
+      // conflict-free selection.
+      final Set<Integer> removeDataNodes =
+          removeDataNodeNum > 1
+              ? selectRemoveDataNodesWithoutRegionConflict(
+                  allDataNodeId, removeDataNodeNum, getAllRegionMap(statement))
+              : selectRemoveDataNodes(allDataNodeId, removeDataNodeNum);
 
       List<DataNodeWrapper> removeDataNodeWrappers =
           removeDataNodes.stream()

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

@@ -35,8 +35,10 @@
 import java.sql.SQLException;
 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.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
@@ -73,6 +75,59 @@ public static Set<Integer> selectRemoveDataNodes(
     return new HashSet<>(shuffledDataNodeIds.subList(0, removeDataNodeNum));
   }
 
+  /**
+   * Select {@code removeDataNodeNum} DataNodes to remove such that no two of them host a replica of
+   * the same consensus group. Removing two DataNodes that share a region group would make the
+   * 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.
+   *
+   * @param allDataNodeId all registered DataNode ids
+   * @param removeDataNodeNum how many DataNodes to remove
+   * @param regionMap regionId -&gt; the set of DataNode ids hosting a replica of that region group
+   *     (all region types)
+   * @return the selected DataNode ids, or throws if no conflict-free selection of the requested
+   *     size exists
+   */
+  public static Set<Integer> selectRemoveDataNodesWithoutRegionConflict(
+      Set<Integer> allDataNodeId, int removeDataNodeNum, Map<Integer, Set<Integer>> regionMap) {
+    // dataNodeId -> the set of region groups it hosts a replica of
+    Map<Integer, Set<Integer>> dataNodeToRegions = new HashMap<>();
+    for (Integer dataNodeId : allDataNodeId) {
+      dataNodeToRegions.put(dataNodeId, new HashSet<>());
+    }
+    regionMap.forEach(
+        (regionId, dataNodeIds) ->
+            dataNodeIds.forEach(
+                dataNodeId ->
+                    dataNodeToRegions
+                        .computeIfAbsent(dataNodeId, id -> new HashSet<>())
+                        .add(regionId)));
+
+    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;
+        }
+      }
+    }
+    throw new IllegalStateException(
+        String.format(
+            "Cannot select %d DataNodes to remove without a same-region-group conflict. "
+                + "allDataNodeId=%s, regionMap=%s",
+            removeDataNodeNum, allDataNodeId, regionMap));
+  }
+
   public static void restartDataNodes(List<DataNodeWrapper> dataNodeWrappers) {
     dataNodeWrappers.parallelStream()
         .forEach(

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

@@ -881,7 +881,6 @@
                 new Pair<>("Coordinator", coordinator)),
             req.getModel());
 
-    ConfigNodeConfig conf = ConfigNodeDescriptor.getInstance().getConf();
     if (configManager
             .getPartitionManager()
             .getAllReplicaSetsMap(regionId.getType())
@@ -1137,7 +1136,7 @@
         return new TSStatus(TSStatusCode.MIGRATE_REGION_ERROR.getStatusCode())
             .setMessage(
                 String.format(
                     "Target DataNode %s does not exist in the cluster",
                     migrateRegionReq.getToId()));
       }
       final TDataNodeLocation originalDataNode = originalDataNodeConfiguration.getLocation();
@@ -1245,10 +1244,29 @@
     return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
   }
 
+  /**
+   * Resolve the location of a registered DataNode, or return {@code null} if the given id does not
+   * belong to any registered DataNode (e.g. it is a ConfigNode id or simply does not exist). {@link
+   * org.apache.iotdb.confignode.manager.node.NodeManager#getRegisteredDataNode} returns an empty
+   * {@link TDataNodeConfiguration} whose location is {@code null} in that case, so callers must not
+   * dereference the result blindly.
+   */
+  private TDataNodeLocation getRegisteredDataNodeLocationOrNull(int dataNodeId) {
+    return configManager.getNodeManager().getRegisteredDataNode(dataNodeId).getLocation();
+  }
+
   public TSStatus reconstructRegion(TReconstructRegionReq req) {
     RegionMaintainHandler handler = env.getRegionMaintainHandler();
     final TDataNodeLocation targetDataNode =
-        configManager.getNodeManager().getRegisteredDataNode(req.getDataNodeId()).getLocation();
+        getRegisteredDataNodeLocationOrNull(req.getDataNodeId());
+    if (targetDataNode == null) {
+      // The target id is not a registered DataNode. Reject here instead of pushing a null down into
+      // checkReconstructRegion, which would otherwise throw a NullPointerException.
+      return new TSStatus(TSStatusCode.RECONSTRUCT_REGION_ERROR.getStatusCode())
+          .setMessage(
+              String.format(
+                  "Target DataNode %s does not exist in the cluster", req.getDataNodeId()));
+    }
     try (AutoCloseableLock ignoredLock =
         AutoCloseableLock.acquire(env.getSubmitRegionMigrateLock())) {
       List<ReconstructRegionProcedure> procedures = new ArrayList<>();
@@ -1352,7 +1370,15 @@
 
       // find target dn
       final TDataNodeLocation targetDataNode =
-          configManager.getNodeManager().getRegisteredDataNode(req.getDataNodeId()).getLocation();
+          getRegisteredDataNodeLocationOrNull(req.getDataNodeId());
+      if (targetDataNode == null) {
+        // The target id is not a registered DataNode. Reject here instead of pushing a null down
+        // into checkExtendRegion, which would otherwise throw a NullPointerException.
+        return new TSStatus(TSStatusCode.EXTEND_REGION_ERROR.getStatusCode())
+            .setMessage(
+                String.format(
+                    "Target DataNode %s does not exist in the cluster", req.getDataNodeId()));
+      }
       // select coordinator for adding peer
       RegionMaintainHandler handler = env.getRegionMaintainHandler();
       // TODO: choose the DataNode which has lowest load